在加载预训练模型的时候,由于用DataParallel训练的模型数据并行方式的,key中会包含”module“关键字,故会出现以下错误:
RuntimeError: Error(s) in loading state_dict for ResNet: Missing key(s) in state_dict……
去掉DataParallel 预训练模型中的module,修改如下:
#model.load_state_dict(checkpoint['state_dict'])
model.load_state_dict({k.replace('module.',''):v for k,v in checkpoint['state_dict'].items()})