我使用PyTorch训练了一个模型并保存了一个state dict文件。我使用下面的代码加载了预训练模型。我收到了关于RuntimeError: 在加载VGG的state_dict时出现错误的消息：
RuntimeError: 在加载VGG的state_dict时出现错误： state_dict中缺少的键："features.0.weight", "features.0.bias", "features.2.weight", "features.2.bias", "features.5.weight", "features.5.bias", "features.7.weight", "features.7.bias", "features.10.weight", "features.10.bias", "features.12.weight", "features.12.bias", "features.14.weight", "features.14.bias", "features.17.weight", "features.17.bias", "features.19.weight", "features.19.bias", "features.21.weight", "features.21.bias", "features.24.weight", "features.24.bias", "features.26.weight", "features.26.bias", "features.28.weight", "features.28.bias", "classifier.0.weight", "classifier.0.bias", "classifier.3.weight", "classifier.3.bias", "classifier.6.weight", "classifier.6.bias"。 state_dict中意外的键："state_dict", "optimizer_state_dict", "globalStep", "train_paths", "test_paths"。 我正在按照这个网站上的说明进行操作： https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-across-devices
非常感谢
import argparseimport datetimeimport globimport osimport randomimport shutilimport timefrom os.path import joinimport numpy as npimport pandas as pdimport torchimport torch.nn as nnfrom torch.utils.data import DataLoaderfrom torch.utils.tensorboard import SummaryWriterfrom torchvision.transforms import ToTensorfrom tqdm import tqdmimport torch.optim as optimfrom convnet3 import Convnetfrom dataset2 import CellsDatasetfrom convnet3 import Convnetfrom VGG import VGGfrom dataset2 import CellsDatasetfrom torchvision import modelsfrom Conv import Conv2dparser = argparse.ArgumentParser('预测像素中的命中')parser.add_argument('name',type=str,help='实验名称')parser.add_argument('data_dir',type=str,help='包含图像和gt.csv的数据目录路径')parser.add_argument('--weight_decay',type=float,default=0.0,help='权重衰减系数（如10^-5）')parser.add_argument('--lr',type=float,default=0.0001,help='学习率')args = parser.parse_args()metadata = pd.read_csv(join(args.data_dir,'gt.csv'))metadata.set_index('filename', inplace=True)# 创建数据集：dataset = CellsDataset(args.data_dir,transform=ToTensor(),return_filenames=True)dataset = DataLoader(dataset,num_workers=4,pin_memory=True)model_path = '/Users/nubstech/Documents/GitHub/CellCountingDirectCount/VGG_model_V1/checkpoints/checkpoint.pth'class VGG(nn.Module): def __init__(self, pretrained=True): super(VGG, self).__init__() vgg = models.vgg16(pretrained=pretrained) # if pretrained: vgg.load_state_dict(torch.load(model_path)) features = list(vgg.features.children()) self.features4 = nn.Sequential(*features[0:23]) self.de_pred = nn.Sequential(Conv2d(512, 128, 1, same_padding=True, NL='relu'), Conv2d(128, 1, 1, same_padding=True, NL='relu')) def forward(self, x): x = self.features4(x) x = self.de_pred(x) return xmodel=VGG()#model.load_state_dict(torch.load(model_path),strict=False)model.eval() #optimizer = torch.optim.Adam(model.parameters(),lr=args.lr,weight_decay=args.weight_decay)for images, paths in tqdm(dataset): targets = torch.tensor([metadata['count'][os.path.split(path)[-1]] for path in paths]) # B targets = targets.float() # 代码用于将训练数据打印到csv文件 #filename=CellsDataset(args.data_dir,transform=ToTensor(),return_filenames=True) output = model(images) # B x 1 x 9 x 9 (类似于热图) preds = output.sum(dim=[1,2,3]) # 预测的细胞计数（长度为B的向量） print(preds) paths_test = np.array([paths]) names_preds = np.hstack(paths) print(names_preds) df=pd.DataFrame({'Image_Name':names_preds, 'Target':targets.detach(), 'Prediction':preds.detach()}) print(df) # 保存图像名称、目标和预测 df.to_csv(r'model.csv', index=False, mode='a')保存state dict的代码
torch.save({'state_dict':model.state_dict(), 'optimizer_state_dict':optimizer.state_dict(), 'globalStep':global_step, 'train_paths':dataset_train.files, 'test_paths':dataset_test.files},checkpoint_path)回答：
问题在于保存的内容与预期加载的内容不一致。代码尝试加载仅一个state_dict；它保存的内容比这多得多——看起来像是一个包含其他信息的字典内的state_dict。加载方法没有逻辑去查看字典内部。
这应该可以工作：
import torch, torchvision.modelsmodel = torchvision.models.vgg16()path = 'test.pth'torch.save(model.state_dict(), path) # 这里没有其他内容model.load_state_dict(torch.load(path))