备注

点击此处下载完整示例代码

使用 Flask 通过 REST API 部署 PyTorch 的 Python 教程

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：Avinash Sajjanshetty

在本教程中，我们将使用 Flask 部署 PyTorch 模型，并公开 REST API 进行模型推理。特别是，我们将部署一个预训练的 DenseNet 121 模型，用于图像检测。

提示

所有在此处使用的代码均以 MIT 许可证发布，可在 GitHub 上获取。

这代表了一系列关于在生产环境中部署 PyTorch 模型的教程中的第一篇。使用 Flask 这种方式是启动服务 PyTorch 模型最简单的方法，但它不适用于高性能要求的用例。对于那种情况：

如果您已经熟悉 TorchScript，可以直接跳转到我们的“在 C++中加载 TorchScript 模型”教程。

如果您首先需要复习 TorchScript，请查看我们的“TorchScript 简介”教程。

API 定义

我们首先定义我们的 API 端点、请求和响应类型。我们的 API 端点位于 /predict ，它接受包含图像的 file 参数的 HTTP POST 请求。响应将是包含预测的 JSON 响应：

{"class_id": "n02124075", "class_name": "Egyptian_cat"}

依赖项

通过运行以下命令安装所需的依赖项：

pip install Flask==2.0.1 torchvision==0.10.0

简单的 Web 服务器 ¶

以下是一个简单的 Web 服务器，摘自 Flask 的文档

from flask import Flask
app = Flask(__name__)


@app.route('/')
def hello():
    return 'Hello World!'

我们还将更改响应类型，使其返回包含 ImageNet 类别 ID 和名称的 JSON 响应。更新后的 app.py 文件现在将是：

from flask import Flask, jsonify
app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    return jsonify({'class_id': 'IMAGE_NET_XXX', 'class_name': 'Cat'})

推理 ¶

在接下来的章节中，我们将专注于编写推理代码。这包括两个部分，一部分是准备图像以便输入到 DenseNet 中，接下来我们将编写从模型获取实际预测的代码。

准备图像

DenseNet 模型需要图像为 3 通道 RGB 图像，大小为 224 x 224。我们还将使用所需的均值和标准差值对图像张量进行归一化。您可以在此处了解更多信息。

我们将使用 torchvision 库中的 transforms ，构建一个转换管道，将我们的图像转换为所需格式。您可以在此处了解更多关于转换的信息。

import io

import torchvision.transforms as transforms
from PIL import Image

def transform_image(image_bytes):
    my_transforms = transforms.Compose([transforms.Resize(255),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize(
                                            [0.485, 0.456, 0.406],
                                            [0.229, 0.224, 0.225])])
    image = Image.open(io.BytesIO(image_bytes))
    return my_transforms(image).unsqueeze(0)

上述方法接收字节形式的图像数据，应用一系列转换并返回一个张量。为了测试上述方法，以字节模式读取图像文件（首先将../_static/img/sample_file.jpeg 替换为您的计算机上文件的实际路径）并查看是否返回一个张量：

with open("../_static/img/sample_file.jpeg", 'rb') as f:
    image_bytes = f.read()
    tensor = transform_image(image_bytes=image_bytes)
    print(tensor)

预测 ¶

现在将使用预训练的 DenseNet 121 模型来预测图像类别。我们将使用来自 torchvision 库的一个模型，加载模型并获取推理。虽然在这个例子中我们将使用预训练模型，但您可以使用这种方法来使用自己的模型。有关加载您自己的模型的信息，请参阅本教程。

from torchvision import models

# Make sure to set `weights` as `'IMAGENET1K_V1'` to use the pretrained weights:
model = models.densenet121(weights='IMAGENET1K_V1')
# Since we are using our model only for inference, switch to `eval` mode:
model.eval()


def get_prediction(image_bytes):
    tensor = transform_image(image_bytes=image_bytes)
    outputs = model.forward(tensor)
    _, y_hat = outputs.max(1)
    return y_hat

张量 y_hat 将包含预测类别的索引。然而，我们需要一个人类可读的类别名称。为此，我们需要类别 ID 到名称的映射。下载此文件为 imagenet_class_index.json ，并记住您保存的位置（或者，如果您正在按照本教程的精确步骤进行，请将其保存在 tutorials/_static 中）。此文件包含 ImageNet 类别 ID 到 ImageNet 类别名称的映射。我们将加载此 JSON 文件并获取预测索引的类别名称。

import json

imagenet_class_index = json.load(open('../_static/imagenet_class_index.json'))

def get_prediction(image_bytes):
    tensor = transform_image(image_bytes=image_bytes)
    outputs = model.forward(tensor)
    _, y_hat = outputs.max(1)
    predicted_idx = str(y_hat.item())
    return imagenet_class_index[predicted_idx]

在使用 imagenet_class_index 字典之前，我们首先将张量值转换为字符串值，因为 imagenet_class_index 字典的键是字符串。我们将测试我们上面的方法：

with open("../_static/img/sample_file.jpeg", 'rb') as f:
    image_bytes = f.read()
    print(get_prediction(image_bytes=image_bytes))

你应该得到如下响应：

['n02124075', 'Egyptian_cat']

数组的第一个元素是 ImageNet 类 ID，第二个元素是人可读的名称。

将模型集成到我们的 API 服务器中

在本节的最后部分，我们将把我们的模型添加到我们的 Flask API 服务器中。由于我们的 API 服务器需要接收一个图片文件，我们将更新我们的 predict 方法以从请求中读取文件：

from flask import request

@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        # we will get the file from the request
        file = request.files['file']
        # convert that to bytes
        img_bytes = file.read()
        class_id, class_name = get_prediction(image_bytes=img_bytes)
        return jsonify({'class_id': class_id, 'class_name': class_name})

import io
import json

from torchvision import models
import torchvision.transforms as transforms
from PIL import Image
from flask import Flask, jsonify, request


app = Flask(__name__)
imagenet_class_index = json.load(open('<PATH/TO/.json/FILE>/imagenet_class_index.json'))
model = models.densenet121(weights='IMAGENET1K_V1')
model.eval()


def transform_image(image_bytes):
    my_transforms = transforms.Compose([transforms.Resize(255),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize(
                                            [0.485, 0.456, 0.406],
                                            [0.229, 0.224, 0.225])])
    image = Image.open(io.BytesIO(image_bytes))
    return my_transforms(image).unsqueeze(0)


def get_prediction(image_bytes):
    tensor = transform_image(image_bytes=image_bytes)
    outputs = model.forward(tensor)
    _, y_hat = outputs.max(1)
    predicted_idx = str(y_hat.item())
    return imagenet_class_index[predicted_idx]


@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        file = request.files['file']
        img_bytes = file.read()
        class_id, class_name = get_prediction(image_bytes=img_bytes)
        return jsonify({'class_id': class_id, 'class_name': class_name})


if __name__ == '__main__':
    app.run()

FLASK_ENV=development FLASK_APP=app.py flask run

用于向我们的应用发送 POST 请求的库：

import requests

resp = requests.post("http://localhost:5000/predict",
                     files={"file": open('<PATH/TO/.jpg/FILE>/cat.jpg','rb')})

打印 resp.json()将显示以下内容：

{"class_id": "n02124075", "class_name": "Egyptian_cat"}

我们编写的服务器相当简单，可能无法满足您生产应用的所有需求。因此，这里有一些您可以做的改进：

端点 /predict 假设请求中始终存在图像文件。这可能并不适用于所有请求。用户可能会发送带有不同参数的图像，或者根本不发送图像。

用户也可能发送非图像类型的文件。由于我们没有处理错误，这会导致我们的服务器崩溃。添加一个显式的错误处理路径，该路径将抛出异常，将使我们能够更好地处理不良输入。

尽管模型可以识别大量图像类别，但它可能无法识别所有图像。增强实现以处理模型无法在图像中识别任何内容的情况。

我们以开发模式运行 Flask 服务器，这不适合在生产环境中部署。您可以查看这篇教程了解如何在生产环境中部署 Flask 服务器。

您也可以通过创建一个包含表单的页面来添加 UI，该表单接收图像并显示预测结果。查看类似项目的演示及其源代码。

在本教程中，我们只展示了如何构建一个一次只能返回单个图像预测结果的服务。我们可以修改我们的服务，使其能够一次返回多个图像的预测结果。此外，service-streamer 库会自动将请求队列到您的服务中，并将它们采样成小批量，以便输入到您的模型中。您可以查看这个教程。

最后，我们鼓励您查看页面顶部链接的我们其他关于部署 PyTorch 模型的教程。

脚本总运行时间：（0 分钟 0.000 秒）

由 Sphinx-Gallery 生成的画廊