保存和加载模型
在完成60分钟入门之后，接下来有六节tutorials和五节关于文本处理的tutorials。争取一天一节。不过重点是关注神经网络构建和数据处理部分。
本节主要是用于解决模型的保存和加载。会的直接跳过就好。我也只是做记录，这篇搞定就直接进入NLP部分。
三个核心函数：
- torch.save：将序列化的对象保存在硬盘上，使用Python的pickle来序列化。
- torch.load：使用pickle的拆包功能将硬盘上的序列化文件导入内存中。
- torch.nn.Module.load_state_dict：加载一个模型的参数字典。
主要目录
- 什么是state_dict?
- Saving & Loading Model for Inference
- Saving & Loading a General Checkpoint
- Saving Multiple Models in One File
- Warmstarting Model Using Parameters from a Different Model
- Saving & Loading Model Across Devices
1.什么是state_dict?
在PyTorch中，torch.nn.Module的可学习参数(即权重和偏差)，模块模型包含在model's参数中(通过model.parameters()访问)。state_dict是个简单的Python dictionary对象，它将每个层映射到它的参数张量。
注意，只有具有可学习参数的层(卷积层、线性层等)才有model's state_dict中的条目。优化器对象(connector .optim)也有一个state_dict，其中包含关于优化器状态以及所使用的超参数的信息。
可以通过遍历模型中的state_dict每一个tensor来查看。
简单构建一个模型，这里我用了两种写法，注释掉的那一种是比较通用的，也是建议使用的。
import torch
import torch.nn as nn
import torch.nn.functional as F
# Define model
class TheModelClass(nn.Module):
def __init__(self):
super(TheModelClass, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def farward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# class TheModelClass(nn.Module):
# def __init__(self):
# super(TheModelClass, self).__init__()
# self.conv = torch.nn.Sequential()
# self.conv.add_module('conv1', nn.Conv2d(3, 6, 5))
# self.conv.add_module('pool', nn.MaxPool2d(2, 2))
# self.conv.add_module('conv2', nn.Conv2d(6, 16, 5))
# self.dense = torch.nn.Sequential()
# self.dense.add_module('fc1', nn.Linear(16 * 5 * 5, 120))
# self.dense.add_module('fc2', nn.Linear(120, 84))
# self.dense.add_module('fc3', nn.Linear(84, 10))
#
# def forward(self, x):
# conv_out = self.conv(x)
# res = conv_out.view(conv_out.size(0), -1)
# out = self.dense(res)
# return out
# Initialize model
model = TheModelClass()
# Initialize optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
print(model)
查看以下模型的state_dict
print('Model state dict:')
for param in model.state_dict():
print(param, "\t", model.state_dict()[param].size())
print('Optimizer state dict:')
for var in optimizer.state_dict():
print(var, "\t", optimizer.state_dict()[var])
TheModelClass(
(conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(fc1): Linear(in_features=400, out_features=120, bias=True)
(fc2): Linear(in_features=120, out_features=84, bias=True)
(fc3): Linear(in_features=84, out_features=10, bias=True)
)
Model state dict:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])
conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])
fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])
fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])
fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])
Optimizer state dict:
state {}
param_groups [{'lr': 0.0001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [2761682742728, 2761682742872, 2761682742944, 2761682743016, 2761682743088, 2761682743160, 2761682743232, 2761682796616, 2761682796688, 2761682796760]}]
2. 保存和加载模型的推理
2.1 保存和加载state_dict
torch.save(model.state_dict(), PATH)
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
需要注意的一点是：必须调用model.eval()，以便在运行推断之前将dropout和batch规范化层设置为评估模式。如果不这样做，将会产生不一致的推断结果。
2.2 保存和加载整个模型
torch.save(model, PATH)
model=torch.load(PATH)
model.eval()
3. 保存和加载一般检查点以进行推理和/或恢复训练
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
...
}, PATH)
加载会稍微麻烦一些
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()
# - or -
# model.train()
保存通用检查点时，用于推理或恢复训练，必须保存的不仅仅是模型的state_dict。保存优化器的state_dict也很重要，因为它包含作为模型训练更新的缓冲区和参数。您可能想要保存的其他项目是您停止使用的纪元，最新记录的训练损失，外部torch.nn.Embedding等。
要保存多个组件，请在字典中组织它们并使用torch.save()来序列化字典。常见的PyTorch约定是使用.tar文件扩展名保存这些检查点。
要加载项目，首先初始化模型和优化器，然后使用torch.load()在本地加载字典。可以通过简单地查询字典来轻松访问已保存的项目。
请记住，在运行推理之前，必须调用model.eval()将dropout和批处理规范化层设置为评估模式。如果不这样做，将导致不一致的推理结果。
如果希望恢复训练，请调用model.train()以确保这些图层处于训练模式。
