告别命令行：用Python脚本封装trtexec，实现ONNX模型批量自动转换Engine文件

在AI模型部署的日常工作中，我们经常需要将训练好的ONNX模型转换为TensorRT的engine文件。这个过程看似简单，但当面对以下场景时，手动操作就显得力不从心：

• 需要批量转换数十个不同结构的ONNX模型
• 同一模型需要测试FP16/INT8等不同精度模式
• 要对比不同workspace大小对转换结果的影响
• 需要记录每次转换的详细日志和性能数据

本文将介绍如何用Python编写一个自动化脚本，将这些重复性工作全部交给程序处理。这个方案特别适合需要管理多个模型版本或频繁测试不同参数的AI工程师。

1. 环境准备与基础工具链

1.1 TensorRT工具链配置

确保系统中已安装以下组件并正确配置环境变量：

bash复制# 验证CUDA安装
nvcc --version

# 验证TensorRT安装
dpkg -l | grep tensorrt  # Ubuntu
# 或检查Windows环境变量中的TensorRT路径

注意：不同版本的TensorRT对CUDA和cuDNN有特定要求，建议参考NVIDIA官方文档匹配版本。

1.2 trtexec基础使用

trtexec是TensorRT自带的命令行工具，基本转换命令如下：

bash复制trtexec --onnx=model.onnx --saveEngine=model.trt --fp16 --workspace=4096

常用参数说明：

2. Python自动化脚本设计

2.1 脚本架构设计

我们设计一个具备以下功能的Python脚本：

1. 模型遍历：自动扫描指定目录下的ONNX文件
2. 参数配置：支持不同精度模式和workspace设置
3. 子进程调用：安全执行trtexec命令
4. 错误处理：捕获并记录转换失败情况
5. 日志生成：保存详细的转换报告

python复制import subprocess
import glob
import os
from datetime import datetime

class TRTConverter:
    def __init__(self, input_dir, output_dir):
        self.input_dir = input_dir
        self.output_dir = output_dir
        self.log_file = f"conversion_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        
    def find_onnx_files(self):
        return glob.glob(os.path.join(self.input_dir, "*.onnx"))
    
    def convert_single_model(self, onnx_path, precision="fp32", workspace=1024):
        # 实现单个模型转换逻辑
        pass
    
    def batch_convert(self, precision_options=None, workspace_options=None):
        # 实现批量转换逻辑
        pass

2.2 核心功能实现

模型转换函数实现：

python复制def convert_single_model(self, onnx_path, precision="fp32", workspace=1024):
    model_name = os.path.splitext(os.path.basename(onnx_path))[0]
    output_path = os.path.join(self.output_dir, f"{model_name}_{precision}.trt")
    
    cmd = [
        "trtexec",
        f"--onnx={onnx_path}",
        f"--saveEngine={output_path}",
        f"--workspace={workspace}",
        "--explicitBatch"
    ]
    
    if precision == "fp16":
        cmd.append("--fp16")
    elif precision == "int8":
        cmd.append("--int8")
    
    try:
        result = subprocess.run(
            cmd,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        self._log_conversion(model_name, True, result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        self._log_conversion(model_name, False, e.stderr)
        return False

批量转换与参数组合：

python复制def batch_convert(self, precision_options=None, workspace_options=None):
    if precision_options is None:
        precision_options = ["fp32"]
    if workspace_options is None:
        workspace_options = [1024]
    
    onnx_files = self.find_onnx_files()
    total = len(onnx_files) * len(precision_options) * len(workspace_options)
    processed = 0
    
    for onnx_path in onnx_files:
        for precision in precision_options:
            for workspace in workspace_options:
                success = self.convert_single_model(
                    onnx_path, precision, workspace
                )
                processed += 1
                print(f"Progress: {processed}/{total} - {'Success' if success else 'Failed'}")

3. 高级功能扩展

3.1 性能基准测试

在转换的同时收集性能数据：

python复制def add_benchmark_params(self, cmd, iterations=100):
    cmd.extend([
        f"--iterations={iterations}",
        "--duration=0",
        "--avgRuns=100"
    ])
    return cmd

3.2 错误处理与恢复

增强鲁棒性的错误处理机制：

python复制def handle_cublas_error(self, cmd):
    # 处理常见的CUBLAS错误
    cmd.append("--tacticSources=-cublasLt,+cublas")
    return cmd

def convert_with_retry(self, cmd, max_retries=3):
    for attempt in range(max_retries):
        try:
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            return True, result.stdout
        except subprocess.CalledProcessError as e:
            if "CUBLAS_STATUS_SUCCESS" in e.stderr:
                cmd = self.handle_cublas_error(cmd)
                continue
            return False, e.stderr
    return False, "Max retries exceeded"

3.3 日志系统设计

详细的日志记录可以帮助后续分析：

python复制def _log_conversion(self, model_name, success, details):
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    status = "SUCCESS" if success else "FAILED"
    
    log_entry = f"""
[{timestamp}] {status} - {model_name}
{'='*50}
{details}
{'='*50}
"""
    
    with open(self.log_file, "a") as f:
        f.write(log_entry)

4. 实战应用案例

4.1 典型使用场景

python复制if __name__ == "__main__":
    # 初始化转换器
    converter = TRTConverter(
        input_dir="models/onnx",
        output_dir="models/trt"
    )
    
    # 执行批量转换
    converter.batch_convert(
        precision_options=["fp32", "fp16", "int8"],
        workspace_options=[1024, 2048, 4096]
    )

4.2 性能对比分析

通过脚本自动生成的日志，我们可以轻松对比不同参数下的转换结果：

4.3 集成到CI/CD流程

将转换脚本集成到自动化部署流程中：

python复制# Jenkins或GitHub Actions示例
def ci_cd_pipeline():
    # 1. 拉取最新模型
    download_latest_models()
    
    # 2. 执行批量转换
    converter = TRTConverter("ci/models", "ci/engines")
    success = converter.batch_convert()
    
    # 3. 验证转换结果
    if not success:
        send_alert("模型转换失败")
        raise Exception("Conversion failed")
    
    # 4. 部署新模型
    deploy_engines()

5. 常见问题解决方案

在实际使用中可能会遇到以下典型问题：

1. 版本兼容性问题

   • TensorRT与CUDA/cuDNN版本必须严格匹配
   • ONNX opset版本需符合TensorRT支持范围

2. 显存不足错误

   • 适当增大--workspace参数值
   • 尝试更小的batch size或更低的精度模式

3. 不支持的ONNX算子

   • 使用polygraphy工具诊断问题算子
   • 考虑自定义插件或修改模型架构

提示：遇到问题时，先尝试使用--verbose参数获取详细日志，这能帮助快速定位问题根源。

6. 工程化建议

为了使这个自动化方案更加健壮，建议：

• 参数配置文件化：使用YAML或JSON管理转换参数
• 异常通知机制：集成邮件或Slack通知
• 版本控制：记录模型与engine文件的对应关系
• 资源监控：跟踪GPU显存和利用率

python复制# 示例：YAML配置文件
conversion_profiles:
  default:
    precision: [fp32, fp16]
    workspace: [1024, 2048]
  performance:
    precision: [int8]
    workspace: [4096, 8192]
    benchmark: true
    iterations: 1000

这个Python封装方案彻底改变了我们团队处理模型转换的方式。从原来每人每天最多处理几个模型，到现在可以轻松完成上百个模型的自动化测试和验证。特别是在模型迭代频繁的项目中，这种自动化工具的价值更加凸显。