saving and loading models - 保存和读取模型

r囧r小猫 2023-02-19 03:28 34阅读 0赞

#  saving and loading models - 保存和读取模型 #

[https://pytorch.org/tutorials/beginner/saving\_loading\_models.html][https_pytorch.org_tutorials_beginner_saving_loading_models.html]  
[https://github.com/pytorch/tutorials/blob/master/beginner\_source/saving\_loading\_models.py][https_github.com_pytorch_tutorials_blob_master_beginner_source_saving_loading_models.py]  
[https://pytorch.org/tutorials/index.html][https_pytorch.org_tutorials_index.html]

This document provides solutions to a variety of use cases regarding the saving and loading of PyTorch models. Feel free to read the whole document, or just skip to the code you need for a desired use case.  
本文档为有关保存和加载 PyTorch 模型的各种用例提供了解决方案。随意阅读整个文档，或者只是跳到所需用例所需的代码。

When it comes to saving and loading models, there are three core functions to be familiar with:  
关于保存和加载模型，有三个核心功能需要熟悉：

*  `torch.save`: Saves a serialized object to disk. This function uses Python’s pickle utility for serialization. Models, tensors, and dictionaries of all kinds of objects can be saved using this function.  
    保存一个序列化对象 (serialized object) 到磁盘中。该函数使用 Python 的 pickle 工具完成序列化。Models, tensors, 以及由各种对象所组成的字典数据都可以通过该函数进行保存。
 *  `torch.load`: Uses pickle’s unpickling facilities to deserialize pickled object files to memory. This function also facilitates the device to load the data into (see `Saving & Loading Model Across Devices`).  
    使用 pickle 的解包工具 (unpickling facilities) 来反序列化 pickled object 文件到内存中。该函数同样可以操作设备 (device) 来加载数据 (see `Saving & Loading Model Across Devices`)。
 *  `torch.nn.Module.load_state_dict`: Loads a model’s parameter dictionary using a deserialized `state_dict`. For more information on `state_dict`, see `What is a state_dict?`.  
    使用反序列化的 `state_dict` 加载模型的参数字典。

## 1. What is a `state_dict`? ##

In PyTorch, the learnable parameters (i.e. `weights` and `biases`) of an `torch.nn.Module` model are contained in the model’s parameters (accessed with `model.parameters()`). A `state_dict` is simply a Python dictionary object that maps each layer to its parameter tensor. Note that only layers with learnable parameters (convolutional layers, linear layers, etc.) and registered buffers (batchnorm’s `running_mean`) have entries in the model’s `state_dict`. Optimizer objects (`torch.optim`) also have a `state_dict`, which contains information about the optimizer’s state, as well as the hyperparameters used.  
在 PyTorch 中，`torch.nn.Module` 模型中的可学习参数 (`weighs` and `biases`) 保存在模型参数中 (accessed with `model.parameters()`)。而 `state_dict` 是一个典型的 Python 字典对象，它将每个层映射到它的参数张量。注意，只有具有可学习参数的层 (convolutional layers, linear layers, etc.) 才会被保存在模型的 `state_dict` 数据结构中。优化器对象 (Optimizer objects `torch.optim`) 也有一个 `state_dict`，其中包含了优化器状态以及所使用的超参数的信息。

Because `state_dict` objects are Python dictionaries, they can be easily saved, updated, altered, and restored, adding a great deal of modularity to PyTorch models and optimizers.  
由于 `state_dict` 对象是 Python 字典，因此可以轻松地保存、更新、更改和还原它们，从而为 PyTorch 模型和优化器增加了很多模块化。

modularity /,mɒdjʊ'lærɪtɪ/：n. 模块性

Let’s take a look at the `state_dict` from the simple model used in the Training a classifier tutorial.  
[https://pytorch.org/tutorials/beginner/blitz/cifar10\_tutorial.html][https_pytorch.org_tutorials_beginner_blitz_cifar10_tutorial.html]

# 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])

#!/usr/bin/env python
    # -*- coding: utf-8 -*-
    # yongqiang cheng
    
    from __future__ import absolute_import
    from __future__ import division
    from __future__ import print_function
    
    import torch.nn as nn
    import torch.nn.functional as F
    import torch.optim as optim
    
    
    # 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])

Output:

/home/yongqiang/miniconda3/envs/pt-1.4_py-3.6/bin/python /home/yongqiang/pytorch_work/end2end-asr-pytorch-example/yongqiang.py
    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': [139998270550880, 139998249461728, 139998249461800, 139998249461872, 139998249461944, 139998249462016, 139998249462088, 139998249462160, 139998249462232, 139998249462304]}]
    
    Process finished with exit code 0

## 2. Saving & Loading Model for Inference ##

### 2.1 Save/Load `state_dict` (Recommended) ###

**Save:**

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

**Load:**

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

When saving a model for inference, it is only necessary to save the trained model’s learned parameters. Saving the model’s `state_dict` with the `torch.save()` function will give you the most flexibility for restoring the model later, which is why it is the recommended method for saving models.  
当为 inference 阶段保存模型时，仅仅保存训练好的模型的可更新参数即可。利用 `torch.save()` 函数来保存模型的`state_dict` 可以在之后恢复模型时提供极大的灵活性，这也是我们推荐使用该方法来保存模型的原因。

A common PyTorch convention is to save models using either a `.pt` or `.pth` file extension.  
常见的 PyTorch 约定是使用 `.pt` 或 `.pth` 文件扩展名保存模型。

Remember that you must call `model.eval()` to set `dropout` and `batch normalization` layers to evaluation mode before running inference. Failing to do this will yield inconsistent inference results.  
请记住，在运行推理之前，必须调用 `model.eval()` 将 `dropout` and `batch normalization` 层设置为评估模式。不这样做将产生不一致的推断结果。

Notice that the `load_state_dict()` function takes a dictionary object, NOT a path to a saved object. This means that you must deserialize the saved `state_dict` before you pass it to the `load_state_dict()` function. For example, you CANNOT load using `model.load_state_dict(PATH)`.  
请注意，`load_state_dict()` 函数接受的参数是一个字典对象, 而不是模型文件的保存路径。这意味着你必须先将模型文件解序列成字典以后，才能将其传给 `load_state_dict()` 函数。例如，您不能使用 `model.load_state_dict(PATH)` 加载。

### 2.2 Save/Load Entire Model ###

**Save:**

torch.save(model, PATH)

**Load:**

# Model class must be defined somewhere
    model = torch.load(PATH)
    model.eval()

This save/load process uses the most intuitive syntax and involves the least amount of code. Saving a model in this way will save the entire module using Python’s pickle module. The disadvantage of this approach is that the serialized data is bound to the specific classes and the exact directory structure used when the model is saved. The reason for this is because pickle does not save the model class itself. Rather, it saves a path to the file containing the class, which is used during load time. Because of this, your code can break in various ways when used in other projects or after refactors.  
此保存/加载过程使用最直观的语法，并且涉及最少的代码。以这种方式保存模型将使用 Python 的 pickle 模块保存整个模块。这种方法的缺点是序列化的数据绑定到特定的类，并且在保存模型时使用确切的目录结构。这样做的原因是因为 pickle 不会保存模型类本身。而是将其保存到包含类的文件的路径，该路径在加载时使用。因此，在其他项目中使用或重构后，您的代码可能会以各种方式崩溃。

A common PyTorch convention is to save models using either a `.pt` or `.pth` file extension.  
常见的 PyTorch 约定是使用 `.pt` 或 `.pth` 文件扩展名保存模型。

## 3. Saving & Loading a General Checkpoint for Inference and/or Resuming Training ##

**Save:**

torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': loss,
                ...
                }, PATH)

**Load:**

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()

When saving a general checkpoint, to be used for either inference or resuming training, you must save more than just the model’s `state_dict`. It is important to also save the optimizer’s `state_dict`, as this contains buffers and parameters that are updated as the model trains. Other items that you may want to save are the epoch you left off on, the latest recorded training loss, external `torch.nn.Embedding` layers, etc.  
保存用于检查或继续训练的常规检查点时，您必须保存的不只是模型的 `state_dict`。保存优化器的 `state_dict` 也很重要，因为它包含随着模型训练而更新的缓冲区和参数。您可能要保存的其他项包括您停止时的 epoch，最新记录的训练损失，外部 `torch.nn.Embedding` 层等。

To save multiple components, organize them in a dictionary and use `torch.save()` to serialize the dictionary. A common PyTorch convention is to save these checkpoints using the `.tar` file extension.  
要保存多个部分，请将它们组织在字典中，然后使用 `torch.save()` 序列化字典。常见的 PyTorch 约定是使用 `.tar` 文件扩展名保存这些检查点。

To load the items, first initialize the model and optimizer, then load the dictionary locally using torch.load(). From here, you can easily access the saved items by simply querying the dictionary as you would expect.  
要加载项，请首先初始化模型和优化器，然后使用 `torch.load()` 在本地加载字典。从这里，您只需按期望查询字典即可轻松访问已保存的项目。

Remember that you must call `model.eval()` to set dropout and batch normalization layers to evaluation mode before running inference. Failing to do this will yield inconsistent inference results. If you wish to resuming training, call `model.train()` to ensure these layers are in training mode.  
请记住，在运行推理之前，必须调用 `model.eval()` 来将 dropout and batch normalization 层设置为评估模式。不这样做将产生不一致的推断结果。如果您希望恢复训练，请调用 `model.train()` 以确保这些层处于训练模式。

## 4. Saving Multiple Models in One File ##

**Save:**

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)

**Load:**

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()
    # - or -
    modelA.train()
    modelB.train()

When saving a model comprised of multiple `torch.nn.Modules`, such as a GAN, a sequence-to-sequence model, or an ensemble of models, you follow the same approach as when you are saving a general checkpoint. In other words, save a dictionary of each model’s `state_dict` and corresponding optimizer. As mentioned before, you can save any other items that may aid you in resuming training by simply appending them to the dictionary.  
保存由多个 `torch.nn.Modules` 组成的模型 (例如 GAN，序列到序列模型或模型集合) 时，您将采用与保存常规检查点相同的方法。换句话说，保存每个模型的 `state_dict` 和相应优化器的字典。如前所述，您可以保存任何其他可以帮助您恢复训练的项，只需将它们添加到字典中即可。

A common PyTorch convention is to save these checkpoints using the `.tar` file extension.  
常见的 PyTorch 约定是使用 `.tar` 文件扩展名保存这些检查点。

To load the models, first initialize the models and optimizers, then load the dictionary locally using `torch.load()`. From here, you can easily access the saved items by simply querying the dictionary as you would expect.  
要加载模型，请首先初始化模型和优化器，然后使用 `torch.load()` 在本地加载字典。从这里，您只需按期望查询字典即可轻松访问已保存的项目。

Remember that you must call `model.eval()` to set dropout and batch normalization layers to evaluation mode before running inference. Failing to do this will yield inconsistent inference results. If you wish to resuming training, call `model.train()` to set these layers to training mode.  
请记住，在运行推理之前，必须调用 `model.eval()` 来将 dropout and batch normalization 层设置为评估模式。不这样做将产生不一致的推断结果。如果您希望恢复训练，请调用 `model.train()` 以确保这些层处于训练模式。

## 5. Warmstarting Model Using Parameters from a Different Model (热启动) ##

**Save:**

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

**Load:**

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

Partially loading a model or loading a partial model are common scenarios when transfer learning or training a new complex model. Leveraging trained parameters, even if only a few are usable, will help to warmstart the training process and hopefully help your model converge much faster than training from scratch.  
在迁移学习或训练新的复杂模型时，部分加载模型或加载部分模型是常见方案。利用经过训练的参数，即使只有少数几个可用的参数，也将有助于热启动训练过程，并希望与从头开始训练相比，可以更快地收敛模型。

Whether you are loading from a partial `state_dict`, which is missing some keys, or loading a `state_dict` with more keys than the model that you are loading into, you can set the `strict` argument to False in the `load_state_dict()` function to ignore non-matching keys.  
无论是从缺少某些键的部分 `state_dict` 加载，还是要使用比要加载的模型更多的键加载 `state_dict`，可以在`load_state_dict()` 中将 `strict` 参数设置为 False。忽略不匹配键的功能。

If you want to load parameters from one layer to another, but some keys do not match, simply change the name of the parameter keys in the state\_dict that you are loading to match the keys in the model that you are loading into.  
如果要将参数从一层加载到另一层，但是某些键不匹配，只需更改要加载的 `state_dict` 中参数键的名称，以匹配要加载到的模型中的键。

## 6. Saving & Loading Model Across Devices ##

### 6.1 Save on GPU, Load on CPU ###

**Save:**

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

Load:

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

When loading a model on a CPU that was trained with a GPU, pass `torch.device('cpu')` to the `map_location` argument in the `torch.load()` function. In this case, the storages underlying the tensors are dynamically remapped to the CPU device using the `map_location` argument.  
在 CPU 上加载 GPU 训练的模型时，将 `torch.device('cpu')` 传递到 `torch.load()` 函数中的 `map_location` 参数。在这种情况下，使用 `map_location` 参数将张量下面的存储动态地重新映射到 CPU 设备。

### 6.2 Save on GPU, Load on GPU ###

**Save:**

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

**Load:**

device = torch.device("cuda")
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH))
    model.to(device)
    # Make sure to call input = input.to(device) on any input tensors that you feed to the model

When loading a model on a GPU that was trained and saved on GPU, simply convert the initialized `model` to a CUDA optimized model using `model.to(torch.device('cuda'))`. Also, be sure to use the `.to(torch.device('cuda'))` function on all model inputs to prepare the data for the model. Note that calling `my_tensor.to(device)` returns a new copy of `my_tensor` on GPU. It does NOT overwrite `my_tensor`. Therefore, remember to manually overwrite tensors: `my_tensor = my_tensor.to(torch.device('cuda'))`.  
在 GPU 上加载 GPU 训练保存的模型时，只需使用 `model.to(torch.device('cuda'))` 将已初始化的 `model` 转换为 CUDA 优化模型即可。同样，请确保在所有模型输入上使用 `.to(torch.device('cuda'))` 函数来为模型准备数据。请注意，调用 `my_tensor.to(device)` 会在 GPU 上返回 `my_tensor` 的新副本。它不会覆盖 `my_tensor`。因此，请记住手动覆盖张量：`my_tensor = my_tensor.to(torch.device('cuda'))`。

`mytensor.to(device)` 实际上是在 GPU 中创建了 `mytensor` 的副本, 而并没有改变 `mytensor` 的值, 因此, 需要写成后面的形式来使的 `mytensor` 的值改变：`my_tensor = my_tensor.to(torch.device('cuda'))`。

### 6.3 Save on CPU, Load on GPU ###

**Save:**

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

**Load:**

device = torch.device("cuda")
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
    model.to(device)
    # Make sure to call input = input.to(device) on any input tensors that you feed to the model

When loading a model on a GPU that was trained and saved on CPU, set the `map_location` argument in the `torch.load()` function to `cuda:device_id`. This loads the model to a given GPU device. Next, be sure to call `model.to(torch.device('cuda'))` to convert the model’s parameter tensors to CUDA tensors. Finally, be sure to use the `.to(torch.device('cuda'))` function on all model inputs to prepare the data for the CUDA optimized model. Note that calling `my_tensor.to(device)` returns a new copy of `my_tensor` on GPU. It does NOT overwrite `my_tensor`. Therefore, remember to manually overwrite tensors: `my_tensor = my_tensor.to(torch.device('cuda'))`.  
在 GPU 上加载 CPU 训练保存的模型时，请将 `torch.load()` 函数中的 `map_location` 参数设置为 `cuda:device_id`。这会将模型加载到指定的 GPU 设备。接下来，请确保调用 `model.to(torch.device('cuda'))` 将模型的参数张量转换为CUDA张量。最后，确保在所有模型输入上使用 `.to(torch.device('cuda'))` 函数为 CUDA 优化模型准备数据。请注意，调用 `my_tensor.to(device)` 会在 GPU 上返回 `my_tensor` 的新副本。它不会覆盖 `my_tensor`。因此，请记住手动覆盖张量：`my_tensor = my_tensor.to(torch.device('cuda'))`。

## 7. Saving `torch.nn.DataParallel` Models ##

**Save:**

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

**Load:**

# Load to whatever device you want

`torch.nn.DataParallel` is a model wrapper that enables parallel GPU utilization. To save a `DataParallel` model generically, save the `model.module.state_dict()`. This way, you have the flexibility to load the model any way you want to any device you want.  
`torch.nn.DataParallel` 是支持模型使用 GPU 并行的封装器。要保存一个一般的 `DataParallel` 模型，请保存 `model.module.state_dict()`。这种方式，可以灵活地以任何方式加载模型到任何设备上。

[https_pytorch.org_tutorials_beginner_saving_loading_models.html]: https://pytorch.org/tutorials/beginner/saving_loading_models.html
[https_github.com_pytorch_tutorials_blob_master_beginner_source_saving_loading_models.py]: https://github.com/pytorch/tutorials/blob/master/beginner_source/saving_loading_models.py
[https_pytorch.org_tutorials_index.html]: https://pytorch.org/tutorials/index.html
[https_pytorch.org_tutorials_beginner_blitz_cifar10_tutorial.html]: https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html