PyTorch在Ubuntu上的模型部署方法

以下是在Ubuntu上部署PyTorch模型的主要方法，步驟基于搜索資源整合：

一、環境準備

1. 安裝系統依賴
   • 更新系統：sudo apt update && sudo apt upgrade
   • 安裝NVIDIA驅動、CUDA和cuDNN（需匹配PyTorch版本）。
2. 創建虛擬環境
   • 使用Conda或venv隔離環境：

     # Conda  
     conda create -n pytorch_env python=3.8  
     conda activate pytorch_env  
     # 或venv  
     python3 -m venv pytorch_env  
     source pytorch_env/bin/activate  
     

3. 安裝PyTorch
   • CPU版本：pip install torch torchvision torchaudio
   • GPU版本：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113（需替換為實際CUDA版本）。

二、模型準備

• 訓練并保存模型為.pt或.pth文件，例如：

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

三、部署方式

1. 直接運行腳本（輕量級）

• 編寫Python腳本加載模型并處理輸入：

  import torch  
  from model import MyModel  # 假設模型定義在此文件中  
  
  model = MyModel()  
  model.load_state_dict(torch.load('model.pth'))  
  model.eval()  
  
  # 示例輸入（需根據模型調整）  
  input_data = torch.randn(1, 3, 224, 224)  
  with torch.no_grad():  
      output = model(input_data)  
  print(output)  
  

  運行：python deploy.py。

2. Web服務部署（Flask/FastAPI）

• Flask示例：
  1. 安裝Flask：pip install flask
  2. 創建app.py：

     from flask import Flask, request, jsonify  
     import torch  
     from model import MyModel  
     
     app = Flask(__name__)  
     model = MyModel()  
     model.load_state_dict(torch.load('model.pth'))  
     model.eval()  
     
     @app.route('/predict', methods=['POST'])  
     def predict():  
         data = request.json['input']  # 假設輸入為JSON格式  
         input_tensor = torch.tensor(data).unsqueeze(0)  
         with torch.no_grad():  
             output = model(input_tensor)  
         return jsonify({'output': output.tolist()})  
     
     if __name__ == '__main__':  
         app.run(host='0.0.0.0', port=5000)  
     

  3. 運行服務：python app.py，通過HTTP請求調用。

3. 生產環境部署

• Gunicorn + Flask：
  安裝Gunicorn：pip install gunicorn，運行：
  gunicorn -w 4 -b 0.0.0.0:5000 app:app（-w指定進程數）。
• Docker容器化：
  1. 編寫Dockerfile：

     FROM pytorch/pytorch:latest  
     COPY . /app  
     WORKDIR /app  
     RUN pip install -r requirements.txt  
     CMD ["python", "app.py"]  
     

  2. 構建并運行：

     docker build -t pytorch-model .  
     docker run -p 5000:5000 pytorch-model  
     ```。  
     
     

4. 高性能部署（TorchServe/ONNX）

• TorchServe（官方推薦）：
  1. 安裝TorchServe：pip install torchserve torch-model-archiver
  2. 將模型轉換為TorchScript格式：

     scripted_model = torch.jit.script(model)  
     scripted_model.save('model.ts')  
     

  3. 啟動服務：
     torchserve --start --model_name=my_model --model_path=/path/to/model.ts。
• ONNX轉換：
  將模型導出為ONNX格式，部署到支持ONNX的推理引擎（如TensorRT）：
  torch.onnx.export(model, input, 'model.onnx')。

四、驗證與監控

• 通過curl或Postman測試API接口：

  curl -X POST -H "Content-Type: application/json" -d '{"input": [1,2,3]}' http://localhost:5000/predict  
  

• 生產環境中建議使用Nginx反向代理、Prometheus監控服務狀態。

注意事項

• GPU部署需確保CUDA/cuDNN版本與PyTorch兼容。
• 大規模部署建議使用Docker + Kubernetes實現彈性擴縮容。

參考來源：