使 pytorch 代码与在 CPU 或 GPU 上运行无关的更好方法?

Posted

技术标签:

【中文标题】使 pytorch 代码与在 CPU 或 GPU 上运行无关的更好方法?【英文标题】:A better way to make pytorch code agnostic to running on a CPU or GPU? 【发布时间】:2019-03-07 21:04:21 【问题描述】:

迁移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)

这看起来并不优雅。有没有更好的办法?

【问题讨论】:

【参考方案1】:

怎么样

if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model = model.to(device)

?

如果您只有一个 GPU,则不需要 DataParallel

【讨论】:

这很有趣。我没有意识到DataParallel 仅适用于多 GPU。

以上是关于使 pytorch 代码与在 CPU 或 GPU 上运行无关的更好方法?的主要内容,如果未能解决你的问题,请参考以下文章

在 Google Colaboratory 上,对于 Pytorch,GPU 的性能比 CPU 慢

在GPU上训练PyTorch代码

在GPU上训练PyTorch代码

在GPU上训练PyTorch代码

是否有带有 CUDA 统一 GPU-CPU 内存分支的 PyTorch?

如何告诉 PyTorch 不使用 GPU?