YOLOv12模型转换与部署：从PyTorch到ONNX再到TensorRT全流程

如果你训练了一个YOLOv12模型，看着它在训练集上表现优异，但一到实际推理时速度就慢得让人着急，那你来对地方了。模型训练只是第一步，把它变成一个在生产环境中能“飞起来”的推理引擎，才是真正体现价值的地方。

今天，我们就来手把手走一遍这个完整的“加速”流程：从PyTorch的.pt文件出发，经过ONNX这个中间站进行标准化和初步测试，最后抵达终点站TensorRT，生成一个高度优化的推理引擎。整个过程就像给模型做一次深度“瘦身”和“健身”，目标是在星图GPU平台上榨干硬件的每一分性能，获得极致的推理速度。我会把每一步的代码、常见坑点以及解决方案都摊开来讲，让你能跟着做，做出结果。

1. 准备工作：理清思路与搭建环境

在开始转换之前，我们得先搞清楚这一路要经过哪几个关键站点，以及各自需要准备什么。

1.1 技术路径全景图

整个流程可以概括为三个阶段：

1. PyTorch -> ONNX：将动态图模型转换为静态的、标准化的中间表示格式。这一步主要是为了“标准化”和“跨平台”。
2. ONNX Runtime 验证：用ONNX Runtime（ORT）加载和运行ONNX模型。这一步至关重要，是为了验证转换是否正确，模型输出是否与PyTorch一致，相当于一次“冒烟测试”。
3. ONNX -> TensorRT：利用TensorRT对ONNX模型进行图优化、层融合、精度校准（如INT8），并编译生成针对特定GPU硬件高度优化的引擎文件（.engine）。这是性能飞跃的关键。

1.2 环境搭建与依赖安装

假设你已经在星图GPU平台上拥有了一个环境（比如预置了CUDA的容器），我们需要安装以下核心库。建议使用conda或venv创建独立的Python环境。

# 1. 基础深度学习框架 (请根据你的训练环境选择合适版本)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 2. ONNX 相关
pip install onnx onnxruntime-gpu  # 注意是 onnxruntime-gpu，以启用CUDA

# 3. TensorRT 相关
# 星图平台可能已预装TensorRT。若需自行安装，请从NVIDIA官网下载对应CUDA版本的Tar包。
# 这里以通过pip安装Python绑定的方式为例（可能需要对应版本的TensorRT已安装在系统路径）
pip install tensorrt
# 或者使用更通用的 pycuda + TensorRT 库路径引用的方式

# 4. 其他工具
pip install opencv-python pycocotools numpy
# 用于YOLO模型加载和可视化的工具包（根据你使用的YOLOv12代码库安装）
# 例如： pip install ultralytics  # 如果使用Ultralytics库

关键检查点：

• CUDA/cuDNN版本一致性：确保PyTorch、ONNX Runtime GPU版、TensorRT使用的CUDA和cuDNN版本兼容。不一致是大多数错误的根源。
• 获取模型：准备好你训练好的YOLOv12模型权重文件（yolov12.pt）以及对应的模型定义代码（model.py）。

2. 第一步：从PyTorch到ONNX

ONNX就像一个通用的翻译官，它把不同框架（PyTorch, TensorFlow等）的模型翻译成一种统一的格式。这一步的准确性直接决定了后续步骤能否成功。

2.1 编写转换脚本

创建一个脚本，比如 export_to_onnx.py。核心是使用 torch.onnx.export 函数。

import torch
import torch.nn as nn
from models.yolo import Model  # 导入你的YOLOv12模型定义
import argparse

def export_onnx(weights, imgsz, opset, simplify, dynamic):
    """
    将PyTorch模型导出为ONNX格式。
    
    参数:
        weights: PyTorch模型权重文件路径 (.pt)
        imgsz: 输入图像尺寸 (高度, 宽度)
        opset: ONNX算子集版本 (推荐13或以上)
        simplify: 是否使用onnx-simplifier简化模型
        dynamic: 是否导出动态尺寸 (batch, height, width维度动态)
    """
    # 加载模型
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    ckpt = torch.load(weights, map_location=device)
    model = ckpt['model'] if 'model' in ckpt else ckpt  # 适应不同的checkpoint格式
    model = model.float().eval().to(device)
    
    # 准备示例输入
    batch_size = 1
    dummy_input = torch.randn(batch_size, 3, imgsz[0], imgsz[1], device=device)
    
    # 定义输入输出的名字
    input_names = ['images']
    output_names = ['output0']  # 根据你的模型输出层名修改，YOLO通常输出一个Tensor
    
    # 动态轴设置 (如果需要)
    dynamic_axes = None
    if dynamic:
        dynamic_axes = {
            'images': {0: 'batch_size', 2: 'height', 3: 'width'},  # 动态batch, height, width
            'output0': {0: 'batch_size'}
        }
    
    # 导出模型
    onnx_path = weights.replace('.pt', '.onnx')
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        verbose=False,
        opset_version=opset,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        do_constant_folding=True  # 常量折叠优化
    )
    print(f'模型已导出至: {onnx_path}')
    
    # 可选：模型简化 (强烈推荐)
    if simplify:
        try:
            import onnx
            from onnxsim import simplify
            onnx_model = onnx.load(onnx_path)
            model_simp, check = simplify(onnx_model)
            assert check, "简化后的模型验证失败"
            simplified_path = onnx_path.replace('.onnx', '_simplified.onnx')
            onnx.save(model_simp, simplified_path)
            print(f'简化模型已保存至: {simplified_path}')
            onnx_path = simplified_path
        except ImportError:
            print("未安装 onnx-simplifier，跳过简化步骤。运行: pip install onnx-simplifier")
        except Exception as e:
            print(f'模型简化失败: {e}')
    
    # 验证ONNX模型格式
    import onnx
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)
    print("ONNX模型格式检查通过。")
    
    return onnx_path

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='yolov12.pt', help='PyTorch权重路径')
    parser.add_argument('--imgsz', nargs='+', type=int, default=[640, 640], help='输入尺寸 [height, width]')
    parser.add_argument('--opset', type=int, default=13, help='ONNX opset版本')
    parser.add_argument('--simplify', action='store_true', help='使用onnx-simplifier简化模型')
    parser.add_argument('--dynamic', action='store_true', help='导出动态batch/尺寸')
    args = parser.parse_args()
    
    export_onnx(**vars(args))

2.2 运行转换并排查常见错误

在终端运行：

python export_to_onnx.py --weights yolov12.pt --imgsz 640 640 --opset 13 --simplify

常见错误与解决方案：

• 错误：Unsupported: ONNX export of operator ...

  • 原因：模型中包含ONNX当前opset版本不支持的操作符。
  • 解决：
    1. 尝试升级 opset_version (例如从11升到13或更高)。
    2. 检查模型定义，看是否有自定义的、非常规的操作。可能需要修改模型代码或寻找替代实现。
    3. 使用 --dynamic 参数有时能绕过某些静态形状相关的限制。

• 错误：RuntimeError: Expected all tensors to be on the same device...

  • 原因：模型权重或输入数据不在同一个设备（CPU/GPU）上。
  • 解决：在导出脚本中，确保模型和示例输入 dummy_input 都移动到同一设备（如 .to(device)）。

• 错误：简化后模型输出不一致

  • 原因：onnx-simplifier 在某些复杂结构上可能过于激进。
  • 解决：暂时关闭 --simplify 选项，先获得一个正确的ONNX模型。或者，尝试不同版本的 onnx-simplifier。

3. 第二步：用ONNX Runtime验证转换正确性

拿到ONNX文件后，千万别急着往下走。先用ONNX Runtime跑一下，确保模型转换没出岔子，输出和PyTorch原模型基本一致。

3.1 编写验证脚本

创建 validate_onnx.py。

import onnxruntime as ort
import torch
import numpy as np
import cv2

def validate_onnx(pytorch_model, onnx_path, test_image_path, imgsz=(640, 640)):
    """
    比较PyTorch模型和ONNX模型在同一输入下的输出。
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 1. 准备输入数据
    # 读取并预处理图像 (保持与训练/推理时一致的预处理)
    img = cv2.imread(test_image_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, imgsz[::-1])  # (width, height)
    # 归一化、转换通道、添加batch维度
    input_tensor = img_resized.astype(np.float32) / 255.0
    input_tensor = input_tensor.transpose(2, 0, 1)  # HWC -> CHW
    input_tensor = np.expand_dims(input_tensor, axis=0)  # CHW -> NCHW
    input_numpy = input_tensor.astype(np.float32)
    
    # 转换为PyTorch Tensor
    input_torch = torch.from_numpy(input_numpy).to(device)
    
    # 2. PyTorch推理
    pytorch_model.eval()
    with torch.no_grad():
        torch_output = pytorch_model(input_torch)
    torch_output = torch_output.cpu().numpy()
    print(f"PyTorch输出形状: {torch_output.shape}")
    
    # 3. ONNX Runtime推理
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider']
    session = ort.InferenceSession(onnx_path, providers=providers)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    
    ort_output = session.run([output_name], {input_name: input_numpy})[0]
    print(f"ONNX Runtime输出形状: {ort_output.shape}")
    
    # 4. 比较结果
    # 使用绝对误差或相对误差
    abs_diff = np.abs(torch_output - ort_output)
    mean_abs_diff = np.mean(abs_diff)
    max_abs_diff = np.max(abs_diff)
    
    print(f"\n--- 输出一致性验证 ---")
    print(f"平均绝对误差: {mean_abs_diff:.6e}")
    print(f"最大绝对误差: {max_abs_diff:.6e}")
    
    # 设定一个可接受的误差阈值 (例如 1e-5)
    tolerance = 1e-5
    if mean_abs_diff < tolerance and max_abs_diff < tolerance * 10:
        print("✅ ONNX模型输出与PyTorch模型基本一致，转换成功。")
        return True
    else:
        print("⚠️  输出存在较大差异，请检查转换过程。")
        # 可以进一步打印前几个元素的差异
        print("PyTorch 前5个值:", torch_output.flatten()[:5])
        print("ONNX    前5个值:", ort_output.flatten()[:5])
        return False

if __name__ == '__main__':
    # 加载你的PyTorch模型 (需要模型定义)
    from models.yolo import Model
    pt_model = torch.load('yolov12.pt', map_location='cpu')
    pt_model = pt_model['model'] if 'model' in pt_model else pt_model
    pt_model.eval()
    
    onnx_path = 'yolov12_simplified.onnx'
    test_img = 'test.jpg'
    
    validate_onnx(pt_model, onnx_path, test_img)

3.2 理解验证结果

• 误差极小（如 < 1e-6）：完美，可以进入下一步。
• 误差稍大（如 1e-5 到 1e-3）：可能由不同框架的数值计算精度（PyTorch vs ONNX Runtime）或不同操作符实现导致。如果下游任务（如目标检测mAP）影响不大，可以接受。
• 误差巨大：说明转换失败。需要回到第2步，检查模型定义、导出参数（尤其是动态轴）、opset版本，或者尝试不简化模型。

4. 第三步：从ONNX到TensorRT引擎

这是性能提升的魔法环节。TensorRT会对计算图进行大量优化，比如层融合、内核自动调优、精度校准等。

4.1 使用trtexec工具（推荐给初学者）

trtexec 是TensorRT自带的一个命令行工具，非常适合快速测试和生成引擎。

# 基础命令：将ONNX转换为TensorRT引擎，并指定优化参数
trtexec --onnx=yolov12_simplified.onnx \
        --saveEngine=yolov12.engine \
        --workspace=4096 \          # 设置最大工作空间大小(MiB)
        --fp16 \                    # 启用FP16精度，大幅提升速度，精度损失通常可接受
        --verbose                   # 输出详细信息

# 更复杂的例子：动态形状 + INT8量化
trtexec --onnx=yolov12_dynamic.onnx \
        --saveEngine=yolov12_dynamic_fp16_int8.engine \
        --minShapes=images:1x3x320x320 \     # 最小输入形状
        --optShapes=images:1x3x640x640 \     # 最优输入形状（常用尺寸）
        --maxShapes=images:1x3x1280x1280 \   # 最大输入形状
        --workspace=4096 \
        --fp16 \
        --int8 \                              # 启用INT8量化，需要校准数据
        --calib=<校准缓存文件路径> \          # INT8校准缓存
        --verbose

关键参数解释：

• --fp16: 启用半精度浮点数推理，速度提升显著，是性价比最高的优化。
• --int8: 启用INT8整数精度推理，速度最快，但需要校准数据来减少精度损失。
• --workspace: GPU显存工作区大小。复杂的优化可能需要更多空间，但受限于你的GPU显存。
• --minShapes/optShapes/maxShapes: 当你的ONNX模型是动态形状时，必须指定这些来让TensorRT为不同尺寸生成优化内核。

4.2 使用Python API进行精细控制

对于需要集成到Python代码流或进行更复杂操作（如自定义插件、逐层分析）的情况，可以使用TensorRT的Python API。

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

def build_engine_from_onnx(onnx_file_path, engine_file_path, fp16_mode=True, int8_mode=False, calibration_data=None):
    """
    使用TensorRT Python API构建引擎。
    """
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)
    
    # 解析ONNX模型
    with open(onnx_file_path, 'rb') as model:
        if not parser.parse(model.read()):
            print('ERROR: 解析ONNX文件失败')
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None
    
    # 构建配置
    config = builder.create_builder_config()
    config.max_workspace_size = 4 << 30  # 4GB
    if fp16_mode:
        config.set_flag(trt.BuilderFlag.FP16)
    if int8_mode:
        config.set_flag(trt.BuilderFlag.INT8)
        # 需要设置校准器 (此处省略，需要实现calibrator类)
        # config.int8_calibrator = MyCalibrator(calibration_data)
    
    # 设置动态形状profile (如果需要)
    profile = builder.create_optimization_profile()
    # 假设输入名为'images', 格式为NCHW
    profile.set_shape('images', min=(1, 3, 320, 320), opt=(1, 3, 640, 640), max=(1, 3, 1280, 1280))
    config.add_optimization_profile(profile)
    
    # 构建引擎
    print('开始构建TensorRT引擎，这可能需要几分钟...')
    serialized_engine = builder.build_serialized_network(network, config)
    
    if serialized_engine is None:
        print('构建引擎失败')
        return None
    
    # 保存引擎文件
    with open(engine_file_path, 'wb') as f:
        f.write(serialized_engine)
    print(f'引擎已保存至: {engine_file_path}')
    return serialized_engine

# 使用函数
build_engine_from_onnx('yolov12_simplified.onnx', 'yolov12_fp16.engine', fp16_mode=True)

4.3 常见TensorRT转换错误与优化

• 错误：[TensorRT] ERROR: ... could not be found in ...

  • 原因：ONNX模型中包含TensorRT不支持的操作符。
  • 解决：
    1. 确保ONNX模型是简化过的，且opset版本合适（通常11-13）。
    2. 查看错误信息中缺失的操作符，可能需要将其实现为TensorRT插件（Plugin）。对于YOLO，常见的如SiLU激活函数、Upsample（Resize）等，新版本TensorRT通常已原生支持。
    3. 可以尝试在导出ONNX时，将某些复杂操作分解为更基础的操作。

• 性能优化建议：

  • 优先启用FP16：在绝大多数GPU上都能带来1.5-3倍的速度提升，精度损失微乎其微。
  • 谨慎使用INT8：需要准备有代表性的校准数据集，并进行校准。能带来最大速度提升，但可能引入精度下降，需要仔细评估。
  • 调整工作空间：在GPU显存允许的情况下，适当增加--workspace大小，让TensorRT有更多空间尝试更激进的优化策略。
  • 利用动态形状：如果你的应用需要处理多种尺寸的输入，务必导出动态ONNX并正确设置min/opt/maxShapes。optShapes应该设置为最常出现的尺寸。

5. 第四步：在星图平台部署与性能测试

生成.engine文件后，就可以在星图GPU平台上部署了。

5.1 编写TensorRT推理脚本

创建一个 infer_tensorrt.py 脚本。

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
import time

class YOLOv12TRTInfer:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(self.logger) as runtime:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        
        # 分配输入输出缓冲区
        self.bindings = []
        self.inputs = []
        self.outputs = []
        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding))
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            # 分配主机和设备内存
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))
            if self.engine.binding_is_input(binding):
                self.inputs.append({'host': host_mem, 'device': device_mem})
            else:
                self.outputs.append({'host': host_mem, 'device': device_mem})
        self.stream = cuda.Stream()
        
    def preprocess(self, image_path, input_shape=(640, 640)):
        """预处理图像，与训练/导出时保持一致"""
        img = cv2.imread(image_path)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_resized = cv2.resize(img_rgb, input_shape[::-1])
        img_normalized = img_resized.astype(np.float32) / 255.0
        img_chw = img_normalized.transpose(2, 0, 1)
        img_batched = np.expand_dims(img_chw, axis=0)
        return img_batched.astype(np.float32), img.shape[:2]  # 返回原始尺寸用于后处理
    
    def infer(self, input_numpy):
        """执行推理"""
        # 将输入数据复制到设备
        np.copyto(self.inputs[0]['host'], input_numpy.ravel())
        cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
        
        # 执行推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        
        # 将输出数据复制回主机
        cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream)
        self.stream.synchronize()
        
        # 重塑输出 (根据你的模型输出结构调整)
        output = self.outputs[0]['host']
        # 例如，YOLOv12输出可能是 [1, 25200, 85]
        output_shape = self.context.get_binding_shape(1)  # 获取输出绑定形状
        output = output.reshape(output_shape)
        return output
    
    def postprocess(self, output, orig_shape, input_shape=(640, 640), conf_thresh=0.25):
        """后处理：将模型输出转换为检测框 (这里需要根据你的YOLO版本实现)"""
        # 这是一个简化的示例，实际需要实现非极大值抑制(NMS)等
        # 假设output是 [batch, num_anchors, 5+num_classes]
        detections = output[0]  # 取batch中的第一个
        # 1. 根据置信度阈值过滤
        scores = detections[..., 4:5] * detections[..., 5:]  # 对象置信度 * 类别置信度
        max_scores = np.max(scores, axis=-1)
        keep = max_scores > conf_thresh
        filtered_dets = detections[keep]
        
        # 2. 将框的坐标从输入尺寸映射回原始图像尺寸
        # ... (实现坐标转换)
        
        # 3. 执行NMS
        # ... (实现或调用cv2.dnn.NMSBoxes)
        
        return boxes, scores, class_ids  # 返回格式化的结果
    
    def benchmark(self, image_path, warmup=10, repeats=100):
        """性能基准测试"""
        input_data, _ = self.preprocess(image_path)
        # 预热
        for _ in range(warmup):
            _ = self.infer(input_data)
        
        # 计时
        times = []
        for _ in range(repeats):
            start = time.perf_counter()
            _ = self.infer(input_data)
            cuda.Context.synchronize()  # 确保CUDA操作完成
            end = time.perf_counter()
            times.append((end - start) * 1000)  # 毫秒
        
        avg_time = np.mean(times)
        fps = 1000 / avg_time
        print(f"平均推理时间: {avg_time:.2f} ms")
        print(f"预估FPS: {fps:.2f}")
        return avg_time, fps

if __name__ == '__main__':
    engine_path = 'yolov12_fp16.engine'
    inferer = YOLOv12TRTInfer(engine_path)
    
    # 单张图片推理
    test_img = 'test.jpg'
    input_data, orig_shape = inferer.preprocess(test_img)
    output = inferer.infer(input_data)
    boxes, scores, class_ids = inferer.postprocess(output, orig_shape)
    print(f"检测到 {len(boxes)} 个目标")
    
    # 性能测试
    print("\n--- 性能基准测试 ---")
    avg_time, fps = inferer.benchmark(test_img)

5.2 在星图平台运行与对比

1. 环境准备：确保你的星图实例有TensorRT库和对应的CUDA驱动。
2. 上传文件：将生成的 .engine 文件、推理脚本和测试图片上传到平台。
3. 运行测试：执行 python infer_tensorrt.py。
4. 性能对比：你可以用类似的方式，编写一个使用原始PyTorch模型（.pt）或ONNX Runtime（.onnx）的推理脚本，在相同输入和硬件下进行速度对比。

预期结果：在星图平台的GPU上，TensorRT（FP16）引擎的推理速度通常会比原始PyTorch模型快2倍到5倍甚至更多，具体取决于模型复杂度和优化选项。

6. 总结与后续建议

走完这一整套流程，你应该已经成功地将一个PyTorch训练的YOLOv12模型，转换成了在星图GPU上高速运行的TensorRT引擎。回顾一下，最关键的三步是正确导出ONNX、严谨验证结果以及合理配置TensorRT优化参数。其中，导出ONNX时的动态形状设置、使用ONNX Runtime做验证、以及TensorRT中启用FP16，是影响成功率和最终性能的几个关键点。

实际部署时，你可能会遇到比教程中更复杂的情况，比如需要处理自定义算子、进行INT8量化校准、或者部署到不同架构的GPU（如Jetson边缘设备）。这时，就需要更深入地查阅TensorRT的官方文档和社区资源。不过，只要掌握了这个基础流程，你就有了解决更复杂问题的抓手。

最后，记得模型部署是一个迭代过程。当你的模型结构改变或者TensorRT版本更新时，可能都需要重新走一遍这个流程。把这个过程脚本化、自动化，能为你节省大量时间。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。