[英]A better way to make pytorch code agnostic to running on a CPU or GPU?
迁移指南建议以下内容使代码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
。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.