2155 - 模型读写

读写模型

state_dict

在PyTorch中，Module的可学习参数(即权重和偏差)，模块模型包含在参数中(通过model.parameters()访问)。state_dict是一个从参数名称隐射到参数Tesnor的字典对象。

class 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()
输出：
{'param_groups': [{'dampening': 0,
   'lr': 0.001,
   'momentum': 0.9,
   'nesterov': False,
   'params': [4736167728, 4736166648, 4736167368, 4736165352],
   'weight_decay': 0}],
 'state': {}}

保存和加载模型

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

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

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

   #保存
   torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth
   加载：
   model = TheModelClass(*args, **kwargs)
   model.load_state_dict(torch.load(PATH))

   2. 保存和加载整个模型

   ```python 保存： 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 输出：

tensor([[1], [1]], dtype=torch.uint8)

因为这`net`和`net2`都有同样的模型参数，那么对同一个输入`X`的计算结果将会是一样的。上面的输出也验证了这一点。<br />此外，还有一些其他使用场景，例如GPU与CPU之间的模型保存与读取、使用多块GPU的模型的存储等等
问题：<br />[Pytorch训练结束后保存模型参数，接着再加载这个模型参数，继续训练但是训练效果变差很多？](https://www.zhihu.com/question/482169025/answer/2081124014)<br />原因：因为你加载后，如果仅恢复模型权重，而忽视了其他超参数等。此时的训练不能认为和中断前的是同一次训练，在恢复初期loss很可能会升高，并且有可能会错过解空间的局部最优值。<br />需要注意到，torch.save函数时可以将你想要保存的东西都存在dict中，保存到磁盘的。所以，针对这种情形，在保存的时候不要只保存模型参数，还应该保存优化器参数。参见以下代码：
```python
net = ResNet()
loss = MSE()
optimizer = torch.optim.SGD(net.parameters(), lr=0.001)
"""
正常训练的代码
"""
net.train()
for epoch in range(total_epochs):
   ...
   for x, y in dataloader:
       ...
       loss = ...
       loss.backward()
       optimizer.step()
   state_dict = {"net": net.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch}
   torch.save(state_dict, "model_path/model.pth")
"""
中断后，加载权重继续训练
"""
checkpoint = torch.load("model_path/model.pth")
current_epoch = checkpoint["epoch"]
net.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
net.train()
for epoch in range(current_epoch , total_epochs):
   ...
   for x, y in dataloader:
       ...
       loss = ...
       loss.backward()
       optimizer.step()
   state_dict = {"net": net.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch}
   torch.save(state_dict, "model_path/model.pth")