新建 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 torchimport torch.nn as nnimport torch.nn.functional as F# Define modelclass 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 modelmodel = TheModelClass()# Initialize optimizeroptimizer = 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,