机器学习详解(14):模型的保存和部署实例

在前面的文章卷积神经网络CNN之手语识别代码详解和CNN图像数据增强(解决过拟合问题)中，我们介绍了用CNN模型来识别手语，同时通过数据增强缓解了过拟合的现象。那么本节就来学习如何部署已经训练好的模型，并用在手语预测中(以PyTorch为例)。

文章目录

[1 保存模型和代码](#1 保存模型和代码)
[2 加载模型](#2 加载模型)
[3 为模型准备输入图像](#3 为模型准备输入图像)
- [3.1 显示图像](#3.1 显示图像)
- [3.2 图像预处理](#3.2 图像预处理)
- [3.3 预测](#3.3 预测)
- [3.4 理解预测结果](#3.4 理解预测结果)
- [3.5 整合预测函数](#3.5 整合预测函数)
[4 常见ML模型的保存和加载](#4 常见ML模型的保存和加载)
[5 总结](#5 总结)

1 保存模型和代码

在上篇文章CNN图像数据增强(解决过拟合问题)的最后，我们通过下面的代码保存了模型：

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

注意一下，PyTorch中不能保存已经编译好的模型，但未编译的模型和已编译的模型是共享权重的

接着我们将上篇文章的代码也整合到util.py文件中：

python 复制代码

import torch
import torch.nn as nn

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

def get_batch_accuracy(output, y, N):
    pred = output.argmax(dim=1, keepdim=True)
    correct = pred.eq(y.view_as(pred)).sum().item()
    return correct / N


def train(model, train_loader, train_N, random_trans, optimizer, loss_function):
    loss = 0
    accuracy = 0

    model.train()
    for x, y in train_loader:
        output = model(random_trans(x))
        optimizer.zero_grad()
        batch_loss = loss_function(output, y)
        batch_loss.backward()
        optimizer.step()

        loss += batch_loss.item()
        accuracy += get_batch_accuracy(output, y, train_N)
    print('Train - Loss: {:.4f} Accuracy: {:.4f}'.format(loss, accuracy))

def validate(model, valid_loader, valid_N, loss_function):
    loss = 0
    accuracy = 0

    model.eval()
    with torch.no_grad():
        for x, y in valid_loader:
            output = model(x)

            loss += loss_function(output, y).item()
            accuracy += get_batch_accuracy(output, y, valid_N)
    print('Valid - Loss: {:.4f} Accuracy: {:.4f}'.format(loss, accuracy))

2 加载模型

在加载模型之前，我们先导入一下需要用的库，并做一些初始化：

python 复制代码

import pandas as pd
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader
import torchvision.io as tv_io
import torchvision.transforms.v2 as transforms
import torchvision.transforms.functional as F
import matplotlib.pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.cuda.is_available()

import torch._dynamo
torch._dynamo.config.suppress_errors = True

由于模型使用了自定义模块，我们需要加载该类的代码。以下是加载代码的方式：

python 复制代码

from utils import MyConvBlock

现在我们已经定义了 MyConvBlock，可以使用 torch.load 从路径加载模型，通过 map_location 指定设备（CPU 或 GPU）。加载后，打印模型以检查是否与上一个 notebook 中的模型一致：

python 复制代码

model = torch.load('model.pth', map_location=device)
model

输出:
  model = torch.load('model.pth', map_location=device)
Sequential(
  (0): MyConvBlock(
    (model): Sequential(
      (0): Conv2d(1, 25, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (1): BatchNorm2d(25, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
      (3): Dropout(p=0, inplace=False)
      (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
  )
  (1): MyConvBlock(
    (model): Sequential(
      (0): Conv2d(25, 50, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (1): BatchNorm2d(50, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
      (3): Dropout(p=0.2, inplace=False)
      (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
  )
  (2): MyConvBlock(
    (model): Sequential(
      (0): Conv2d(50, 75, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (1): BatchNorm2d(75, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
      (3): Dropout(p=0, inplace=False)
      (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
  )
  (3): Flatten(start_dim=1, end_dim=-1)
  (4): Linear(in_features=675, out_features=512, bias=True)
  (5): Dropout(p=0.3, inplace=False)
  (6): ReLU()
  (7): Linear(in_features=512, out_features=24, bias=True)
)

我们可以验证模型是否加载到了 GPU 上：

python 复制代码

next(model.parameters()).device

3 为模型准备输入图像

3.1 显示图像

现在我们将使用模型对从未见过的新图像进行预测，这一过程被称为推理（inference）。在对新图像进行预测时，我们先使用 matplotlib 库来显示图像：

python 复制代码

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

def show_image(image_path):
    image = mpimg.imread(image_path)
    plt.imshow(image, cmap='gray')
    
show_image('b.png')

图像如下：

我们观察一下图像的特点：

图像分辨率远高于训练数据中的图像。
图像是彩色的，而训练数据集中的图像是灰度图像，大小为 28 × 28 28 \times 28 28×28 像素。

注意

在对模型进行预测时，输入的形状必须与模型训练数据的形状一致。
模型的训练数据集形状为 ( 27455 , 28 , 28 , 1 ) (27455, 28, 28, 1) (27455,28,28,1)，即 27455 张 28 × 28 28 \times 28 28×28 像素的灰度图像。

3.2 图像预处理

数据集中使用的图像大小为 28 × 28 28 \times 28 28×28 像素，且为灰度图像。在进行预测时，输入图像必须具有相同的大小和灰度格式。Python 提供多种方法可以编辑图像，而 TorchVision 提供了 read_image 函数，可通过 ImageReadMode 指定读取的图像类型。

python 复制代码

image = tv_io.read_image('b.png', tv_io.ImageReadMode.GRAY)

检查图像的形状

python 复制代码

image.shape

输出:
torch.Size([1, 184, 186])

可以看到该图像比训练时使用的图像大得多，因此我们可以再次使用 TorchVision 的 Transforms 来将数据转换为模型所需的格式，我们转换图像的步骤如下：

将图像转换为浮点数（ToDtype）。
- 设置 scale=True，将像素值从 [ 0 , 255 ] [0, 255] [0,255] 缩放到 [ 0 , 1 ] [0, 1] [0,1]。
将图像大小调整为 28 × 28 28 \times 28 28×28 像素（Resize）。
将图像转换为灰度图像（Grayscale）。
- 实际上刚刚我们已经将图片转为灰度图像，但展示了另一种获取灰度图像的方式。

python 复制代码

IMG_WIDTH = 28
IMG_HEIGHT = 28

preprocess_trans = transforms.Compose([
    transforms.ToDtype(torch.float32, scale=True), # Converts [0, 255] to [0, 1]
    transforms.Resize((IMG_WIDTH, IMG_HEIGHT)),
    transforms.Grayscale()  # From Color to Gray
])

现在我们使用 preprocess_trans 测试图像，确保其正确工作：

python 复制代码

processed_image = preprocess_trans(image)
processed_image

输出:
tensor([[[0.6960, 0.6279, 0.6348, 0.6493, 0.6584, 0.6599, 0.6638, 0.6630,
          0.6652, 0.6677, 0.6711, 0.6696, 0.6661, 0.6677, 0.7103, 0.6559,
          0.6369, 0.6507, 0.6464, 0.6379, 0.6293, 0.6216, 0.6139, 0.6074,
          0.5984, 0.5895, 0.5836, 0.6977],
         [0.7013, 0.6388, 0.6480, 0.6578, 0.6647, 0.6687, 0.6709, 0.6734,
          0.6746, 0.6794, 0.6814, 0.6784, 0.6902, 0.6801, 0.7648, 0.6606,
          0.5610, 0.6348, 0.6418, 0.6509, 0.6417, 0.6308, 0.6252, 0.6197,
          0.6095, 0.6007, 0.5923, 0.6773],
          ...略
          [0.7822, 0.6822, 0.6749, 0.6725, 0.6757, 0.6784, 0.6790, 0.6801,
          0.6824, 0.6579, 0.6527, 0.6838, 0.6824, 0.6756, 0.6680, 0.6401,
          0.6141, 0.5686, 0.5584, 0.5850, 0.6220, 0.6086, 0.6050, 0.6293,
          0.6379, 0.6134, 0.5933, 0.7271]]])

检查数值和形状

python 复制代码

processed_image.shape

输出:
torch.Size([1, 28, 28])

最后我们绘制图像，检查其是否符合训练时的预期格式：

python 复制代码

plot_image = F.to_pil_image(processed_image)
plt.imshow(plot_image, cmap='gray')

输出：

3.3 预测

现在我们已经准备好使用模型进行预测。需要注意的是，模型仍然期望输入为一批（batch）图像。如果 squeeze 函数用于移除大小为 1 1 1 的维度，那么 unsqueeze 函数可以在指定的索引位置添加大小为 1 1 1 的维度。对于模型，第一维通常是批次维度，因此可以使用 .unsqueeze(0) 添加批次维度。

python 复制代码

batched_image = processed_image.unsqueeze(0)
batched_image.shape

输出:
torch.Size([1, 1, 28, 28])

接下来，需要确保输入张量与模型位于相同的设备（CPU 或 GPU）上。

python 复制代码

batched_image_gpu = batched_image.to(device)
batched_image_gpu.device

输出:
device(type='cpu')

现在可以将数据输入到模型中进行预测了：

python 复制代码

output = model(batched_image_gpu)
output

输出:
tensor([[-29.3448,  14.1301,  -8.4348, -18.9694,  -1.4105, -11.8612, -10.7342,
         -32.7700,  -7.4740, -18.4512, -11.2153,  -6.2315, -16.0828, -21.6887,
          -8.1797, -18.2011, -12.8640, -29.7611, -18.7097,  -5.7279, -24.3577,
         -11.8029,  -6.6998, -28.0620]], grad_fn=<AddmmBackward0>)

3.4 理解预测结果

模型的预测结果是一个长度为 24 24 24 的数组。数组中的每个值表示输入图像属于对应类别的可能性，值越大表示图像属于该类别的概率越高。为了更直观地理解结果，可以通过 numpy 库的 argmax 函数找到概率最大的索引：

python 复制代码

prediction = output.argmax(dim=1).item()
prediction

输出:
1

解读预测结果

数组的每个元素对应手语字母表中的一个字母。
注意，手语字母表中不包含字母 j 和 z，因为它们需要手部移动，而我们只处理静态图像。

以下是索引与字母的映射关系：

python 复制代码

# 手语字母表（不包含 j 和 z）
alphabet = "abcdefghiklmnopqrstuvwxy"

现在可以通过预测的索引获取对应的字母：

python 复制代码

alphabet[prediction]

输出:
'b'

3.5 整合预测函数

def predict_letter(file_path):
    show_image(file_path)
    image = tv_io.read_image(file_path, tv_io.ImageReadMode.GRAY)
    image = preprocess_trans(image)
    image = image.unsqueeze(0)
    image = image.to(device)
    output = model(image)
    prediction = output.argmax(dim=1).item()
    # convert prediction to letter
    predicted_letter = alphabet[prediction]
    return predicted_letter

测试：

predict_letter("b.png")

输出:
'b'

4 常见ML模型的保存和加载

以下是 PyTorch、TensorFlow 和 Keras 中模型保存与加载的主要函数总结：

框架	保存模型	加载模型	说明
PyTorch	`torch.save(model.state_dict(), path)`	`model.load_state_dict(torch.load(path))`	保存和加载模型参数（推荐方式）；需先定义模型结构。
PyTorch	`torch.save(model, path)`	`model = torch.load(path)`	保存和加载整个模型（包括结构和参数）；适合实验性用途，但对环境依赖较强。
TensorFlow	`model.save(path)`	`tf.keras.models.load_model(path)`	保存和加载完整的模型（包括结构、权重和优化器配置）；支持 `.h5` 或 SavedModel 格式。
Keras	`model.save('model.h5')`	`tf.keras.models.load_model('model.h5')`	使用 HDF5 格式保存和加载 Keras 模型，便于兼容性。
TensorFlow/Keras	`model.save_weights(path)`	`model.load_weights(path)`	保存和加载模型的权重；需先定义模型结构。

5 总结

通过保存模型，可以将训练成果持久化，避免重复训练节省时间和计算资源。加载模型则使得我们能够方便地在不同环境中复用模型，如本地测试、云端部署或边缘设备运行。此外，保存和加载还便于团队协作和版本管理，确保模型的可追溯性和持续优化能力。这些优势使模型的应用更高效、更灵活。