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

我有一些现有的PyTorch代码，其中包含cuda()，如下所示，而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代码转换为CPU版本。我尝试查看了关于PyTorch CPU/GPU切换的一些相关帖子，但它们都涉及到device的使用，因此不适用于我的情况。

回答：

正如kHarshit在他的评论中指出的那样，你可以简单地用.cpu()替换.cuda()调用：

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)

附注，
请注意，VariableAPI已被弃用，不再需要使用。

学技术