Pytorch: как отобразить предварительно обученную модель зоопарка в новый графический процессор

Question

Я пытаюсь загрузить одну из предварительно обученных моделей

model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth'}

Когда я использовал следующий код, он всегда загружает модель в cuda: 0.Что если я захочу загрузить его в cuda: 3?

model = ResNet(BasicBlock, [3, 4, 6, 3]) 
device = 3
model.load_state_dict(model_zoo.load_url(model_urls['resnet34'], 
                      map_location=lambda storage, loc: storage.cuda(device)))

Answer 1

Это должно сделать работу за вас:

device = torch.device('cuda')
model = ResNet(BasicBlock, [3, 4, 6, 3]) 
with torch.cuda.device(3):
    model.load_state_dict(model_zoo.load_url(model_urls['resnet34'], 
                          map_location=lambda storage, loc: storage.cuda(device)))

Я думаю, что это работает для версии 0.4.0 и выше, вы можете проверить некоторые другие примеры в 0.4.0.руководство по миграции: https://pytorch.org/2018/04/22/0_4_0-migration-guide.html