Linux上PyTorch模型如何部署

在Linux上部署PyTorch模型有多种方法，以下提供两种常见的部署流程：

1. 使用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)

保存序列化模型：

使用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();

2. 使用ONNX进行部署

ONNX（Open Neural Network Exchange）是一种开放格式，用于表示深度学习模型。通过将PyTorch模型转换为ONNX格式，可以在多种框架上进行部署。

步骤：

安装ONNX和ONNX Runtime：

pip install onnx onnxruntime

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

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

在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模型的常见方法，具体选择哪种方法取决于你的应用场景和需求。

1. 使用TorchScript进行部署

步骤：

2. 使用ONNX进行部署

步骤：

最新问答

相关标签