在機(jī)器學(xué)習(xí)項(xiàng)目的實(shí)際應(yīng)用中，將訓(xùn)練好的模型部署為服務(wù)，使其能夠接收外部請(qǐng)求并返回預(yù)測(cè)結(jié)果，是實(shí)現(xiàn)模型價(jià)值的關(guān)鍵一步。Flask 作為 Python 的輕量級(jí) Web 框架，憑借其簡(jiǎn)潔易用的特性，成為部署 PyTorch 模型的理想選擇之一。本文將詳細(xì)指導(dǎo)您如何使用 Flask 將 PyTorch 模型部署為 REST API 服務(wù)，以預(yù)訓(xùn)練的 DenseNet 121 模型為例，實(shí)現(xiàn)圖像分類功能。

一、環(huán)境搭建與依賴安裝

在開(kāi)始部署之前，確保已安裝所需的依賴庫(kù)。運(yùn)行以下命令以安裝 Flask 和 torchvision：

pip install Flask torchvision

二、創(chuàng)建簡(jiǎn)單的 Web 服務(wù)器

首先，我們創(chuàng)建一個(gè)基本的 Flask Web 服務(wù)器，后續(xù)將在此基礎(chǔ)上添加模型推理功能。

from flask import Flask


app = Flask(__name__)


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


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

保存上述代碼為 app.py，運(yùn)行 Flask 開(kāi)發(fā)服務(wù)器：

FLASK_ENV=development FLASK_APP=app.py flask run

訪問(wèn) http://localhost:5000/，您將看到 "Hello World!" 文字，這表明服務(wù)器已成功啟動(dòng)。

三、定義 API 端點(diǎn)與推理邏輯

我們將定義一個(gè) /predict 端點(diǎn)，用于接收包含圖像文件的 HTTP POST 請(qǐng)求，并返回預(yù)測(cè)結(jié)果。

（一）圖像預(yù)處理

DenseNet 121 模型要求輸入圖像為 224 x 224 的 3 通道 RGB 圖像，且需進(jìn)行歸一化處理。我們使用 torchvision.transforms 構(gòu)建圖像預(yù)處理管道：

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)

（二）加載預(yù)訓(xùn)練模型

加載預(yù)訓(xùn)練的 DenseNet 121 模型，并設(shè)置為評(píng)估模式：

from torchvision import models


model = models.densenet121(pretrained=True)
model.eval()

（三）圖像分類預(yù)測(cè)

編寫函數(shù)以獲取圖像的預(yù)測(cè)類別：

import json


imagenet_class_index = json.load(open('imagenet_class_index.json'))  # 請(qǐng)?zhí)鎿Q為實(shí)際文件路徑


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

四、整合 Flask API 服務(wù)器

將模型推理功能整合到 Flask 服務(wù)器中，完成 API 的定義：

from flask import jsonify, request


@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()

完整的 app.py 文件如下：

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('imagenet_class_index.json'))  # 請(qǐng)?zhí)鎿Q為實(shí)際文件路徑
model = models.densenet121(pretrained=True)
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)
    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()

五、測(cè)試與驗(yàn)證

運(yùn)行 Flask 服務(wù)器：

FLASK_ENV=development FLASK_APP=app.py flask run

使用 requests 庫(kù)發(fā)送 POST 請(qǐng)求進(jìn)行測(cè)試：

import requests


resp = requests.post("http://localhost:5000/predict", files={"file": open('cat.jpg', 'rb')})  # 請(qǐng)?zhí)鎿Q為實(shí)際圖像文件路徑
print(resp.json())

成功返回結(jié)果示例：

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

六、總結(jié)與優(yōu)化建議

通過(guò)本文，您已成功使用 Flask 部署了一個(gè) PyTorch 模型，并通過(guò) REST API 提供圖像分類服務(wù)。然而，當(dāng)前的實(shí)現(xiàn)較為基礎(chǔ)，對(duì)于生產(chǎn)環(huán)境，您可以考慮以下優(yōu)化措施：

• 增強(qiáng)錯(cuò)誤處理機(jī)制 ：對(duì)請(qǐng)求中未包含圖像文件、文件類型錯(cuò)誤等情況進(jìn)行處理，返回友好的錯(cuò)誤提示信息。
• 提升性能與擴(kuò)展性 ：采用更高效的服務(wù)器架構(gòu)（如 Gunicorn）替代開(kāi)發(fā)服務(wù)器，結(jié)合模型量化、剪枝等技術(shù)優(yōu)化模型性能，以適應(yīng)高并發(fā)場(chǎng)景。
• 添加認(rèn)證與授權(quán) ：為 API 添加身份驗(yàn)證，確保服務(wù)的安全性。
• 構(gòu)建前端界面 ：開(kāi)發(fā)一個(gè)簡(jiǎn)單的 Web 界面，允許用戶上傳圖像并顯示預(yù)測(cè)結(jié)果，提升用戶體驗(yàn)。

模型部署是連接模型開(kāi)發(fā)與實(shí)際應(yīng)用的橋梁，掌握這一技能，能夠使您的模型真正發(fā)揮價(jià)值，解決實(shí)際問(wèn)題。編程獅將持續(xù)為您帶來(lái)更多模型部署與應(yīng)用開(kāi)發(fā)的實(shí)用教程。