迁移 guide推荐以下内容以使代码与 CPU/GPU 无关:
> # at beginning of the script
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# then whenever you get a new Tensor or Module
# this won't copy if they are already on the desired device
input = data.to(device)
model = MyModule(...).to(device)
我这样做并在纯 CPU 设备上运行了我的代码,但是当输入数组时我的模型崩溃了,因为它说它期待的是 CPU 张量而不是 GPU 张量。我的模型以某种方式自动将 CPU 输入数组转换为 GPU 数组。最后我在我的代码中找到了这个命令:
model = torch.nn.DataParallel(model).to(device)
即使我将模型转换为“cpu”,nn.DataParallel 也会覆盖它。我想到的最佳解决方案是有条件的:
if device.type=='cpu':
model = model.to(device)
else:
model = torch.nn.DataParallel(model).to(device)
这看起来并不优雅。有没有更好的办法?
最佳答案
怎么样
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model = model.to(device)
?
如果您只有一个 GPU,则不需要 DataParallel
。
关于python - 使 pytorch 代码与在 CPU 或 GPU 上运行无关的更好方法?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/52613383/