Pythoch 中的state_dict() 和 load_state_dict()

• Pytorch 中一种模型保存和加载的方式, 其中model.state_dict()其实返回的是一个OrderDict，存储了网络结构的名字和对应的参数

import torch

# 保存自己训练好的模型
torch.save(model.state_dict(), path)

# 将训练好的模型加载出来
model = mymodelf(*args, **kwargs)
model.load_state_dict(torch.load(path))
model.eval()

• load_state_dict() 中还有一个关键的参数 strict, 当strict=True,要求预训练权重层数的键值与新构建的模型中的权重层数名称完全吻合；如果新构建的模型在层数上进行了部分微调，则上述代码就会报错：key对应不上, 此时，采用strict=False 就能够解决这个问题

import torch
import torchvision

net = torchvision.models.resnet50(pretrained = True)
net.load_state_dict(torch.load(path), strict = True)