Лучший способ сделать код Pytorch независимым от запуска на CPU или GPU? - PullRequest
Новая
       53

Лучший способ сделать код Pytorch независимым от запуска на CPU или GPU?

0 голосов
/ 02 октября 2018

Руководство по миграции рекомендует сделать код независимым от процессора / графического процессора: 

> # at beginning of the script
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# then whenever you get a new Tensor or Module
# this won't copy if they are already on the desired device
input = data.to(device)
model = MyModule(...).to(device)

Я сделал это и запустил свой код на устройстве только с процессором, но моя модель потерпела крахпри подаче входного массива, как говорится, ожидал, что тензор ЦП, а не ГП.Каким-то образом моя модель автоматически преобразовывала входной массив CPU в массив GPU.Наконец, я отследил это до этой команды в своем коде: 

model = torch.nn.DataParallel(model).to(device)

Даже если я преобразовал модель в 'CPU', nn.DataParallel переопределяет это.Лучшее решение, которое я придумал, было условным: 

if device.type=='cpu':
    model = model.to(device)
else:
    model = torch.nn.DataParallel(model).to(device)

Это не кажется элегантным.Есть ли лучший способ?

1 Ответ

0 голосов
/ 05 октября 2018

Как насчет 

if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model = model.to(device)

?

Вам не нужно DataParallel, если у вас только один графический процессор.

...