在 PyTorch 中，如何将 cuda() 相关代码转换为 CPU 版本？答案

【问题标题】：In PyTorch, how to convert the cuda() related codes into CPU version?在 PyTorch 中，如何将 cuda() 相关代码转换为 CPU 版本？
【发布时间】：2023-03-03 05:15:20
【问题描述】：

我有一些现有的带有cuda() 的PyTorch 代码如下，而net 是MainModel.KitModel 对象：

net = torch.load(model_path)
net.cuda()

和

im = cv2.imread(image_path)
im = Variable(torch.from_numpy(im).unsqueeze(0).float().cuda())

我想在没有任何 GPU 的机器上测试代码，所以我想将 cuda-code 转换为 CPU 版本。我试图查看一些有关 PyTorch 的 CPU/GPU 切换的相关帖子，但它们与 device 的使用有关，因此不适用于我的情况。

【问题讨论】：

您可以简单地使用.cpu() 代替.cuda() 或根据gpu 可用性移动张量的设备变量。

标签： machine-learning image-processing computer-vision pytorch gpu

【解决方案1】：

   net = torch.load(model_path, map_location=torch.device('cpu'))

Pytorch 文档：https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-on-cpu-load-on-gpu

【讨论】：

【解决方案2】：

正如 kHarshit 在他的 comment 中指出的那样，您可以简单地将 .cuda() 调用替换为 .cpu()：

net.cpu()
# ...
im = torch.from_numpy(im).unsqueeze(0).float().cpu()

但是，每次您想从 GPU 迁移到 CPU 时，这都需要在多个位置更改代码，反之亦然。

为了缓解这个困难，pytorch 有一个更“通用”的方法.to()。
您可能有一个 device 变量来定义您希望 pytorch 运行的位置，这个 device 也可以是 CPU (!)。
例如：

if torch.cuda.is_available():
  device = torch.device("cuda")
else:
  device = torch.device("cpu")

一旦您在代码中确定一次想要/可以运行，只需使用.to() 将您的模型/变量发送到那里：

net.to(device)
# ...
im = torch.from_numpy(im).unsqueeze(0).float().to(device)

顺便说一句，
您也可以使用.to() 来控制数据类型（.float()）：

im = torch.from_numpy(im).unsqueeze(0).to(device=device, dtype=torch.float)

PS，
请注意，Variable API 已为 deprecated，不再需要。

【讨论】：