pytorch模型保存与加载

模型保存相关的三个核心功能

torch.save: 将序列化对象保存到磁盘。此函数使用Python的pickle模块进行序列化，使用此模型可以保存如模型、tensor、字典等各种对象。
torch.load: 使用pickle的unpicking功能将pickle对象文件反序列化到内存。此功能还可以有助于设备加载数据。
torch.nn.Moudle.load_state_dict: 使用反序列化函数state_dict来加载模型的参数字典。

状态字典

在pytorch中，torch.nn.Module模型的可学习参数（即权重和偏差）包含在模型的parameters中，（使用model.parameters()可以进行访问）。state_dict仅仅是python字典对象，它将每一层映射到其参数张量。注意，只有具有可学习参数的层（如卷积层、线性层等）的模型才具有state_dict这一项。优化目标torch.optim也有state_dict属性，它包含有关优化器的状态信息，以及使用的超参数。

示例

# 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 forward(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
# Initialize model
model = TheModelClass()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# Print optimizer's state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

输出：

Model's 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's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

保存和加载推断模型

保存/加载`state_dict`(推荐使用)

保存：

1	torch.save(model.state_dict(), PATH)

加载:

1
2
3

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

用保存的模型进行推断的时候，只需要保存模型学习到的参数，使用torch.save()函数来保存模型state_dict，所用的资源要少于保存完整模型。在进行推断之前，要调用model.eval()去设置dropout和batch normalization层为评估模式。在传入load_state_dict()函数之前，需要使用torch.load()对state_dict进行反序列化。

保存/加载完整模型

保存：

1	torch.save(model, PATH)

加载：

1 2	model = torch.load(PATH) model.eval()

保存`torch.nn.DataParallel`模型

保存：

1
2
3

model = TheModelClass(*args, **kwargs)
model = torch.nn.DataParallel(model)
torch.save(model.state_dict(), PATH)

加载：

model = TheModelClass(*args, **kwargs)
model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load(PATH))
model.eval()

在加载模型继续训练的时候，加载了两次torch.nn.DataParallel，保存的模型进行推断也需要加载两次才能进行推断。可以通过以下方法将保存的模型转化为非DataParallel模式的模型（所有key的名字前去掉modules）

from collections import OrderedDict
from efficientnet import efficientnet_b0b
import torch.nn as nn
import torch

model_path = 'Result/efficientnet/07-22_13-15-51/1net_params.pkl'
state_dict = torch.load(model_path)
new_state_dict = OrderedDict()

for k, v in state_dict.items():
    name = k[7:]
    new_state_dict[name] = v

two_state_dict = OrderedDict()

for k, v in new_state_dict.items():
    name = k[7:]
    two_state_dict[name] = v

net = efficientnet_b0b((224, 224), num_classes=1852)
net = nn.DataParallel(net)
net.load_state_dict(new_state_dict)

模型保存相关的三个核心功能

状态字典

示例

保存和加载推断模型

保存/加载state_dict(推荐使用)

保存/加载完整模型

保存torch.nn.DataParallel模型

保存/加载`state_dict`(推荐使用)

保存`torch.nn.DataParallel`模型