PyTorch：如何将模型库预训练模型映射至指定GPU（cuda:3）

嘿，我来帮你搞定这个问题！你现在的代码里，权重加载的逻辑没完全生效，还漏掉了一个关键步骤——把模型本身迁移到目标设备上。下面给你两种靠谱的解决方法：

方法一：用torch.device直接指定（推荐） #

这种写法更清晰，也不容易出错：

import torch
from torchvision.models.resnet import ResNet, BasicBlock
from torch.utils.model_zoo import load_url

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

# 定义目标CUDA设备
target_device = torch.device('cuda:3')

# 初始化模型
model = ResNet(BasicBlock, [3, 4, 6, 3])

# 加载预训练权重时直接指定map_location到目标设备
pretrained_weights = load_url(model_urls['resnet34'], map_location=target_device)
model.load_state_dict(pretrained_weights)

# **关键步骤**：把模型本身的参数和计算逻辑都移到目标设备
model = model.to(target_device)

方法二：保留lambda写法但补全模型迁移 #

如果你更习惯原来的lambda写法，记得最后一定要把模型迁移到cuda:3：

import torch
from torchvision.models.resnet import ResNet, BasicBlock
from torch.utils.model_zoo import load_url

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

device_num = 3
model = ResNet(BasicBlock, [3, 4, 6, 3])

# 加载权重到指定CUDA设备
pretrained_weights = load_url(model_urls['resnet34'], map_location=lambda storage, loc: storage.cuda(device_num))
model.load_state_dict(pretrained_weights)

# **关键步骤**：务必把模型整体移到目标设备
model = model.to(f'cuda:{device_num}')

额外检查点 #

运行前可以先确认cuda:3是否可用，避免设备不存在的报错：

print("CUDA是否可用：", torch.cuda.is_available())
print("可用CUDA设备数量：", torch.cuda.device_count())

为什么之前的代码总是到cuda:0？ #

大概率是因为你只处理了预训练权重的加载位置，但模型初始化的参数还留在默认设备（比如CPU或者cuda:0），必须调用model.to(device)才能把整个模型的参数和计算逻辑都迁移到cuda:3上。

内容的提问来源于stack exchange，提问作者Tsuan