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

Question

我有一些现有的 PyTorch 代码cuda()如下， whilenet是一个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代码转换为CPU版本。我尝试查看一些有关 PyTorch CPU/GPU 切换的相关帖子，但它们与使用有关device，因此不适用于我的情况。

Answer 1

正如 kHarshit 在他的评论中指出的，您可以简单地将.cuda()call 替换为.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，
请注意，该VariableAPI 已被弃用，不再需要。