如何在PyTorch中加载MMFashion预训练模型并解决相关报错

报错原因 #

1. TypeError: 'strict' is an invalid keyword argument for load()：strict是model.load_state_dict()的专属参数，被错误放置在torch.load()的传参中。
2. 权重key不匹配错误：MMFashion的预训练权重文件为完整训练 checkpoint，包含meta、state_dict、optimizer三个顶级字段，直接将整个checkpoint传入load_state_dict无法匹配模型的参数key。

操作步骤 #

• 从torch.load返回的checkpoint中提取state_dict字段，该字段存储实际的模型参数
• 将strict=False参数放在model.load_state_dict的传参中，适配MMFashion预训练模型和torchvision原生ResNet50的结构差异
• 若需用于服装属性预测，需根据实际属性数量替换ResNet50末尾的全连接层

正确代码示例 #

import io
import torch
import torch.nn as nn
from torchvision import models
from PIL import Image
import torchvision.transforms as transforms

checkpoint_path = 'C:/venvs/ai/aiproduct/latest.pth'
checkpoint = torch.load(checkpoint_path, map_location='cpu')
pretrained_weights = checkpoint['state_dict']

model = models.resnet50(pretrained=False)
# 如需适配服装属性预测任务，取消注释下方代码并修改属性数量
# num_attrs = 48 # 可替换为你需要的属性类别数
# model.fc = nn.Linear(model.fc.in_features, num_attrs)

model.load_state_dict(pretrained_weights, strict=False)
model.eval()
print(model)

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