PyTorch读写模型（state_dict、torch.save、torch.load）

1. state_dict

在PyTorch中，state_dict 是一个简单的python的字典对象，将每一层与它的对应参数建立映射关系。(如model的每一层的weights及bias等)

首先，我们来定义一个MLP模型：

import torch.nn as nnclass MLP(nn.Module):def __init__(self):super(MLP, self).__init__()self.hidden = nn.Linear(3, 2)self.act = nn.ReLU()self.output = nn.Linear(2, 1)def forward(self, x):a = self.act(self.hidden(x))return self.output(a)net = MLP()
net.state_dict()

输出：

OrderedDict([('hidden.weight', tensor([[ 0.2448,  0.1856, -0.5678],[ 0.2030, -0.2073, -0.0104]])),('hidden.bias', tensor([-0.3117, -0.4232])),('output.weight', tensor([[-0.4556,  0.4084]])),('output.bias', tensor([-0.3573]))])

注意：只有具有可学习参数的层(卷积层、线性层等)才有state_dict中的条目。优化器(optim)也有一个state_dict，其中包含关于优化器状态以及所使用的超参数的信息。

optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
optimizer.state_dict()

输出：

{
'state': {}, 
'param_groups': [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'maximize': False, 'foreach': None, 'differentiable': False, 'params': [0, 1, 2, 3]}]
}

2. 保存和加载模型

PyTorch中保存和加载训练模型有两种常见的方法:

仅保存和加载模型参数(state_dict)；
保存和加载整个模型。

1. 保存和加载state_dict(推荐方式)
保存：

torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth

加载：

model = ModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

2. 保存和加载整个模型
保存：

torch.save(model, PATH)

加载：

model = torch.load(PATH)

我们采用推荐的方法一来实验一下:

X = torch.randn(2, 3)
Y = net(X)PATH = "./net.pt"
torch.save(net.state_dict(), PATH)net2 = MLP()
net2.load_state_dict(torch.load(PATH))
Y2 = net2(X)
Y2 == Y