训练Faster RCNN时已调用.to(device)仍报张量跨设备错误如何解决

问题根因 #

你的模型权重存在跨设备的核心原因是模型修改顺序错误：

• 你先调用了model.to(device)将原始预训练结构的参数移到GPU
• 之后才替换了model.roi_heads.box_predictor为新初始化的FastRCNNPredictor，这个新生成的预测头参数默认是在CPU上，没有被迁移到GPU，前向传播时就出现了CPU/GPU张量不匹配的报错。

可行解决方案 #

方案1：调整模型初始化和迁移的顺序 #

把修改模型预测头的代码放到model.to(device)之前执行，保证所有模型参数初始化完成后再统一迁移到GPU，修改后的代码片段如下：

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
# 先修改模型结构
num_classes = 3 # eyelids, iris + background
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# 再统一把所有参数迁移到GPU
model.to(device)

方案2：修改完预测头后再调用一次to(device) #

如果你不想调整原有代码顺序，也可以在替换完预测头后再执行一次迁移操作：

model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
model.to(device)

额外排查点（可选） #

如果修改后仍然报错，可以检查以下场景：

• 确认你的FastRCNNPredictor类内部没有硬编码指定CPU生成张量
• 暂时把dataloader的num_workers设为0测试是否还报错，排除多进程加载时的设备问题
• 确认训练循环里的图像、target张量迁移逻辑没有遗漏字段，你当前的写法是正确的

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