DAMOYOLO-S模型导出与部署全流程：从PyTorch到ONNX再到TensorRT加速

张

张建站

2026/7/9 10:11:06

10分钟阅读

DAMOYOLO-S模型导出与部署全流程从PyTorch到ONNX再到TensorRT加速如果你已经用DAMOYOLO-S训练出了一个不错的模型接下来最头疼的恐怕就是怎么把它用起来了。尤其是在需要实时处理大量图片或者视频流的场景里PyTorch原生的推理速度常常让人等得心焦。这时候模型部署和加速就成了关键。今天咱们就来手把手走一遍这个流程把训练好的PyTorch模型先变成通用的ONNX格式再用TensorRT这把“利器”进行深度优化和加速。整个过程走下来模型推理速度提升个几倍是很有希望的特别适合用在那些对延迟要求苛刻的生产环境里。别担心我会尽量用人话把每个步骤讲清楚。1. 准备工作与环境搭建在开始“折腾”模型之前得先把台子搭好。这一步主要是准备好模型文件和安装必要的工具链。首先你得有一个训练好的DAMOYOLO-S模型权重文件通常是一个.pth或.pt文件。假设你的文件叫damoyolo_s_model.pth。接下来是安装环境。我们主要需要三个东西PyTorch用来加载和导出模型、ONNX相关工具用来转换和优化模型、以及TensorRT用来做终极加速。# 1. 安装PyTorch (请根据你的CUDA版本选择合适命令这里以CUDA 11.8为例) pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 2. 安装ONNX和onnx-simplifier pip install onnx onnxsim onnxruntime-gpu # 3. 安装TensorRT # 这是最复杂的一步。强烈建议从NVIDIA官网下载对应你系统环境的TensorRT tar包进行安装。 # 例如下载 TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz # 解压后将lib路径加入环境变量并安装Python包 export LD_LIBRARY_PATH$LD_LIBRARY_PATH:/path/to/TensorRT-8.6.1.6/lib pip install /path/to/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp3x-none-linux_x86_64.whl # 4. 安装pycuda用于TensorRT的一些底层操作 pip install pycuda安装TensorRT可能会遇到一些依赖问题比如缺少libnvinfer.so。如果遇到请确保安装了对应的libnvinfer系统包或者直接将TensorRT库文件拷贝到系统库路径下。2. 从PyTorch到ONNX模型导出与简化ONNX就像一个“中间翻译”它能把不同框架比如PyTorch、TensorFlow训练出来的模型转换成一种通用的格式。我们先完成这一步。2.1 编写模型导出脚本你需要一个Python脚本来加载你的PyTorch模型并调用PyTorch的导出功能。关键点在于定义好模型的输入一个示例输入张量和输出。创建一个文件叫export_to_onnx.py内容如下import torch import onnx from models import build_model # 假设这是你的DAMOYOLO模型构建函数 from configs.damoyolo_s_config import get_config # 假设这是你的配置文件 def main(): # 加载模型配置 cfg get_config() cfg.model.backbone.pretrained False # 导出时不需要预训练权重 cfg.model.head.num_classes 80 # 改成你训练时的类别数例如COCO是80 # 构建模型并加载权重 model build_model(cfg) checkpoint torch.load(damoyolo_s_model.pth, map_locationcpu) model.load_state_dict(checkpoint[model]) model.eval() # 非常重要切换到评估模式 model.cuda() # 将模型放到GPU上 # 定义输入尺寸 (batch_size, channels, height, width) # DAMOYOLO-S的典型输入是640x640 dummy_input torch.randn(1, 3, 640, 640).cuda() # 设置输入输出名 input_names [images] output_names [output0, output1, output2] # 根据DAMOYOLO的实际输出层数修改 # 导出模型为ONNX格式 torch.onnx.export( model, dummy_input, damoyolo_s.onnx, verboseFalse, # 设为True可以看到导出详情 opset_version12, # ONNX算子集版本11或12比较稳定 do_constant_foldingTrue, # 优化常量 input_namesinput_names, output_namesoutput_names, dynamic_axes{ images: {0: batch_size}, # 支持动态batch output0: {0: batch_size}, output1: {0: batch_size}, output2: {0: batch_size} } ) print(ONNX model exported to damoyolo_s.onnx) if __name__ __main__: main()注意output_names和dynamic_axes需要根据DAMOYOLO-S模型的实际结构进行调整。最准确的方法是先运行一次模型打印出输出的元组看看有几个元素每个元素的形状是什么。例如输出可能是三个不同尺度的特征图。2.2 简化ONNX模型直接导出的ONNX模型可能包含一些冗余的算子不利于后续优化。我们可以用onnx-simplifier来简化它。python -m onnxsim damoyolo_s.onnx damoyolo_s_sim.onnx这个命令会尝试合并冗余的算子优化模型结构。生成的damoyolo_s_sim.onnx通常更小、更干净是进行TensorRT转换的理想输入。3. 使用TensorRT进行极致优化ONNX模型已经可以在很多推理引擎上跑了但要榨干NVIDIA GPU的性能还得看TensorRT。它会针对你的具体GPU型号比如RTX 4090, A100进行内核优化、层融合、精度校准生成一个高度优化的推理引擎。3.1 构建TensorRT引擎我们可以使用TensorRT的Python API来构建引擎。这里有两种精度模式可选FP32单精度精度高和FP16半精度速度快内存占用少精度略有损失。对于目标检测FP16通常是速度和精度的一个很好平衡。创建一个build_trt_engine.py文件import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit TRT_LOGGER trt.Logger(trt.Logger.WARNING) # 设置日志级别WARNING减少输出 def build_engine(onnx_file_path, engine_file_path, fp16_modeTrue): 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: Failed to parse the ONNX file.) for error in range(parser.num_errors): print(parser.get_error(error)) return None # 配置构建参数 config builder.create_builder_config() config.max_workspace_size 1 30 # 1GB工作空间可根据需要调整 if fp16_mode and builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) print(Using FP16 mode) else: print(Using FP32 mode) # 设置动态形状profile如果导出时指定了动态batch profile builder.create_optimization_profile() # 假设最小、最优、最大batch size分别为1, 8, 16 profile.set_shape(images, (1, 3, 640, 640), (8, 3, 640, 640), (16, 3, 640, 640)) config.add_optimization_profile(profile) # 构建引擎 print(Building engine... This may take a while.) serialized_engine builder.build_serialized_network(network, config) if serialized_engine is None: print(Failed to build engine) return None # 保存引擎到文件 with open(engine_file_path, wb) as f: f.write(serialized_engine) print(fEngine saved to {engine_file_path}) return serialized_engine if __name__ __main__: onnx_path damoyolo_s_sim.onnx engine_path_fp16 damoyolo_s_fp16.engine build_engine(onnx_path, engine_path_fp16, fp16_modeTrue)运行这个脚本你会得到一个.engine文件。这就是为你的GPU和模型量身定制的、优化后的推理引擎。3.2 使用TensorRT引擎进行推理引擎建好了怎么用呢我们需要写一个加载引擎并执行推理的脚本。这里涉及到GPU内存的分配和数据拷贝。创建一个inference_with_trt.py文件import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np import cv2 import time class TRTInference: 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: # 获取每个binding的shape注意是元组 shape self.engine.get_binding_shape(binding) size trt.volume(shape) dtype trt.nptype(self.engine.get_binding_dtype(binding)) # 在GPU上分配内存 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, shape: shape}) print(fInput binding: {binding}, shape: {shape}) else: self.outputs.append({host: host_mem, device: device_mem, shape: shape, name: binding}) print(fOutput binding: {binding}, shape: {shape}) self.stream cuda.Stream() def preprocess(self, image_path): # 简单的预处理读取、Resize、归一化、转置维度 img cv2.imread(image_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (640, 640)) img img.astype(np.float32) / 255.0 # 归一化到[0,1] img img.transpose(2, 0, 1) # HWC - CHW img np.expand_dims(img, axis0) # 添加batch维度 return img def infer(self, image_np): # 将输入数据拷贝到GPU np.copyto(self.inputs[0][host], image_np.ravel()) cuda.memcpy_htod_async(self.inputs[0][device], self.inputs[0][host], self.stream) # 设置动态输入形状如果用了动态batch # self.context.set_binding_shape(0, image_np.shape) # 执行推理 self.context.execute_async_v2(bindingsself.bindings, stream_handleself.stream.handle) # 将输出数据从GPU拷贝回CPU for out in self.outputs: cuda.memcpy_dtoh_async(out[host], out[device], self.stream) self.stream.synchronize() # 等待流中所有操作完成 # 整理输出 results [] for out in self.outputs: # 根据原始shape重塑输出数组 results.append(out[host].reshape(out[shape])) return results def postprocess(self, trt_outputs, conf_threshold0.5): # 这里需要根据DAMOYOLO的后处理逻辑来写 # trt_outputs是一个列表包含多个输出层的结果 # 你需要将这些结果解码成边界框、置信度和类别 # 由于后处理逻辑较长这里仅给出框架 boxes, scores, class_ids [], [], [] # ... 你的后处理代码 ... # 可能包括解码anchor、应用置信度阈值、非极大值抑制(NMS)等 return boxes, scores, class_ids if __name__ __main__: # 初始化推理器 trt_engine TRTInference(damoyolo_s_fp16.engine) # 预处理图片 img_np trt_engine.preprocess(test.jpg) # 预热 for _ in range(10): _ trt_engine.infer(img_np) # 正式测速 num_runs 100 start time.time() for _ in range(num_runs): outputs trt_engine.infer(img_np) end time.time() avg_latency (end - start) * 1000 / num_runs print(fAverage inference latency: {avg_latency:.2f} ms) # 后处理并可视化 boxes, scores, class_ids trt_engine.postprocess(outputs) print(fDetected {len(boxes)} objects.)这个类封装了TensorRT推理的核心步骤。postprocess函数需要你根据DAMOYOLO-S的后处理逻辑来填充这通常涉及将模型输出的特征图解码成具体的边界框。4. 性能对比与实用技巧走完整个流程你可能最关心的是到底快了多少这里有一些经验之谈。在我的测试环境RTX 3090, CUDA 11.8下对比纯PyTorch推理经过TensorRT FP16优化后的DAMOYOLO-S引擎速度通常能有2到5倍的提升具体取决于输入尺寸和batch size。batch size越大TensorRT的优化效果往往越明显因为它能更好地并行化计算。几个实用小技巧动态Shape如果你需要处理不同尺寸的图片在导出ONNX和构建TensorRT引擎时务必设置好动态形状范围dynamic_axes和optimization_profile。这会给引擎构建增加一些复杂度但能带来部署的灵活性。精度校准对于INT8量化能进一步提速并减少内存占用但需要准备一个校准数据集并且可能会有轻微的精度损失。对于大多数检测任务FP16已经是个很好的选择。多流推理在高并发场景下可以为每个CPU线程或处理流创建一个独立的TensorRT上下文IExecutionContext利用GPU的并发执行能力。引擎缓存.engine文件的构建比较耗时。在生产环境中应该预先构建好引擎并缓存起来服务启动时直接加载而不是每次运行时都构建。5. 总结从PyTorch到ONNX再到TensorRT这条路虽然步骤不少但每一步都有明确的目标。ONNX负责格式的统一和初步简化而TensorRT则承担了针对特定硬件深度优化的重任是提升NVIDIA GPU上模型推理性能的“标准答案”。实际操作下来最花时间的可能不是跑脚本而是根据自己模型的具体结构调整导出和后处理的代码。一旦打通看到推理速度大幅提升那种感觉还是很棒的。对于真的要部署到线上服务里的模型这点投入绝对值得。如果过程中遇到问题多看看TensorRT的日志输出里面通常会有比较详细的错误信息。另外确保你的CUDA、cuDNN、TensorRT版本之间是兼容的这也是减少莫名错误的关键。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。