突破动态瓶颈：TensorRT ONNX动态形状全自动化测试实践指南

【免费下载链接】TensorRT NVIDIA® TensorRT™ 是一个用于在 NVIDIA GPU 上进行高性能深度学习推理的软件开发工具包（SDK）。此代码库包含了 TensorRT 的开源组件 项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT

引言：动态形状的工业级挑战

在深度学习推理部署中，固定输入形状的模型往往无法满足实际业务需求。用户上传的图片尺寸不一、文本长度各异，这些现实场景都要求模型能够处理动态变化的输入形状。NVIDIA® TensorRT™作为高性能深度学习推理SDK，提供了强大的动态形状支持，但如何确保动态形状下的推理准确性和性能稳定性，一直是工程师们面临的棘手问题。

本文将以MNIST手写数字识别为案例，详细介绍如何利用TensorRT构建支持动态形状的ONNX模型推理引擎，并实现全自动化测试流程。通过本文，你将掌握：

• 动态形状推理的核心原理与实现方法
• TensorRT优化配置文件的编写技巧
• 自动化测试框架的搭建与验证策略
• 性能调优与常见问题解决方案

动态形状推理原理与实现

核心概念：动态维度与优化配置文件

TensorRT通过优化配置文件（PluginConfig.yaml）定义动态形状的范围，如sampleDynamicReshape所示：

auto profile = builder->createOptimizationProfile();
profile->setDimensions(input->getName(), OptProfileSelector::kMIN, Dims4{1, 1, 1, 1});
profile->setDimensions(input->getName(), OptProfileSelector::kOPT, Dims4{1, 1, 28, 28});
profile->setDimensions(input->getName(), OptProfileSelector::kMAX, Dims4{1, 1, 56, 56});
preprocessorConfig->addOptimizationProfile(profile);

这段代码定义了输入形状的最小值(1x1)、最优值(28x28)和最大值(56x56)，TensorRT将针对该范围内的所有形状进行优化。

动态形状预处理网络构建

在实际应用中，我们通常需要构建一个预处理网络来调整输入形状，使其符合模型要求：

auto input = preprocessorNetwork->addInput("input", nvinfer1::DataType::kFLOAT, Dims4{-1, 1, -1, -1});
auto resizeLayer = preprocessorNetwork->addResize(*input);
resizeLayer->setOutputDimensions(mPredictionInputDims);
preprocessorNetwork->markOutput(*resizeLayer->getOutput(0));

其中-1表示动态维度，将在运行时由实际输入决定。

自动化测试框架搭建

测试用例设计

为确保动态形状推理的正确性，我们需要设计覆盖各种边界情况的测试用例：

test_cases = [
    {"input_shape": (1, 1, 28, 28), "expected_output": 7},  # 标准MNIST尺寸
    {"input_shape": (1, 1, 14, 14), "expected_output": 3},  # 小于最优尺寸
    {"input_shape": (1, 1, 56, 56), "expected_output": 9},  # 等于最大尺寸
    {"input_shape": (1, 1, 32, 32), "expected_output": 2},  # 中间尺寸
]

自动化测试流程实现

利用TensorRT Python API构建自动化测试框架：

def test_dynamic_shape(engine_path, test_cases):
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    runtime = trt.Runtime(TRT_LOGGER)
    
    with open(engine_path, "rb") as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    
    results = []
    for case in test_cases:
        input_shape = case["input_shape"]
        expected = case["expected_output"]
        
        # 创建执行上下文并设置输入形状
        context = engine.create_execution_context()
        context.set_binding_shape(0, input_shape)
        
        # 生成随机输入数据
        input_data = np.random.rand(*input_shape).astype(np.float32)
        
        # 执行推理
        inputs, outputs, bindings, stream = allocate_buffers(engine, context)
        inputs[0].host = input_data
        [output] = do_inference(context, engine, bindings, inputs, outputs, stream)
        
        # 验证结果
        predicted = np.argmax(output)
        results.append({
            "input_shape": input_shape,
            "expected": expected,
            "predicted": predicted,
            "success": predicted == expected
        })
        
        free_buffers(inputs, outputs, stream)
    
    return results

性能调优与最佳实践

动态形状下的INT8量化

在动态形状场景下使用INT8量化需要特别注意校准过程，如demo/BERT/builder.py所示：

# dynamic shape not working with calibration, so we need generate a calibration cache first using fulldims network
def generate_calibration_cache(sequence_lengths, workspace_size, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num):
    if not config.use_int8 or os.path.exists(calibrationCacheFile):
        return calibrationCacheFile
    
    # 生成校准缓存
    saved_use_fp16 = config.use_fp16
    config.use_fp16 = False
    config.is_calib_mode = True
    
    with build_engine([1], workspace_size, sequence_lengths, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, False) as serialized_engine:
        TRT_LOGGER.log(TRT_LOGGER.INFO, "校准缓存已生成: {:}".format(calibrationCacheFile))
    
    config.use_fp16 = saved_use_fp16
    config.is_calib_mode = False

多配置文件管理

对于复杂模型，建议为不同的动态形状范围创建独立的优化配置文件，如plugin/fcPlugin/CustomFCPluginDynamic_PluginConfig.yaml所示：

dynamic_shape_profile:
  - input_shape:
      min: [1, 1, 28, 28]
      opt: [1, 1, 28, 28]
      max: [1, 1, 56, 56]
    output_shape:
      min: [1, 1, 10]
      opt: [1, 1, 10]
      max: [1, 1, 10]

常见问题与解决方案

动态形状与INT8校准冲突

问题描述：在使用INT8量化时，动态形状可能导致校准失败。

解决方案：如demo/BERT/builder.py所示，先使用固定形状生成校准缓存，再用于动态形状引擎构建：

# dynamic shape not working with calibration, so we need generate a calibration cache first using fulldims network

性能优化建议

1. 合理设置优化配置文件，避免过度宽泛的形状范围
2. 对常用形状进行单独优化
3. 使用Timing Cache加速引擎构建：

if (!mParams.timingCacheFile.empty()) {
    timingCache = samplesCommon::buildTimingCacheFromFile(
        sample::gLogger.getTRTLogger(), *preprocessorConfig, mParams.timingCacheFile, sample::gLogError);
}

总结与展望

本文详细介绍了TensorRT ONNX动态形状推理的实现方法和全自动化测试框架搭建，包括核心原理、代码实现、测试策略和性能优化。通过合理利用TensorRT的动态形状功能，我们可以构建更加灵活和通用的推理引擎，满足多样化的业务需求。

未来，随着深度学习模型的不断发展，动态形状支持将变得越来越重要。TensorRT在这一领域持续创新，如最新版本中引入的动态形状推理优化和内存管理改进，都为构建高效、灵活的推理系统提供了强大支持。

完整的示例代码和测试框架可参考：

• sampleDynamicReshape
• BERT动态形状实现
• Polygraphy动态形状工具

通过这些资源，你可以快速上手并应用TensorRT动态形状推理技术，为你的应用带来更高的性能和灵活性。

【免费下载链接】TensorRT NVIDIA® TensorRT™ 是一个用于在 NVIDIA GPU 上进行高性能深度学习推理的软件开发工具包（SDK）。此代码库包含了 TensorRT 的开源组件 项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT