Linux上PyTorch模型如何部署

在Linux上部署PyTorch模型有多種方法，以下提供兩種常見的部署流程：

1. 使用TorchScript進行部署

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

步驟：

1. 模型轉換：

• 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)
  

2. 保存序列化模型：

使用save方法將scriptmodule保存為文件。

sm.save("model.pt")

3. 在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
}

4. 執行Script Module：

調用module的forward方法進行推理。

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

2. 使用ONNX進行部署

ONNX（Open Neural Network Exchange）是一種開放格式，用于表示深度學習模型。通過將PyTorch模型轉換為ONNX格式，可以在多種框架上進行部署。

步驟：

1. 安裝ONNX和ONNX Runtime：

pip install onnx onnxruntime

2. 將PyTorch模型轉換為ONNX格式：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.onnx import export

# 定義一個簡單的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

# 實例化模型、優化器和損失函數
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# 導出模型為ONNX格式
dummy_input = torch.randn(1, 1, 28, 28)
export(model, dummy_input, "simple_model.onnx", verbose=True)

3. 在Python中使用ONNX Runtime進行推理：

import onnx
import onnxruntime as ort

# 加載ONNX模型
onnx_model = onnx.load("simple_model.onnx")
onnx.checker.check_model(onnx_model)

# 創建ONNX Runtime會話
ort_session = ort.InferenceSession("simple_model.onnx")

# 準備輸入數據
input_data = {ort_session.get_inputs()[0].name: dummy_input.numpy()}

# 運行推理
outputs = ort_session.run(None, input_data)

4. 在C++中使用ONNX Runtime進行推理：

#include <onnxruntime_cpp.h>

int main() {
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    session_options.SetInterOpNumThreads(1);

    Ort::Session session(env, "simple_model.onnx", session_options);

    // 準備輸入數據
    std::vector<float> input_data = {/* 輸入數據 */};
    std::vector<int64_t> input_shape = {1, 1, 28, 28};

    // 運行推理
    std::vector<float> output_data(10);
    session.Run(Ort::RunOptions{nullptr}, input_names, &input_data[0], input_shape.data(), input_shape.size(), output_names, &output_data[0], output_data.size());

    // 處理輸出數據
}

以上是兩種在Linux上部署PyTorch模型的常見方法，具體選擇哪種方法取決于你的應用場景和需求。