TorchScript与ONNX的推理速度及部署方案对比咨询

一、关于配置难度的纠正 #

你的初步印象部分正确，补充两点关键信息：

• ONNX确实需要编写转换代码（torch.onnx.export），且要处理算子兼容性问题（部分PyTorch算子ONNX原生不支持，需替换逻辑或自定义导出规则），推理时还需集成ONNX Runtime等第三方引擎，整体配置步骤更多。
• TorchScript也需要导出步骤（torch.jit.trace/script），但属于PyTorch原生方案，无需处理跨框架兼容，配置更轻量化，适合纯PyTorch技术栈的场景。

二、推理速度差异 #

两者的速度表现取决于模型复杂度、推理硬件及优化策略：

• CPU场景：简单模型下，TorchScript（依托LibTorch）和ONNX Runtime速度差距极小；复杂模型（如大卷积网络、Transformer）中，ONNX Runtime通过图优化、算子融合、量化等技术，通常能实现比TorchScript更高的推理效率。
• GPU场景：若用NVIDIA GPU，TorchScript依托CUDA原生支持，性能稳定；ONNX Runtime通过TensorRT后端优化，可能在大模型上跑出比原生TorchScript更快的速度，但配置复杂度更高。
• 总结：追求开箱即用的稳定性能选TorchScript；追求极致性能且愿意做额外优化配置，选ONNX+对应推理引擎。

三、Python存模型、C++加载推理的具体方案 #

方案1：TorchScript + LibTorch #

1. Python端导出模型

import torch
import torchvision.models as models

# 加载预训练/自定义模型并设为评估模式
model = models.resnet18(pretrained=True)
model.eval()

# 静态结构模型用trace导出，动态结构模型用script导出
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
# 保存为.pt格式的TorchScript模型
traced_model.save("resnet18_traced.pt")

2. C++端加载推理（依赖LibTorch库）

#include <torch/torch.h>
#include <iostream>

int main() {
    // 加载TorchScript模型
    torch::jit::script::Module module;
    try {
        module = torch::jit::load("resnet18_traced.pt");
    } catch (const c10::Error& e) {
        std::cerr << "模型加载失败: " << e.what() << std::endl;
        return -1;
    }
    // 切换到GPU（可选，需CUDA环境）
    module.to(torch::kCUDA);
    module.eval();

    // 构造输入张量
    torch::Tensor input = torch::randn({1, 3, 224, 224}).to(torch::kCUDA);
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(input);

    // 执行推理并处理输出
    at::Tensor output = module.forward(inputs).toTensor();
    int pred = output.argmax(1).item<int>();
    std::cout << "预测类别: " << pred << std::endl;

    return 0;
}

方案2：ONNX + ONNX Runtime #

1. Python端导出模型

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

example_input = torch.randn(1, 3, 224, 224)
# 导出ONNX模型，可指定动态轴支持可变batch size
torch.onnx.export(
    model,
    example_input,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

2. C++端加载推理（依赖ONNX Runtime库）

#include <onnxruntime_cxx_api.h>
#include <iostream>
#include <vector>

int main() {
    // 初始化ONNX Runtime环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ResNet18Inference");
    Ort::SessionOptions session_options;
    // 开启GPU加速（可选，需ONNX Runtime GPU版本）
    // session_options.SetGraphOptimizationLevel(ORT_ENABLE_ALL);
    // Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0));

    // 加载ONNX模型
    Ort::Session session(env, "resnet18.onnx", session_options);

    // 获取输入输出名称与形状
    Ort::AllocatorWithDefaultOptions allocator;
    const char* input_name = session.GetInputName(0, allocator);
    const char* output_name = session.GetOutputName(0, allocator);
    auto input_shape = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
    input_shape[0] = 1; // 设置batch size为1

    // 准备输入数据（示例用随机数）
    size_t input_size = 1;
    for (auto dim : input_shape) input_size *= dim;
    std::vector<float> input_data(input_size);
    for (size_t i = 0; i < input_size; ++i) {
        input_data[i] = static_cast<float>(rand()) / RAND_MAX;
    }

    // 创建输入张量并执行推理
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_data.data(), input_size, input_shape.data(), input_shape.size()
    );
    std::vector<const char*> input_names = {input_name};
    std::vector<const char*> output_names = {output_name};
    auto output_tensor = session.Run(
        Ort::RunOptions{nullptr}, input_names.data(), &input_tensor, 1, output_names.data(), 1
    );

    // 处理输出结果
    float* output_data = output_tensor[0].GetTensorMutableData<float>();
    int pred = std::max_element(output_data, output_data + 1000) - output_data;
    std::cout << "预测类别: " << pred << std::endl;

    return 0;
}

选择建议 #

• 若技术栈以PyTorch为主，且无需跨框架部署，优先选TorchScript+LibTorch，配置简单、兼容性拉满。
• 若需要跨框架共享模型，或追求极致推理性能，选ONNX+ONNX Runtime/TensorRT，灵活性更高。

内容的提问来源于stack exchange，提问作者Reza