介紹

在這個教程中，我們將學習如何將訓練好的 PyTorch 模型部署到生產環境中，並創建一個簡單的在線推理服務。我們將使用 FastAPI 框架和 Uvicorn 服務器來搭建一個 REST API，實現對模型的在線推理。本教程將涵蓋模型導出、Web服務搭建、API設計以及在線推理等關鍵步驟。

步驟 1：導出 PyTorch 模型

首先，我們需要導出已經訓練好的 PyTorch 模型，以便在生產環境中使用。使用 torch.jit 模塊將模型轉換為 TorchScript 格式：

import torch
# 示例模型，你需要替換成你自己的模型
class SampleModel(torch.nn.Module):
    def forward(self, x):
        return x * 2
# 創建模型實例
model = SampleModel()
# 保存模型為 TorchScript 格式
torch.jit.save(torch.jit.script(model), 'sample_model.pt')

步驟 2：安裝 FastAPI 和 Uvicorn

在搭建 Web 服務之前，我們需要安裝 FastAPI 和 Uvicorn。使用以下命令安裝：

pip install fastapi uvicorn

步驟 3：創建 FastAPI 應用

from fastapi import FastAPI
# 創建 FastAPI 應用
app = FastAPI()
# 定義模型推理接口
@app.post("/predict")
async def predict(input_data: int):
    # 調用模型進行推理
    result = model(torch.tensor(input_data))
    return {"prediction": result.item()}

步驟 4：運行 FastAPI 服務

使用 Uvicorn 啟動 FastAPI 服務：

uvicorn your_app_name:app --reload

步驟 5：測試推理服務

使用工具如 curl 或者 Postman 來測試推理服務：

curl -X POST "http://127.0.0.1:8000/predict" -H "Content-Type: application/json" -d '{"input_data": 5}'

總結

通過這個教程，你學會了如何將 PyTorch 模型導出為 TorchScript 格式，並使用 FastAPI 和 Uvicorn 創建一個簡單的在線推理服務。在實際項目中，你可以根據需求擴展接口、增加模型加載邏輯、加入輸入參數驗證等功能，以滿足生產環境的需求。