时间戳

首页 > 解决方案 > Pytorch：如何将模型动物园预训练模型映射到新 GPU

pytorch - Pytorch：如何将模型动物园预训练模型映射到新 GPU

问题描述

我正在尝试加载其中一个预训练模型

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)))

标签： pytorch

解决方案

这应该为您完成工作：

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

推荐阅读