FunAudioLLM/CosyVoice项目：动态输入ONNX模型转TensorRT引擎的解决方案

概述

在深度学习模型部署过程中，将ONNX模型转换为TensorRT引擎是常见的优化手段。然而，当模型包含动态输入时，转换过程会遇到特殊挑战。本文以FunAudioLLM/CosyVoice项目为例，探讨动态输入ONNX模型转换为TensorRT引擎的解决方案。

问题背景

TensorRT作为NVIDIA推出的高性能推理引擎，能够显著提升模型推理速度。但在处理动态输入模型时，需要特别注意优化配置。典型的错误提示为："Network has dynamic or shape inputs, but no optimization profile has been defined"。

动态输入处理要点

1. 优化配置文件：TensorRT需要明确的优化配置文件来处理动态输入
2. 输入维度范围：必须为每个动态维度指定最小、最优和最大尺寸
3. 工作空间配置：适当的工作空间大小对转换成功至关重要

解决方案实现

针对FunAudioLLM/CosyVoice项目的flow.onnx模型转换，完整的TensorRT转换脚本应包含以下关键部分：

import tensorrt as trt

def build_engine(onnx_file_path, engine_file_path):
    logger = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB工作空间
    
    # 创建网络定义
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    
    # 解析ONNX模型
    parser = trt.OnnxParser(network, logger)
    with open(onnx_file_path, "rb") as model:
        if not parser.parse(model.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None
    
    # 设置优化配置
    profile = builder.create_optimization_profile()
    input_tensor = network.get_input(0)
    input_shape = input_tensor.shape
    
    # 假设输入为[batch, sequence, features]格式
    # 设置动态维度范围
    min_shape = (1, 1, input_shape[2])  # 最小输入尺寸
    opt_shape = (1, 256, input_shape[2])  # 典型/最优输入尺寸
    max_shape = (1, 1024, input_shape[2])  # 最大输入尺寸
    
    profile.set_shape(input_tensor.name, min_shape, opt_shape, max_shape)
    config.add_optimization_profile(profile)
    
    # 构建引擎
    engine_bytes = builder.build_serialized_network(network, config)
    if engine_bytes is None:
        print("引擎构建失败")
        return None
    
    # 保存引擎文件
    with open(engine_file_path, "wb") as f:
        f.write(engine_bytes)
    return engine_bytes

关键参数说明

1. 优化配置文件(Optimization Profile)：定义模型处理不同输入尺寸的能力
2. 三种输入尺寸：
   • 最小尺寸(min_shape)：模型能处理的最小输入
   • 最优尺寸(opt_shape)：最常见的输入尺寸，TensorRT会针对此尺寸优化
   • 最大尺寸(max_shape)：模型能处理的最大输入
3. 工作空间(Workspace)：TensorRT优化过程中可用的临时内存

实际应用建议

1. 尺寸范围确定：根据实际应用场景确定合理的输入尺寸范围
2. 多优化配置：对于变化较大的输入，可考虑设置多个优化配置
3. 性能测试：转换后应在不同输入尺寸下测试推理性能
4. 版本兼容性：注意TensorRT与CUDA、cuDNN版本的兼容性

总结

处理动态输入ONNX模型转TensorRT引擎时，正确配置优化参数是关键。通过合理设置优化配置文件中的输入尺寸范围，可以确保模型在不同输入尺寸下都能高效运行。FunAudioLLM/CosyVoice项目的实践经验表明，这种方法能有效解决动态输入转换问题。