Pytorch预训练模型导入、查看模型参数及模型保存

---

一、预训练模型权重导入

1.普通用法

模型的权重文件一般以.pth为后缀，当我们想导入模型权重时，首先需要是用torch.load函数读入权重文件，获得预训练权重的键值对，然后使用model.load_state_dict函数为模型读入权重 。

# 创建模型实例
model = MyModel()

# 读取预训练权重文件
pretrained_weights_path = 'a.pth'
pretrained_weights = torch.load(pretrained_weights_path)

# 加载权重到模型
# 如果模型定义和权重完全匹配，使用strict=True
model.load_state_dict(pretrained_weights, strict=True)

在load_state_dict函数中，有一个参数为strict，这个单词的中文意思为“严格”，当strict=True时，只有模型的定义结构和权重的结构完全相同时才会成功导入，否则便会报错。当strict=False时，函数会自动忽略不匹配的键，只会导入匹配的键，并且不会报错。因此如果我们有一个预训练权重文件，他和我们设计的模型结构不是完全一样，但是我们仍然想要结构一样的层的参数，这时我们需要把strict设置为False.

2.只导入部分预训练权重

如果模型包含编码器解码器部分，我们只想要为编码器部分导入预训练权重，那么应该怎么写呢？我们可以先把编码器部分的键值对单独取出来导入到模型中便可以了。

import torch
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.encoder = ...
        self.decoder = ...
    def forward(self, x):
        # 定义前向传播
        
model = MyModel()

# 加载预训练权重文件
pretrained_weights_path = 'path_to_your_pretrained_weights.pth'
pretrained_weights = torch.load(pretrained_weights_path)

# 获得编码器的state_dict
encoder_state_dict = pretrained_weights['encoder']

# 将预训练权重加载到编码器
model.load_state_dict(encoder_state_dict, strict=False)

如果模型的参数是

二、查看模型每层的权重参数

有时我们再训练模型时会需要进行dubug去查看一下模型的参数，这时候便需要用到model.named_parameters函数，这是一个模型的方法，它返回一个生成器，可以用for循环去迭代遍历模型的每一层参数（parameter）及其名称（name)。其用法如下：

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.fc1 = nn.Linear(10, 5)

    def forward(self, x):
        x = self.conv1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        return x

# 创建模型实例
model = SimpleModel()

# 使用named_parameters()方法迭代模型的参数和它们的名称
for name, param in model.named_parameters():
    print(name, param.size())

代码中的name便是模型每一层的名字，param便是模型的参数。通过输出我们可以看出，模型的每一层参数都包括权重（weight)和偏置（bias)，且每一层参数的名字(name)与你在定义模型结构时的命名有关。
对于上述代码，如果我们只想查看卷积层的参数，可以这么使用：

for name, param in model.conv1.named_parameters():
    print(name, param)

三、模型的保存

当我们训练好一个模型时，这时候便需要对模型进行保存，模型的保存一般来说有两种常用的方法，一种是保存整个模型，包括模型架构和权重。第二种方法是只保存模型的每层参数及其名字。一般来说第二种方法用的更多，因为其保存后的文件规模更小，并且其除了可以保存模型外，还可以保存优化器的参数，训练的轮数等参数，更加灵活。
第一种方法：

torch.save(model, 'complete_model.pth')

第二种方法：

model =...
optimizer = ...
epoch = 5

checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'epoch': epoch
}
torch.save(checkpoint, 'checkpoint.pth')

另外多提一嘴，model.state_dict()和model.named_parameters()是用于访问模型参数的两个不同的方法，这两种方法有一些区别：

1. model.state_dict() 返回一个有序字典（collections.OrderedDict），其中包含模型中所有参数和它们对应的值。model.named_parameters() 返回一个生成器，生成器中每个元素是一个元组，包含参数的名称和参数本身。
2. state_dict() 包含了模型中所有参数的值，其参数值require_grad为False。named_parameters() 提供了可学习可更新的参数的名称和数值，其迭代出来的参数值require_grad为True。
3. state_dict() 可以直接用于torch.save()进行序列化保存到文件。 named_parameters() 生成的参数名称和对象不能直接序列化保存模型权重，但可以用来检查或修改参数（不推荐此方法修改参数）。
   因此总结一下就是如果我们在训练过程中想查看模型的可训练参数大小可以用model.named_parameters(),如果是想保存模型的参数可以使用model.state_dict()。