Как в PyTorch преобразовать коды, связанные с cuda (), в версию процессора? - PullRequest
Новая
       62

Как в PyTorch преобразовать коды, связанные с cuda (), в версию процессора?

0 голосов
/ 27 мая 2020

У меня есть некоторые существующие коды 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())

Я хочу протестировать код на машине без процессора, поэтому я хочу преобразовать код cuda в версию процессора. Я попытался просмотреть некоторые соответствующие сообщения, касающиеся переключателя CPU / GPU в PyTorch, но они связаны с использованием device и поэтому не относятся к моему случаю.

1 Ответ

2 голосов
/ 27 мая 2020

Как указал kHarshit в его комментарии , вы можете просто заменить вызов .cuda() на .cpu(): 

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

Однако для этого потребуется изменить код в нескольких местах. каждый раз, когда вы хотите перейти с GPU на CPU и наоборот.

Чтобы облегчить эту трудность, pytorch предлагает более «общий» метод .to(). У вас может быть переменная device, определяющая, где вы хотите запустить pytorch, это device также может быть ЦП (!). например: 

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 не рекомендуется и больше не требуется.

...