In PyTorch, how to convert the cuda() related codes into CPU version?

user_9096553 0
Viewed 3499

I have some existing PyTorch codes with cuda() as below, while net is a MainModel.KitModel object:

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

and

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

I want to test the code in a machine without any GPU, so I want to convert the cuda-code into CPU version. I tried to look at some relevant posts regarding the CPU/GPU switch of PyTorch, but they are related to the usage of device and thus doesn't apply to my case.

pytorch computer-vision machine-learning image-processing gpu
2 Answers

As pointed out by kHarshit in his comment, you can simply replace .cuda() call with .cpu():

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

However, this requires changing the code in multiple places every time you want to move from GPU to CPU and vice versa.

To alleviate this difficulty, pytorch has a more "general" method .to().
You may have a device variable defining where you want pytorch to run, this device can also be the CPU (!).
for instance:

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

Once you determined once in your code where you want/can run, simply use .to() to send your model/variables there:

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

BTW,
You can use .to() to control the data type (.float()) as well:

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

PS,
Note that the Variable API has been deprecated and is no longer required.

user_17144100
user_148327070
Related