PyTorch模型保存与加载

Bertha 。 2024-04-17 16:16 92阅读 0赞

1.  `torch.save`：保存序列化的对象到磁盘，使用了Python的pickle进行序列化，模型、张量、所有对象的字典。
2.  `torch.load`：使用了pickle的unpacking将pickled的对象反序列化到内存中。
3.  `torch.nn.Module.load_state_dict`：使用反序列化的state\_dict加载模型的参数字典。

state\_dict 是一个Python字典，将每一层映射成它的参数张量。注意只有带有可学习参数的层（卷积层、全连接层等），以及注册的缓存（batchnorm的运行平均值）在state\_dict 中才有记录。state\_dict同样包含优化器对象，存储了优化器的状态，所使用到的超参数。

## 一个简单的例子 ##

# 定义模型
    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
    
    # 初始化模型
    model = TheModelClass()
    
    # 初始化优化器
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    
    # 打印模型的 state_dict
    print("Model's state_dict:")
    for param_tensor in model.state_dict():
        print(param_tensor, "\t", model.state_dict()[param_tensor].size())
    
    # 打印优化器的 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（推荐） ##

保存：

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

加载：

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

要注意这个细节，如果使用`nn.DataParallel`在一台电脑上使用了多个GPU，那么加载模型的时候也必须先进行nn.DataParallel。

保存模型的推理过程的时候，只需要保存模型训练好的参数，使用`torch.save()`保存`state_dict`，能够方便模型的加载。因此推荐使用这种方式进行模型保存。

记住一定要使用`model.eval()`来固定`dropout`和归一化层，否则每次推理会生成不同的结果。

>     注意，load_state_dict()需要传入字典对象，因此需要先反序列化state_dict再传入load_state_dict()

## 保存/加载整个模型 ##

保存：

torch.save(model, PATH)

加载：

# 模型类必须在别的地方定义
    model = torch.load(PATH)
    model.eval()

这种保存/加载模型的过程使用了最直观的语法，所用代码量少。这使用`Python`的`pickle`保存所有模块。这种方法的缺点是，保存模型的时候，序列化的数据被绑定到了特定的类和确切的目录。这是因为`pickle`不保存模型类本身，而是保存这个类的路径，并且在加载的时候会使用。因此，当在其他项目里使用或者重构的时候，加载模型的时候会出错。

一般来说，`PyTorch`的模型以.pt或者.pth文件格式保存。

一定要记住在评估模式的时候调用`model.eval()`来固定`dropout`和批次归一化。否则会产生不一致的推理结果。

## 保存加载用于推理的常规Checkpoint/或继续训练 ##

保存：

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()
    # - 或者 -
    model.train()

在保存用于推理或者继续训练的常规检查点的时候，除了模型的`state_dict`之外，还必须保存其他参数。保存优化器的`state_dict`也非常重要，因为它包含了模型在训练时候优化器的缓存和参数。除此之外，还可以保存停止训练时`epoch`数，最新的模型损失，额外的`torch.nn.Embedding`层等。

要保存多个组件，则将它们放到一个字典中，然后使用`torch.save()`序列化这个字典。一般来说，使用.tar文件格式来保存这些检查点。

加载各个组件，首先初始化模型和优化器，然后使用`torch.load()`加载保存的字典，然后可以直接查询字典中的值来获取保存的组件。

同样，评估模型的时候一定不要忘了调用`model.eval()`。

## 保存多个模型到一个文件 ##

保存：

torch.save({
                'modelA_state_dict': modelA.state_dict(),
                'modelB_state_dict': modelB.state_dict(),
                'optimizerA_state_dict': optimizerA.state_dict(),
                'optimizerB_state_dict': optimizerB.state_dict(),
                ...
                }, PATH)

加载：

modelA = TheModelAClass(*args, **kwargs)
    modelB = TheModelBClass(*args, **kwargs)
    optimizerA = TheOptimizerAClass(*args, **kwargs)
    optimizerB = TheOptimizerBClass(*args, **kwargs)
    
    checkpoint = torch.load(PATH)
    modelA.load_state_dict(checkpoint['modelA_state_dict'])
    modelB.load_state_dict(checkpoint['modelB_state_dict'])
    optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
    optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])
    
    modelA.eval()
    modelB.eval()
    # - 或者 -
    modelA.train()
    modelB.train()

保存的模型包含多个`torch.nn.Modules`时，比如GAN，一个序列-序列模型，或者组合模型，使用与保存常规检查点的方式来保存模型。也就是说，保存每个模型的state\_dict和对应的优化器到一个字典中。我们可以保存任何能帮助我们继续训练的东西到这个字典中。

## 使用其他模型来预热当前模型 ##

保存：

torch.save(modelA.state_dict(), PATH)

加载：

modelB = TheModelBClass(*args, **kwargs)
    modelB.load_state_dict(torch.load(PATH), strict=False)

在迁移学习或者训练新的复杂模型时，加载部分模型是很常见的。利用经过训练的参数，即使只有少数参数可用，也将有助于预热训练过程，并且使模型更快收敛。

在加载部分模型参数进行预训练的时候，很可能会碰到键不匹配的情况（模型权重都是按键值对的形式保存并加载回来的）。因此，无论是缺少键还是多出键的情况，都可以通过在load\_state\_dict()函数中设定strict参数为False来忽略不匹配的键。

如果想将某一层的参数加载到其他层，但是有些键不匹配，那么修改state\_dict中参数的key可以解决这个问题。

## 跨设备保存与加载模型 ##

GPU上保存，CPU上加载

保存：

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

加载：

device = torch.device('cpu')
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH, map_location=device))

当在CPU上加载一个GPU上训练的模型时，在torch.load()中指定map\_location=torch.device(‘cpu’)，此时，map\_location动态地将tensors的底层存储重新映射到CPU设备上。

上述代码只有在模型是在一块GPU上训练时才有效，如果模型在多个GPU上训练，那么在CPU上加载时，会得到类似如下错误：

KeyError: ‘unexpected key “module.conv1.weight” in state_dict’

原因是在使用多GPU训练并保存模型时，模型的参数名都带上了module前缀，因此可以在加载模型时，把key中的这个前缀去掉：

# 原始通过DataParallel保存的文件
    state_dict = torch.load('myfile.pth.tar')
    # 创建一个不包含`module.`的新OrderedDict
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:] # 去掉 `module.`
        new_state_dict[name] = v
    # 加载参数
    model.load_state_dict(new_state_dict)

GPU上保存，GPU上加载

保存：

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

加载：

device = torch.device("cuda")
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH))
    model.to(device)
    
    # 往模型中输入数据的时候不要忘记在任意tensor上调用input = input.to(device)

在把GPU上训练的模型加载到GPU上时，只需要使用model.to(torch.devie(‘cuda’))将初始化的模型转换为CUDA优化模型。同时确保在模型所有的输入上使用.to(torch.device(‘cuda’))。注意，调用my\_tensor.to(device)会返回一份在GPU上的my\_tensor的拷贝。不会覆盖原本的my\_tensor，因此要记得手动将tensor重写：my\_tensor = my\_tensor.to(torch.device(‘cuda’))。

CPU上保存，GPU上加载

保存：

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

加载：

device = torch.device("cuda")
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # 选择希望使用的GPU
    model.to(device)

## 保存torch.nn.DataParallel模型 ##

保存：

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