使用 Flask 进行部署¶

在这个教程中,您将学习:

如何将训练好的 PyTorch 模型封装到 Flask 容器中,通过 Web API 暴露出去
如何将传入的 Web 请求转换为 PyTorch 张量,以供您的模型使用
如何为 HTTP 响应打包您模型的输出

环境设置¶

您需要一个安装了以下软件包(及其依赖项)的 Python 3 环境:

PyTorch 1.5
TorchVision 0.6.0
Flask 1.1

另外,如果需要获取一些支持文件,您还需要 git。

安装 PyTorch 和 TorchVision 的说明在 pytorch.org_ 上有介绍。安装 Flask 请查看 Flask 官网_ 。

什么是 Flask?¶

Flask 是一个用 Python 编写的轻量级 Web 服务器。它为您提供了一种便捷的方式,快速建立一个 Web API, 用于您训练好的 PyTorch 模型的预测,可直接使用,或作为更大系统中的 Web 服务。

设置和支持文件¶

我们将创建一个 Web 服务,接收图像,并将其映射到 ImageNet 数据集的 1000 个类别之一。为此,您需要一个用于测试的图像文件。另外,您还可以获取一个文件,将模型输出的类索引映射为可读的类名。

选项 1: 快速获取文件¶

可以通过检出 TorchServe 仓库并将文件复制到您的工作文件夹来快速获取这两个支持文件。 (注意:本教程不依赖于 TorchServe - 这只是快速获取文件的一种方式。) 从您的 shell 提示符下发出以下命令:

git clone https://github.com/pytorch/serve
cp serve/examples/image_classifier/kitten.jpg .
cp serve/examples/image_classifier/index_to_name.json .

And you’ve got them!

选项 2: 使用您自己的图像¶

index_to_name.json 文件在下面的 Flask 服务中是可选的。您可以使用自己的图像测试您的服务 - 需确保是一个 3 色 JPEG 图像。

构建您的 Flask 服务¶

Flask 服务的完整 Python 脚本在本教程的最后展示;您可以复制并粘贴到您自己的 app.py 文件中。下面我们将查看各个部分,以明确它们的功能。

导入¶

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

按顺序:

将使用来自 torchvision.models 的预训练 DenseNet 模型
torchvision.transforms 包含用于操作图像数据的工具
Pillow (PIL) 是我们最初加载图像文件时将使用的库
当然我们还需要从 flask 导入一些类

预处理¶

def transform_image(infile):
    input_transforms = [transforms.Resize(255),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
            [0.229, 0.224, 0.225])]
    my_transforms = transforms.Compose(input_transforms)
    image = Image.open(infile)
    timg = my_transforms(image)
    timg.unsqueeze_(0)
    return timg

Web 请求给了我们一个图像文件,但我们的模型期望一个形状为 (N, 3, 224, 224) 的 PyTorch 张量, 其中 N 是输入批次的数量。(我们将只使用批量大小为 1。)我们首先要做的是组合一组 TorchVision 转换, 调整图像大小和裁剪图像,将其转换为张量,然后对张量中的值进行归一化。 (有关此归一化的更多信息,请参阅 torchvision.models_ 的文档。)

之后,我们打开文件并应用转换。转换返回一个形状为 (3, 224, 224) 的张量 - 224x224 图像的 3 个颜色通道。因为我们需要将这个单个图像变成一个批次,所以我们使用 unsqueeze_(0) 调用通过添加一个新的第一维来就地修改张量。张量包含相同的数据,但现在形状为 (1, 3, 224, 224)。

一般来说,即使您不是在处理图像数据,您也需要将来自 HTTP 请求的输入转换为 PyTorch 可以使用的张量。

推理¶

def get_prediction(input_tensor):
    outputs = model.forward(input_tensor)
    _, y_hat = outputs.max(1)
    prediction = y_hat.item()
    return prediction

推理本身是最简单的部分:当我们将输入张量传递给模型时,我们会得到一个张量值,代表模型估计图像属于特定类别的可能性。 max() 调用找到具有最大可能性值的类别,并返回该值及其 ImageNet 类索引。最后,我们使用 item() 调用从包含它的张量中提取该类索引,并返回它。

后处理¶

def render_prediction(prediction_idx):
    stridx = str(prediction_idx)
    class_name = 'Unknown'
    if img_class_map is not None:
        if stridx in img_class_map is not None:
            class_name = img_class_map[stridx][1]

    return prediction_idx, class_name

The render_prediction() method maps the predicted class index to a human-readable class label. It’s typical, after getting the prediction from your model, to perform post-processing to make the prediction ready for either human consumption, or for another piece of software.

render_prediction() 方法将预测的类索引映射为人类可读的类标签。在从您的模型获得预测之后,通常需要进行后处理, 使预测可供人类使用或供另一个软件使用。

运行完整的 Flask 应用¶

将以下内容粘贴到名为 app.py 的文件中:

import io
import json
import os

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


app = Flask(__name__)
model = models.densenet121(pretrained=True)               # 在 ImageNet 的 1000 个类别上训练
model.eval()                                              # 关闭自动梯度计算



img_class_map = None
mapping_file_path = 'index_to_name.json'                  # ImageNet 类别的可读名称
if os.path.isfile(mapping_file_path):
    with open (mapping_file_path) as f:
        img_class_map = json.load(f)



# 将输入转换为模型期望的形式
def transform_image(infile):
    input_transforms = [transforms.Resize(255),           # 我们使用多个 TorchVision 转换来准备图像
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],       # ImageNet 模型输入的标准归一化
            [0.229, 0.224, 0.225])]
    my_transforms = transforms.Compose(input_transforms)
    image = Image.open(infile)                            # 打开图像文件
    timg = my_transforms(image)                           # 将 PIL 图像转换为合适形状的 PyTorch 张量
    timg.unsqueeze_(0)                                    # PyTorch 模型期望批量输入;创建批量大小为 1
    return timg


# 获取预测
def get_prediction(input_tensor):
    outputs = model.forward(input_tensor)                 # 获取所有 ImageNet 类别的可能性
    _, y_hat = outputs.max(1)                             # 提取最可能的类别
    prediction = y_hat.item()                             # 从 PyTorch 张量中提取 int 值
    return prediction

# 使预测结果可读
def render_prediction(prediction_idx):
    stridx = str(prediction_idx)
    class_name = 'Unknown'
    if img_class_map is not None:
        if stridx in img_class_map is not None:
            class_name = img_class_map[stridx][1]

    return prediction_idx, class_name


@app.route('/', methods=['GET'])
def root():
    return jsonify({'msg' : 'Try POSTing to the /predict endpoint with an RGB image attachment'})


@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        file = request.files['file']
        if file is not None:
            input_tensor = transform_image(file)
            prediction_idx = get_prediction(input_tensor)
            class_id, class_name = render_prediction(prediction_idx)
            return jsonify({'class_id': class_id, 'class_name': class_name})


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

从 shell 提示符启动服务器,请执行以下命令:

FLASK_APP=app.py flask run

默认情况下,您的 Flask 服务器监听 5000 端口。服务器运行后,打开另一个终端窗口,测试您新的推理服务器:

curl -X POST -H "Content-Type: multipart/form-data" http://localhost:5000/predict -F "file=@kitten.jpg"

如果一切设置正确,您应该会收到类似如下的响应:

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

重要资源¶

pytorch.org 提供安装说明,以及更多文档和教程
`Flask 官网`_ 有一个 `快速入门指南`_ ,对设置一个简单的 Flask 服务有更详细的介绍