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
        my_module = MyModule(10, 20)
        sm = torch.jit.script(my_module)
    

• 保存序列化模型： 使用save方法将scriptmodule保存为文件。

  sm.save("model.pt")
  

• 在C++中加载序列化的PyTorch模型： 使用LibTorch库加载保存的模型文件并进行推理。

  #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： 调用module的forward方法进行推理。

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

使用Flask或FastAPI等Web框架部署

如果你想通过网络提供服务，可以使用Flask或FastAPI等Web框架来部署你的模型。

• 使用Flask：

  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()
  
  if torch.cuda.is_available():
      model.cuda()
  
  @app.route('/predict', methods=['POST'])
  def predict():
      data = request.json['input_data']
      input_tensor = torch.tensor(data).unsqueeze(0)
      if torch.cuda.is_available():
          input_tensor = input_tensor.cuda()
      with torch.no_grad():
          output = model(input_tensor)
      return jsonify(output.tolist())
  
  if __name__ == '__main__':
      app.run(host='0.0.0.0', port=5000)
  

  运行Flask应用：

  python app.py
  

  然后你可以通过HTTP请求发送数据到http://<your-server-ip>:5000/predict来获取预测结果。

使用Docker进行部署

Docker是一个容器化平台，可以让你轻松地打包和部署应用程序。

• 安装Docker：

  sudo apt update
  sudo apt install docker.io
  

• 创建Dockerfile：

  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
  

使用TorchServe进行部署

TorchServe是PyTorch的模型服务工具，可以部署TorchScript模型。

• 安装TorchServe：

  pip install torchserve
  torchserve --start --model_name=model --model_path=/path/to/model.pt --runtime_version=1.8
  

• 使用Flask或Django： 你可以将模型部署为一个Web服务，使用Flask或Django框架。

以上步骤可以帮助你在Linux系统下成功部署PyTorch模型。根据你的具体需求，可以选择使用TorchServe、Flask或Django等进行部署。