Forward项目ONNX模型推理使用指南

前言

Forward是一个高性能的深度学习推理框架，支持多种模型格式的推理加速。本文将重点介绍如何使用Forward框架进行ONNX模型的推理，包括环境准备、模型导出、项目构建以及各种精度模式下的推理实现。

环境准备

在使用Forward进行ONNX模型推理前，需要确保系统满足以下环境要求：

• GPU环境：

  • NVIDIA CUDA >= 10.0
  • CuDNN >= 7
  • 推荐版本：CUDA 10.2

• 推理框架：

  • TensorRT >= 7.0.0.11
  • 推荐版本：TensorRT-7.2.1.6

• 构建工具：

  • CMake >= 3.12.2
  • GCC >= 5.4.0
  • ld >= 2.26.1

ONNX模型导出

ONNX(Open Neural Network Exchange)是一种开放的模型格式，支持跨框架的模型转换。我们可以将训练好的PyTorch、TensorFlow等框架的模型导出为ONNX格式。

PyTorch模型导出示例

import torch
import torch.onnx
import torchvision.models as models

# 加载预训练模型
model = models.resnet50(pretrained=True)
model.cpu()
model.eval()

# 准备示例输入
inputs = torch.randn(1, 3, 224, 224)

# 导出JIT模型用于验证
traced_model = torch.jit.trace(model, inputs)
torch.jit.save(traced_model, 'resnet50.pth')

# 导出ONNX模型
input_names = ["input"]  # 输入节点名称
output_names = ["output"]  # 输出节点名称
torch.onnx.export(model, inputs, 'resnet50.onnx', 
                 verbose=True, 
                 input_names=input_names, 
                 output_names=output_names)

注意事项：

1. 导出前务必将模型设置为eval模式
2. 输入输出名称可通过可视化工具查看
3. 建议同时导出JIT模型用于精度验证

项目构建

Forward支持灵活的构建选项，可以根据需求启用或禁用特定功能：

mkdir build
cd build
cmake .. \
-DTensorRT_ROOT="TensorRT安装路径" \
-DENABLE_ONNX=ON \  # 启用ONNX支持
-DENABLE_DYNAMIC_BATCH=ON \  # 启用动态批量支持
-DENABLE_UNIT_TESTS=ON  # 启用单元测试

make -j

关键构建选项说明：

• TensorRT_ROOT：指定TensorRT的安装路径
• ENABLE_ONNX：启用ONNX模型支持
• ENABLE_DYNAMIC_BATCH：启用动态批量处理功能
• BUILD_PYTHON_LIB：是否构建Python接口

动态批量处理

Forward支持动态批量输入，这在处理可变大小输入时非常有用。

使用要求

1. TensorRT版本需大于7.1.x.x
2. INT8模式下也可使用动态批量

配置方式

CMake配置：

-DENABLE_DYNAMIC_BATCH=ON

Builder配置：

• C++接口：

  onnx_builder.SetOptBatchSize(opt_batch_size);
  

• Python接口：

  onnx_builder.set_opt_batch_size(opt_batch_size)
  

ONNX模型要求

1. 导出ONNX模型时需要指定动态维度：

   dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
   torch.onnx.export(..., dynamic_axes=dynamic_axes)
   

2. 仅支持batch_size作为动态维度
3. batch_size必须是输入的第一个维度

C++接口使用

基本使用流程

// 1. 创建Builder
fwd::OnnxBuilder builder;
builder.SetInferMode("float32");  // 设置推理精度

// 2. 构建Engine
auto engine = builder.Build("model.onnx");

// 3. 准备输入数据
std::vector<float> input_data(1*3*224*224, 0);
fwd::Tensor input = {input_data.data(), {1,3,224,224}, fwd::DataType::FLOAT, fwd::DeviceType::CPU};

// 4. 执行推理
std::vector<fwd::Tensor> outputs;
engine->Forward({input}, outputs);

// 5. 处理输出
for(auto& out : outputs) {
    std::vector<float> host_data(out.dims[0]*out.dims[1]*out.dims[2]*out.dims[3]);
    MemcpyDeviceToHost(host_data.data(), out.data, host_data.size());
}

INT8量化推理

// 1. 实现数据流接口
class MyBatchStream : public IBatchStream {
    // 实现next、getBatch等方法
};

// 2. 创建量化器
auto stream = std::make_shared<MyBatchStream>();
auto calibrator = std::make_shared<TrtInt8Calibrator>(stream, "calib.cache", "entropy");

// 3. 构建INT8引擎
builder.SetCalibrator(calibrator);
builder.SetInferMode("int8");
auto engine = builder.Build("model.onnx");

BERT模型INT8量化

BERT模型量化需要分两步进行：

// 1. 生成校准文件
builder.SetInferMode("int8_calib");
auto calib_engine = builder.Build("bert.onnx");

// 2. 使用校准文件构建推理引擎
builder.SetInferMode("int8");
auto infer_engine = builder.Build("bert.onnx");

Python接口使用

基本使用示例

import forward
import numpy as np

# 1. 创建Builder
builder = forward.OnnxBuilder()
builder.set_mode("float32")  # 设置推理精度

# 2. 构建Engine
engine = builder.build("model.onnx")

# 3. 执行推理
input_data = np.random.rand(1,3,224,224).astype(np.float32)
outputs = engine.forward([input_data])

INT8量化推理

class MyBatchStream(forward.IPyBatchStream):
    # 实现数据流接口
    pass

# 创建量化器
stream = MyBatchStream()
calibrator = forward.TrtInt8Calibrator(stream, "calib.cache", "minmax")

# 构建INT8引擎
builder.set_calibrator(calibrator)
builder.set_mode("int8")
engine = builder.build("model.onnx")

手动量化Scale设置

在某些场景下，我们可能需要手动设置各层的量化scale值：

1. 准备scale文件，格式为：

   LayerName: scale_value
   

2. C++使用方式：

   auto calibrator = std::make_shared<TrtInt8Calibrator>("cache.calib", "entropy", 1);
   calibrator->setScaleFile("scales.txt");
   

3. Python使用方式：

   calibrator = forward.TrtInt8Calibrator("cache.calib", "entropy", 1)
   calibrator.set_scale_file("scales.txt")
   

最佳实践建议

1. 模型导出：

   • 导出ONNX模型时建议同时导出JIT模型用于精度验证
   • 确保输入输出名称设置正确

2. 量化校准：

   • 校准数据集应具有代表性
   • BERT模型建议使用minmax校准方法
   • 普通CNN模型可以使用entropy校准方法

3. 性能优化：

   • 合理设置opt_batch_size以获得最佳性能
   • FP16模式通常能在保持精度的同时提升性能
   • INT8模式可显著提升性能但可能影响精度

4. 错误排查：

   • 检查TensorRT版本是否匹配
   • 确保CUDA环境配置正确
   • 验证模型输入输出是否符合预期

通过本文介绍，您应该已经掌握了使用Forward框架进行ONNX模型推理的完整流程。根据实际应用场景选择合适的精度模式和配置，可以获得最佳的推理性能。