Pytorch学习记录-使用Pytorch进行深度学习，保存和加

新建 Microsoft PowerPoint 演示文稿 (2).jpg

保存和加载模型
在完成60分钟入门之后，接下来有六节tutorials和五节关于文本处理的tutorials。争取一天一节。不过重点是关注神经网络构建和数据处理部分。

本节主要是用于解决模型的保存和加载。会的直接跳过就好。我也只是做记录，这篇搞定就直接进入NLP部分。

三个核心函数：

• torch.save：将序列化的对象保存在硬盘上，使用Python的pickle来序列化。
• torch.load：使用pickle的拆包功能将硬盘上的序列化文件导入内存中。
• torch.nn.Module.load_state_dict：加载一个模型的参数字典。

主要目录

1. 什么是state_dict?
2. Saving & Loading Model for Inference
3. Saving & Loading a General Checkpoint
4. Saving Multiple Models in One File
5. Warmstarting Model Using Parameters from a Different Model
6. 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()以确保这些图层处于训练模式。