前言

在英伟达深度学习的教程中，我们学习了对美国手语数据集的模型训练，我们需要用到保存已经写的模型，然后加载它，进行图像的预测。

正文

首先，要知道2个概念：
保存模型 = 保存权重
加载模型 = 恢复权重

其次，保存并加载模型，有两种方法：
① 保存 model，保存内容有结构+权重+路径
② 保存 model.state_dict()，保存内容只有权重

第一种保存加载模型方式：保存model（不推荐）

首先，创建一个save.py文件，用来定义模型和保存模型
我们定义了一个名为 BaiQing类 的 卷积块

class BaiQing(nn.Module):
    def __init__(self, in_ch, out_ch, dropout_p):
        kernel_size = 3
        super().__init__()

        self.model = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size, stride=1, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(),
            nn.Dropout(dropout_p),
            nn.MaxPool2d(2, stride=2)
        )

    def forward(self, x):
        return self.model(x)

要注意的是，该类初始化了3个量，后续会提到
然后，我们定义一个基于 BaiQing卷积块 的 基础模型base_model

IMG_CHS = 1     # 图片通道数(Channels)
n_classes = 25  # 分类类别数量
flattened_img_size = 75 * 3 * 3
base_model = nn.Sequential(
    BaiQing(IMG_CHS, 25, 0),
    BaiQing(25, 50, 0.2),
    BaiQing(50, 75, 0.2),
    nn.Flatten(),
    nn.Linear(flattened_img_size, 512),
    nn.Dropout(0.3),
    nn.ReLU(),
    nn.Linear(512, n_classes)
)

没错，用同一个类 BaiQing 创建了 3 个不同的层实例，请注意，后续会提到

重点来了，第一种保存模型方法，直接保存model

1	torch.save(base_model, 'model.pth', weights_only=False)

我们将 基础模型base_model 直接保存，命名为 model.pth 文件
然后，我们创建一个 load.py 文件用来调用模型
开始调用 model.pth 之前，因为base_model中有用到BaiQing类，所以需要现在新的文件内再次定义，或者直接导包

class BaiQing(nn.Module):
    def __init__(self, in_ch, out_ch, dropout_p):
        kernel_size = 3
        super().__init__()

        self.model = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size, stride=1, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(),
            nn.Dropout(dropout_p),
            nn.MaxPool2d(2, stride=2)
        )

    def forward(self, x):
        return self.model(x)

model = torch.load("model.pth")

于是我们就加载了模型了

‼️ 但是，这种方法不推荐：虽然看起来非常方便，但是问题很多

如果改了文件名或者类名，就直接无法运行
更换conda环境，PyTorch版本不对，也可能加载失败
模型的结构更改需要回到save文件，不可控，像薛定谔的猫
写起来比较省事，可以用于临时代码实验

第二种保存加载模型方式：保存model.state_dict()（推荐，权威）

我们同样创建 save.py文件 和 load.py文件
在 save.py文件 里创建好 卷积块 和 基础模型

使用第二种保存模型方法，只保存权重，保存model.state_dict()

1	torch.save(model.state_dict(), "model.pth")

然后进入 load.py文件
由于model.pth 内现在只保存了模型的权重，所以我们需要把模型实例化：

# 模型实例化
model = BaiQing()
# 载入模型
model.load_state_dict(torch.load("model.pth", map_location=device, weights_only=False))
# 模型评估模式
model.eval()

这样是正常的保存和载入的步骤
但是，我们很快会发现：

class BaiQing(nn.Module):
    def __init__(self, in_ch, out_ch, dropout_p):
·
·
·
model = BaiQing()
·
·

没错，我们的卷积块，初始化了3个变量，而实例化的时候，没有传入任何参数，这会导致报错：

1	TypeError: __init__() missing 3 required positional arguments: 'in_ch', 'out_ch', and 'dropout_p'

因为我们在save.py文件中的模型，BaiQing类初始化了3个量，所以我们不能直接在load.py文件中，直接传入一个量的参数，进行加载模型
所以需要再load.py文件中添加一个步骤：再次定义一个 base_model基础模型

完整代码如下：

class BaiQing(nn.Module):
    def __init__(self, in_ch, out_ch, dropout_p):
        kernel_size = 3
        super().__init__()

        self.model = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size, stride=1, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(),
            nn.Dropout(dropout_p),
            nn.MaxPool2d(2, stride=2)
        )

    def forward(self, x):
        return self.model(x)

IMG_CHS = 1  # 图片通道数(Channels)
n_classes = 25  # 分类类别数量

base_model = nn.Sequential(
    BaiQing(IMG_CHS, 25, 0),
    BaiQing(25, 50, 0.2),
    BaiQing(50, 75, 0.2),
    nn.Flatten(),
    nn.Linear(75 * 3 * 3, 512),
    nn.Dropout(0.3),
    nn.ReLU(),
    nn.Linear(512, n_classes)
)

state_dict = torch.load('model.pth', map_location=device, weights_only=False)
base_model.load_state_dict(state_dict)