Linux下PyTorch模型如何部署

在Linux下部署PyTorch模型有多種方法，以下是一些常見的部署流程：

使用TorchScript進行部署

TorchScript是PyTorch的一種序列化格式，可用于在C++、Python等平臺上進行部署。

• 模型轉換：
  • Tracing：適用于模型中沒有控制流的情況。通過給模型一組輸入，走一遍推理網絡，然后用torch.jit.trace記錄路徑上的信息并保存為torch.jit.scriptmodule。

    import torch
    import torchvision
    
    model = torchvision.models.resnet18()
    example = torch.rand(1, 3, 224, 224)
    traced_script_module = torch.jit.trace(model, example)
    

  • Scripting：適用于有控制流的模型。在Torch腳本中編寫模型并相應地注釋模型，通過torch.jit.script編譯模塊，將其轉換為scriptmodule。

    class MyModule(torch.nn.Module):
        def __init__(self, n, m):
            super(MyModule, self).__init__()
            self.weight = torch.nn.Parameter(torch.rand(n, m))
    
        def forward(self, input):
            if input.sum() > 0:
                output = self.weight.mv(input)
            else:
                output = self.weight + input
            return output
    
    my_module = MyModule(10, 20)
    sm = torch.jit.script(my_module)
    

  • 保存序列化模型：

    sm.save("model.pt")
    

  • 在C++中加載序列化的PyTorch模型：

    #include <torch/torch.h>
    
    int main() {
        torch::jit::script::Module module;
        try {
            module = torch::jit::load("model.pt");
        }
        catch (const c10::Error& e) {
            std::cerr << "error loading the model\n";
            return -1;
        }
        // Use the module for inference
    }
    

  • 執行Script Module：

    at::Tensor input = torch::rand({1, 3, 224, 224});
    at::Tensor output = module.forward({input}).toTensor();
    

使用Web框架部署（可選）

如果你想通過網絡提供服務，可以使用Flask或FastAPI等Web框架來部署你的模型。以下是一個使用Flask的簡單示例：

• 安裝Flask：

  pip install flask
  

• 創建Flask應用：

  from flask import Flask, request, jsonify
  import torch
  import torchvision.transforms as transforms
  from PIL import Image
  
  app = Flask(__name__)
  # 加載模型
  model = torch.load('path_to_your_model.pt', map_location=torch.device('cpu'))
  model.eval()
  # 定義圖像預處理
  transform = transforms.Compose([
      transforms.Resize((224, 224)),
      transforms.ToTensor(),
      transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
  ])
  
  @app.route('/predict', methods=['POST'])
  def predict():
      if request.method == 'POST':
          file = request.files['image']
          image = Image.open(file.stream)
          image = transform(image).unsqueeze(0)  # 添加batch維度
          with torch.no_grad():
              output = model(image)
          _, predicted_idx = torch.max(output, 1)
          return jsonify({'prediction': int(predicted_idx.item())})
  
  if __name__ == '__main__':
      app.run(host='0.0.0.0', port=5000)
  

• 運行Flask應用：

  python app.py
  

• 測試模型： 使用curl或Postman來測試你的模型。

  curl -X POST -F "image=@path_to_your_image.jpg" http://localhost:5000/predict
  

使用Docker進行部署

Docker是一個容器化平臺，可以讓你輕松地打包和部署應用程序。

• 安裝Docker：

  sudo apt update
  sudo apt install docker.io
  

• 創建Dockerfile：

  # 使用官方的PyTorch鏡像作為基礎鏡像
  FROM pytorch/pytorch:latest
  # 設置工作目錄
  WORKDIR /app
  # 復制當前目錄下的所有文件到工作目錄
  COPY . /app
  # 安裝依賴項
  RUN pip install --no-cache-dir -r requirements.txt
  # 暴露端口(如果需要)
  EXPOSE 80
  # 運行應用
  CMD ["python", "your_script.py"]
  

• 創建requirements.txt：

  torch
  torchvision
  numpy
  

• 構建Docker鏡像：

  docker build -t your-image-name .
  

• 運行Docker容器：

  docker run -p 80:80 your-image-name
  

以上是Linux下部署PyTorch模型的幾種常見方法，你可以根據自己的需求選擇合適的方式進行部署。