1. state_dict
在PyTorch中,state_dict 是一个简单的python的字典对象,将每一层与它的对应参数建立映射关系。(如model的每一层的weights及bias等)
首先,我们来定义一个MLP模型:
import torch.nn as nn
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()
输出:
{
'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
输出:
tensor([[1],
[1]], dtype=torch.uint8)
因为这net和net2都有同样的模型参数,那么对同一个输入X的计算结果将会是一样的。上面的输出也验证了这一点。