0. 前言

我们已系统性地探讨了如何使用 PyTorch 训练和测试各类机器学习模型。我们从 PyTorch 的基础组件入手，掌握了高效完成深度学习任务的必备工具；随后深入研究了基于 PyTorch 实现的多种深度学习模型架构及其应用场景。在本节中，我们将重点讨论如何将这些模型投入生产环境。简单来说，就是讨论将一个已经训练并测试过的模型对象部署到一个独立的环境中，使其能够对新输入数据进行预测或推理。这称为模型的生产化，即将模型部署到生产系统中。
本节我们将首先构建一个简易的 PyTorch 推理管道：通过输入数据和预训练模型的存储路径即可完成预测。随后将该推理管道部署至模型服务器，使其能够接收数据请求并返回预测结果。

1. 构建 PyTorch 模型推理管道

我们将使用在 PyTorch 深度学习概述一节中构建的 MNIST 手写数字分类模型。利用这个训练好的模型，构建一个推理管道，该管道能够根据给定的手写数字输入图像预测一个 0 到 9 之间的数字。

1.1 保存和加载训练模型

在本节中，我们将演示如何高效地加载一个已保存的预训练 PyTorch 模型，随后该模型将用于服务请求。
假设我们已完成模型训练并在测试数据样本上进行评估，在实际应用中，我们希望训练脚本关闭后仍能调用这个精心训练的模型进行手写数字识别——这正是模型服务的核心价值所在。
接下来，我们将实现无需重新训练即可在独立程序中使用该预训练模型。关键步骤是将模型对象序列化保存为可恢复的文件。PyTorch 提供两种主要保存方式：

• 第一种(非推荐方案)是完整保存模型对象：：

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

  随后，我们可以通过以下方式读取保存的模型：：

  model = torch.load(PATH_TO_MODEL, weights_only=False)
  

  虽然这种方法看似最直接，但在某些情况下可能会出现问题。这是因为我们不仅保存了模型参数，还保存了源代码中使用的模型类和目录结构。如果后续类签名或目录结构发生变化，加载模型时可能会出现无法修复的错误。

• 第二种(更推荐的方法)是仅保存模型的参数：

  torch.save(model.state_dict(), 'model_params.pth')
  

  当需要恢复模型时，我们首先实例化一个空的模型对象，然后将模型参数加载到该对象中：

  model = ConvNet()
  model.load_state_dict(torch.load(PATH_TO_MODEL, weights_only=True))
  

(1) 使用更推荐的方式保存模型：

PATH_TO_MODEL = "./convnet.pth"
torch.save(model.state_dict(), 'convnet.pth')

convnet.pth 文件实际上是一个 pickle 文件，包含了模型的参数。

至此，我们可以安全地关闭当前程序，并编写另一程序文件。

(2) 首先，导入所需库：

import torch
import torch.nn as nn

(3) 接下来，实例化一个空的卷积神经网络模型。理想情况下，步骤 1 中定义的模型应该编写在 Python 脚本中(例如 cnn_model.py)：

from cnn_model import ConvNet
model = ConvNet()

否则，需要重新编写模型定义并实例化：

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.cn1 = nn.Conv2d(1, 16, 3, 1)
        self.cn2 = nn.Conv2d(16, 32, 3, 1)
        self.dp1 = nn.Dropout2d(0.10)
        self.dp2 = nn.Dropout2d(0.25)
        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32
        self.fc2 = nn.Linear(64, 10)
 
    def forward(self, x):
        x = self.cn1(x)
        x = F.relu(x)
        x = self.cn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dp1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dp2(x)
        x = self.fc2(x)
        op = F.log_softmax(x, dim=1)
        return op
    
model = ConvNet()

(4) 接下来，我们可以将保存的模型参数加载到这个实例化的模型对象中：

PATH_TO_MODEL = "./convnet.pth"
model.load_state_dict(torch.load(PATH_TO_MODEL, weights_only=True, map_location="cpu"))

输出结果如下所示：

<All keys matched successfully>

这实质上表明参数加载成功——我们实例化的模型结构与当初保存参数时的原始模型完全一致。通过指定 map_location=torch.device('cpu')，我们明确将模型加载到 CPU 设备而非 GPU 上。

(5) 最后通过以下代码声明不更新已加载模型的参数值：

model.eval()

输出结果如下所示：

该操作再次验证了我们正在使用的模型架构与训练时完全一致。

1.2 构建推理管道

在上一节成功将预训练模型加载到新环境，我们现在将构建模型推理管道并执行预测。

(1) 此时我们已经完整恢复了预训练模型对象。接下来使用以下代码加载待预测图像：

image = Image.open("./digit_image.jpg")

在任何推理管道中，核心有三个主要组件：数据预处理组件，模型推理组件(对于神经网络来说是前向传播)，以及后处理组件。

(2) 首先从第一部分开始，定义一个函数，该函数接受一张图像并将其转换为要输入到模型中的张量：

def image_to_tensor(image):
    gray_image = transforms.functional.to_grayscale(image)
    resized_image = transforms.functional.resize(gray_image, (28, 28))
    input_image_tensor = transforms.functional.to_tensor(resized_image)
    input_image_tensor_norm = transforms.functional.normalize(input_image_tensor, (0.1302,), (0.3069,))
    return input_image_tensor_norm

该预处理流程包含四个标准化步骤：

• 色彩空间转换：将 RGB 图像转换为灰度图像
• 尺寸规范化：将图像调整为 28x28 像素的大小，即模型训练时使用的图像大小
• 数据结构转换：将图像数组转换为 PyTorch 张量。
  数值归一化：张量中的像素值按照模型训练期间使用的均值和标准差值进行归一化

(3) 定义完预处理函数后，调用它来处理加载的图像：

image_tensor = image_to_tensor(image)

(4) 接下来，我们定义模型推理功能。该功能接收张量形式的输入图像，并输出 0 到 9 之间的数字预测结果：

def run_model(input_tensor):
    model_input = input_tensor.unsqueeze(0)
    with torch.no_grad():
        model_output = model(model_input)[0]
    model_prediction = model_output.detach().numpy().argmax()
    return model_prediction

model_output 包含模型的原始预测结果，其中每个输入图像对应一组预测值。由于我们只输入单张图像，预测列表仅包含索引 0 处的一个条目。该原始预测是由 10 个概率值组成的张量，按顺序对应数字 0-9 的预测概率，将该张量将被转换为一个 numpy 数组，选择概率最高的数字。

(5) 使用此函数生成模型预测：

output = run_model(input_tensor)
print(output)
print(type(output))

输出结果如下所示：

2
<class 'numpy.int64'>

模型输出的是一个 numpy 整数。对照图像，模型预测结果准确无误。

(6) 除了输出最终预测值，我们还可以编写调试函数来深入分析原始预测概率等指标：

def debug_model(input_tensor):
    model_input = input_tensor.unsqueeze(0)
    with torch.no_grad():
        model_output = model(model_input)[0]
    model_prediction = model_output.detach().numpy()
    return np.exp(model_prediction)

该函数与 run_model 函数的核心逻辑相同，但其特殊之处在于会返回每个数字类别的原始概率列表。由于模型最后一层采用 Log_softmax 层，原始输出实际上是 softmax 概率的对数值。因此我们需要通过指数运算还原 softmax 概率分布。

(7) 通过这个调试函数，我们可以更详细地查看模型的表现，比如概率分布是否平滑，或是否有明显的峰值：

print(debug_model(input_tensor))

输出结果如下所示：

[2.6426747e-05 5.4441585e-07 9.9744761e-01 7.3382223e-05 4.2181364e-05
 1.5351664e-07 5.4828524e-06 7.8664209e-07 2.4013412e-03 2.2161457e-06]

可以看到，列表中的第三个概率是最高的，且明显大于其他的概率，这对应的数字是 2。

(8) 最后，我们需要对模型预测结果进行后处理，使其能够被其他应用程序直接使用。在本节中，我们只需将模型预测的数字从整型转换为字符串类型：

def post_process(output):
    return str(output)

后处理步骤在其他场景可能更为复杂，例如语音识别任务中可能需要对输出波形进行平滑处理、去除异常值等操作。字符串作为可序列化格式，能够方便地在不同服务器和应用程序间传递预测结果。我们可以验证后处理结果是否符合预期：

final_output = post_process(output)
print(final_output)
print(type(final_output))

输出结果如下所示：

2
<class 'str'>

可以看到，输出已成功转换为字符串类型。至此，我们已完成以下完整流程的实践：加载已保存的模型架构、恢复其训练权重，并使用加载后的模型为样本输入数据(一张图片)生成预测。我们加载了样本图像，通过预处理将其转换为 PyTorch 张量，将其作为输入传递给模型以获得预测结果，最后对预测结果进行后处理生成最终输出。
本节中，输入是外部提供的图像文件，输出是生成的 0 到 9 之间的数字字符串，系统可以嵌入到任何需要手写数字识别功能的应用程序中。在下一节中，我们将进一步深入研究模型服务，构建一个无需复制粘贴代码、任何应用程序都能直接调用的手写数字识别系统。

2. 构建模型服务器

我们已经构建了一个完整的模型推理管道，包含了从预训练模型独立执行预测所需的所有代码。接下来，我们将构建模型服务器，它本质上是一个托管模型推理管道的设备，能够通过接口主动监听任何输入数据，并通过接口输出模型预测结果。

2.1 使用 Flask 编写基础应用

接下来，我们将采用流行的 Python 库 Flask 来开发模型服务器。借助 Flask，仅需数行代码即可构建模型服务器：

from flask import Flask
app = Flask(__name__)
 
@app.route('/')
def hello_world():
    return 'Hello, World!'
 
if __name__ == '__main__': 
    app.run(host='localhost', port=8890)

将此 Python 脚本保存为 example.py 并在命令行中运行：

$ python example.py

将显示以下输出结果：

这表明启动一个 Flask 服务器，该服务器将提供一个名为 example 的应用。打开浏览器并访问以下 URL：

http://localhost:8890

在浏览器中输出结果如下所示：

本质上，该 Flask 服务器正在监听 IP 地址 0.0.0.0 (本地主机)上 8890 端口的 “/” 端点。当我们在浏览器地址栏输入 localhost:8890/ 并按下回车时，服务器就会接收到请求。随后服务器执行 hello_world 函数，根据 example.py 中的函数定义返回字符串 “Hello, World!”。

2.2 使用 Flask 构建模型服务器

了解了 Flask 服务器运行原理后。接下来，我们将利用模型推理管道创建模型服务器。在本节中，我们将启动一个能监听传入请求(图像数据输入)的服务器。此外，我们还将编写另一个 Python 脚本，通过发送样本图像向该服务器发起请求。Flask 服务器将对该图像执行模型推理，并输出后处理后的预测结果。

2.2.1 配置 Flask 服务所需的模型推理服务

在本节中，我们将加载一个预训练的模型并编写模型推理管道。

(1) 首先构建 Flask 服务器。从导入所需库开始：

import os
import json
import numpy as np
from flask import Flask, request

import torch
import torch.nn as nn
import torch.nn.functional as F

flask 和 torch 是本任务的核心依赖库，此外还需要 numpy 和 json 等基础库。

(2) 接下来，需要定义模型类(架构)：

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.cn1 = nn.Conv2d(1, 16, 3, 1)
        self.cn2 = nn.Conv2d(16, 32, 3, 1)
        self.dp1 = nn.Dropout(0.10)
        self.dp2 = nn.Dropout(0.25)
        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32
        self.fc2 = nn.Linear(64, 10)
 
    def forward(self, x):
        x = self.cn1(x)
        x = F.relu(x)
        x = self.cn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dp1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dp2(x)
        x = self.fc2(x)
        op = F.log_softmax(x, dim=1)
        return op

(3) 定义了模型类之后，实例化模型对象并加载预训练参数：

model = ConvNet()
PATH_TO_MODEL = "./convnet.pth"
model.load_state_dict(torch.load(PATH_TO_MODEL, weights_only=True, map_location="cpu"))
model.eval()

(4) 定义 run_model 函数，函数接收张量化的输入图像，输出模型预测结果——0 到 9 之间的任意数字：

def run_model(input_tensor):
    model_input = input_tensor.unsqueeze(0)
    with torch.no_grad():
        model_output = model(model_input)[0]
    model_prediction = model_output.detach().numpy().argmax()
    return model_prediction

(5) 定义 post_process 函数，该函数将 run_model 输出的整型数字转换为字符串：

def post_process(output):
    return str(output)

构建 Flask 应用提供模型服务

完成推理流水线搭建后，我们将构建 Flask 应用来部署加载的模型。

(1) 实例化 Flask 应用，这将创建一个与 Python 脚本同名的 Flask 应用(本节中为 server.py)：

app = Flask(__name__)

(2) 定义 Flask 服务器的端点功能，开放 /test 端点，并定义当服务器收到该端点的 POST 请求时的处理逻辑：

@app.route("/test", methods=["POST"])
def test():
    data = request.files['data'].read()
    md = json.load(request.files['metadata'])
    input_array = np.frombuffer(data, dtype=np.float32)
    input_image_tensor = torch.from_numpy(input_array).view(md["dims"])
    output = run_model(input_image_tensor)
    final_output = post_process(output)
    return final_output

首先，为函数 test 添加一个装饰器(该函数定义在装饰器下方)。这个装饰器指示 Flask 应用：当收到向 /test 端点发送的 POST 请求时执行此函数。
然后定义 test 函数的具体逻辑。首先，从 POST 请求中读取数据和元数据。因为数据是序列化的，需先转为 numpy 数组，再转换为 PyTorch 张量。接下来，使用元数据中提供的图像尺寸来调整张量的形状。最后用加载好的模型对该张量执行前向传播，获得预测结果后经后处理由 test 函数返回。

(3) 至此已完成启动 Flask 应用的所有准备工作，在 server.py 末尾添加以下两行代码：

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8890)

表示 Flask 服务器将运行在 IP 地址 0.0.0.0 (即本地主机)的 8890 端口。保存脚本后，执行以下命令：

$ python server.py

输出结果如下所示：

3. 使用 Flask 服务器执行预测

我们已成功启动模型服务器并开始监听请求。接下来，编写请求服务代码。

(1) 首先，导入相关库：

import io
import json
import requests
from PIL import Image

from torchvision import transforms

requests 库用于向 Flask 服务器发起实际的 POST 请求。Image 模块用于读取样本输入图像文件，transforms 则用于对输入图像数组进行预处理。

(2) 接下来，读取图像文件：

image = Image.open("./digit_image.jpg")

此处读取的是 RGB 图像，其尺寸不必符合模型预期的 28x28 输入尺寸。

(3) 定义预处理函数，将读取的图像转换为模型可识别的格式：

def image_to_tensor(image):
    gray_image = transforms.functional.to_grayscale(image)
    resized_image = transforms.functional.resize(gray_image, (28, 28))
    input_image_tensor = transforms.functional.to_tensor(resized_image)
    input_image_tensor_norm = transforms.functional.normalize(input_image_tensor, (0.1302,), (0.3069,))
    return input_image_tensor_norm

(4) 定义完函数后即可执行预处理过程：

image_tensor = image_to_tensor(image)

image_tensor 是我们需要发送给 Flask 服务器的输入数据。

(5) 将数据打包并发送。需要同时发送图像的像素值及其形状 (28x28)，以便接收端的 Flask 服务器知道如何将像素值流重建为图像：

dimensions = io.StringIO(json.dumps({'dims': list(image_tensor.shape)}))
data = io.BytesIO(bytearray(image_tensor.numpy()))

将张量的形状转换为字符串，并将图像数组转换为字节，确保所有数据都可序列化。

(6) 客户端发起 POST 请求：

r = requests.post('http://localhost:8890/test',
                  files={'metadata': dimensions, 'data' : data})

通过 requests 库向 URL localhost:8890/test (Flask 服务器监听该地址)发起 POST 请求。我们以字典形式同时发送图像数据(字节流格式)和元数据(字符串格式)。

(7) 上述代码中的变量 r 将接收 Flask 服务器的响应，该响应应包含经后处理的模型预测结果。接下来，读取该输出：

response = json.loads(r.content)

response 变量包含了 Flask 服务器的输出结果，也就是一个介于 0 和 9 之间的数字(以字符串形式表示)。

(8) 打印响应，确保结果正确：

print("Predicted digit :", response)

(9) 将该 Python 脚本保存为 make_request.py，并执行以下命令：

$ python make_request.py

输出结果如下所示：

Predicted digit : 2

根据输入图像，模型返回的响应结果正确。
至此，我们已从编写简单的推理函数进阶到创建可远程部署、通过网络提供预测服务的模型服务器。

小结

在本节中，我们探讨了如何将训练好的 PyTorch 深度学习模型部署到生产环境中，成功构建了一个独立运行的模型服务器，能够对手写数字图像进行预测。这套方法可以轻松扩展至其他机器学习模型，这为使用 PyTorch 和 Flask 开发机器学习应用开辟了无限可能。

系列链接

PyTorch实战（1）——深度学习（Deep Learning）
PyTorch实战（2）——使用PyTorch构建神经网络
PyTorch实战（3）——PyTorch vs. TensorFlow详解
PyTorch实战（4）——卷积神经网络（Convolutional Neural Network，CNN）
PyTorch实战（5）——深度卷积神经网络
PyTorch实战（6）——模型微调详解
PyTorch实战（7）——循环神经网络
PyTorch实战（8）——图像描述生成
PyTorch实战（9）——从零开始实现Transformer
PyTorch实战（10）——从零开始实现GPT模型
PyTorch实战（11）——随机连接神经网络（RandWireNN）
PyTorch实战（12）——图神经网络（Graph Neural Network，GNN）
PyTorch实战（13）——图卷积网络（Graph Convolutional Network，GCN）
PyTorch实战（14）——图注意力网络（Graph Attention Network，GAT）
PyTorch实战（15）——基于Transformer的文本生成技术
PyTorch实战（16）——基于LSTM实现音乐生成
PyTorch实战（17）——神经风格迁移
PyTorch实战（18）——自编码器（Autoencoder，AE）
PyTorch实战（19）——变分自编码器（Variational Autoencoder，VAE）
PyTorch实战（20）——生成对抗网络（Generative Adversarial Network，GAN）
PyTorch实战（21）——扩散模型（Diffusion Model）
PyTorch实战（22）——MuseGAN详解与实现
PyTorch实战（23）——基于Transformer生成音乐
PyTorch实战（24）——深度强化学习
PyTorch实战（25）——使用PyTorch构建DQN模型
PyTorch实战（26）——PyTorch分布式训练
PyTorch实战（27）——自动混合精度训练