GPU训练的Torchvision Mask R-CNN模型无GPU环境加载失败解决问询

一、处理已保存的GPU模型 #

如果你已经用torch.save(model, model_name)保存了整个模型，在无GPU机器上加载时，只需在torch.load中指定map_location参数，强制将模型参数映射到CPU即可解决设备识别失败的问题：

import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn

# 加载时指定参数映射到CPU
model = torch.load(model_name, map_location=torch.device('cpu'))
# 模型已在CPU上，可直接进入评估模式使用
model.eval()

二、最佳保存实践 #

保存前将模型迁移到CPU是合理的操作，但更推荐的是只保存模型的状态字典（state_dict），而非整个模型对象。原因是state_dict仅包含模型的参数权重，体积更小、兼容性更强，不受模型类定义或PyTorch版本的额外限制，是PyTorch官方推荐的标准保存方式。

GPU训练后的正确保存流程： #

1. 将模型迁移至CPU：

model.to('cpu')

2. 保存模型的state_dict：

torch.save(model.state_dict(), model_name)

无GPU环境的加载流程： #

1. 初始化与训练时结构完全一致的Mask R-CNN模型（注意匹配num_classes等自定义参数）：

model = maskrcnn_resnet50_fpn(pretrained=False, num_classes=你的类别数)

2. 加载state_dict并设置为评估模式：

model.load_state_dict(torch.load(model_name, map_location='cpu'))
model.eval()

补充说明 #

• 若模型包含自定义组件（如自定义检测头），加载时必须保证初始化的模型结构与训练时完全一致，否则会出现参数不匹配的错误。
• 提前将模型迁移到CPU再保存，能避免后续加载时手动指定map_location的步骤，更适合直接给无GPU环境分发模型。

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