如何在PyTorch中修改模块部分参数的设备类型

问题原因 #

循环中迭代得到的p是临时变量，直接执行p = p.to('cuda')仅修改了临时变量的指向，没有改动模型内部实际存储的参数对象，因此迁移操作不会生效。

正确解决方法 #

方案1：全模型批量迁移（优先推荐） #

直接调用模型自带的to()方法，可一次性将所有参数、缓冲区迁移到指定设备，无需手动遍历参数：

net = net.to('cuda')

方案2：仅迁移位于CPU的参数 #

如果你确实需要仅迁移当前在CPU上的参数，需要直接修改参数的底层数据属性，若存在梯度也要同步迁移：

for p in net.parameters():
    if p.device == torch.device('cpu'):
        p.data = p.data.to('cuda')
        # 同步迁移梯度数据
        if p.grad is not None:
            p.grad.data = p.grad.data.to('cuda')

修改完成后再打印参数设备，就能看到CPU的参数已经成功迁移到CUDA。

内容的提问来源于stack exchange，提问作者skytree