DAMO-YOLO模型在边缘计算中的应用：Jetson平台部署指南-智慧文博士

DAMO-YOLO模型在边缘计算中的应用：Jetson平台部署指南

最近和不少做智能硬件的朋友聊天，大家普遍有个头疼的问题：想把最新的目标检测模型塞进摄像头、无人机或者工控机里，但模型太大、算力不够，跑起来不是卡顿就是发热严重。这就像想用一台小排量家用车去拉重货，不是拉不动，就是跑得慢还费油。

我前阵子正好在NVIDIA Jetson Nano和Jetson Orin NX上折腾了一下DAMO-YOLO，这个模型在精度和速度的平衡上做得挺有意思，特别适合边缘端。今天就把从环境准备到实际跑起来的完整过程，结合我们遇到的那些“坑”，跟大家详细聊聊。如果你手头有Jetson设备，或者正在为边缘设备选型发愁，这篇文章应该能给你一些直接的参考。

1. 为什么选择DAMO-YOLO上边缘？

在边缘设备上跑目标检测，我们最关心三件事：准不准、快不快、省不省。DAMO-YOLO在这几个方面，相比一些大家熟悉的模型，有它独特的优势。

首先说说“准”。DAMO-YOLO在公开数据集上的表现，比如COCO，精度（mAP）是相当能打的。它用了不少新的设计思路，比如更高效的网络结构和对小目标更友好的检测头。这意味着在监控场景里，远处的人脸或者车牌，它识别出来的可能性更高。

其次是“快”。模型速度不光看参数量，还得看实际在硬件上的推理延迟。DAMO-YOLO系列提供了从轻量级到高性能的不同版本（比如Tiny、S、M、L）。像Tiny版本，模型本身很小，在Jetson Nano这种入门级设备上，也能跑到比较理想的帧率，满足实时性要求不高的场景。

最关键的是“省”。这里的省，一是省内存，二是省算力。边缘设备的内存通常很有限，动辄几百兆的大模型根本加载不进去。DAMO-YOLO通过模型压缩和剪枝，体积控制得不错。同时，它的计算操作对Jetson平台的GPU加速比较友好，能充分利用Tensor Core，让每一分算力都花在刀刃上，从而降低功耗和发热。

简单来说，如果你需要一个在资源紧张的设备上，既能保持不错精度，又能流畅运行的检测模型，DAMO-YOLO是个值得优先考虑的选择。

2. 部署前的环境准备与检查

在开始动手之前，得先把“地基”打好。Jetson设备的系统环境和我们常用的x86服务器不太一样，需要一些特别的配置。

2.1 Jetson系统基础配置

拿到一台新的Jetson设备，无论是Nano、Xavier NX还是Orin系列，第一步都是确保系统是最新状态。打开终端，运行下面这几条命令：

# 更新软件源列表 sudo apt-get update # 升级所有已安装的包 sudo apt-get upgrade -y # 如果有系统映像更新，也一并升级 sudo apt-get dist-upgrade -y

这个过程可能会花点时间，但能避免很多因为库版本过旧导致的奇怪错误。更新完成后，建议重启一下设备。

接下来是安装一些我们后续肯定会用到的开发工具和库：

# 安装编译工具、Python开发环境及常用工具 sudo apt-get install -y build-essential cmake git libopencv-dev python3-dev python3-pip python3-venv

这里重点提一下OpenCV。Jetson系统通常预装了OpenCV，但可能版本较旧或者缺少某些功能。libopencv-dev这个包会确保我们拥有开发所需的头文件和库。如果后续需要从源码编译特定版本的OpenCV，可以再单独处理。

2.2 深度学习环境搭建

Jetson平台的核心是它的GPU，所以深度学习框架必须能够支持它的CUDA和TensorRT。NVIDIA为Jetson提供了优化好的PyTorch和TensorFlow轮子（wheel），直接安装这些预编译版本最省事。

首先，访问NVIDIA的官方开发者网站，找到对应你JetPack版本（可以用cat /etc/nv_tegra_release命令查看）的PyTorch下载链接。假设你的JetPack版本是5.1.2，安装命令可能类似这样：

# 安装PyTorch (具体URL请根据实际JetPack版本替换) wget https://developer.download.nvidia.com/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl pip3 install torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl # 安装TorchVision pip3 install torchvision

重要提示：务必使用NVIDIA官方为对应JetPack版本提供的PyTorch wheel。自己从源码编译不仅耗时数小时，还极易出错。

然后是TensorRT，这是NVIDIA用于高性能推理的SDK。它通常已经随JetPack系统安装好了。你可以用以下命令验证：

dpkg -l | grep tensorrt

如果显示已安装，通常就不需要再动它。TensorRT的Python接口包可以通过pip安装：

pip3 install nvidia-tensorrt

环境准备好后，创建一个独立的工作目录，并把DAMO-YOLO的代码拉取下来：

mkdir ~/damo-yolo-workspace && cd ~/damo-yolo-workspace git clone https://github.com/tinyvision/DAMO-YOLO.git cd DAMO-YOLO

3. 模型获取与转换实战

直接从PyTorch训练出来的模型（.pt文件）不能直接在TensorRT上跑，需要经过一个“翻译”过程，也就是模型转换。

3.1 下载与验证预训练模型

DAMO-YOLO官方提供了在COCO数据集上预训练好的各种尺寸模型。我们以轻量级的damo-yolo-tiny为例：

# 假设你在DAMO-YOLO项目根目录下 import torch # 模型名称，可根据需要替换为 'damo-yolo-s', 'damo-yolo-m' 等 model_name = 'damo-yolo-tiny' model = torch.hub.load('DAMO-YOLO', model_name, pretrained=True) # 保存为PyTorch的jit trace格式，方便后续转换 example_input = torch.randn(1, 3, 640, 640) # 标准输入尺寸 traced_model = torch.jit.trace(model, example_input) traced_model.save(f"{model_name}_traced.pt") print(f"模型已保存为 {model_name}_traced.pt")

运行这段代码，它会自动下载预训练权重并保存为可追踪的格式。记得检查一下生成的文件大小，确保下载完整。

3.2 关键一步：转换为ONNX格式

ONNX是一种开放的模型格式，是PyTorch模型通往TensorRT的“桥梁”。转换时需要注意几点：

固定输入尺寸：TensorRT在构建引擎时，固定尺寸能获得更好的优化。所以我们指定一个常用的尺寸，比如640x640。
操作集版本：使用支持的opset，比如12或13。
简化网络：使用onnx-simplifier工具可以优化ONNX模型结构，去掉一些冗余操作，这对后续TensorRT转换成功至关重要。

转换脚本如下：

import torch import onnx # 加载刚才保存的追踪模型 model = torch.jit.load("damo-yolo-tiny_traced.pt") model.eval() # 定义输入尺寸 input_shape = (1, 3, 640, 640) dummy_input = torch.randn(input_shape) # 导出为ONNX onnx_path = "damo-yolo-tiny.onnx" torch.onnx.export( model, dummy_input, onnx_path, input_names=["images"], output_names=["output"], opset_version=12, # 选择一个合适的opset版本 dynamic_axes=None # 使用静态尺寸，简化转换 ) print(f"ONNX模型已导出: {onnx_path}") # 验证ONNX模型 onnx_model = onnx.load(onnx_path) onnx.checker.check_model(onnx_model) print("ONNX模型检查通过。")

导出后，强烈建议使用onnx-simplifier进行简化：

pip3 install onnx-simplifier python3 -m onnxsim damo-yolo-tiny.onnx damo-yolo-tiny_sim.onnx

简化后的damo-yolo-tiny_sim.onnx文件，就是我们要交给TensorRT的“原料”。

4. 使用TensorRT优化与部署

这是核心环节，目标是把ONNX模型变成在Jetson上跑得飞快的TensorRT引擎。

4.1 构建TensorRT引擎

我们可以使用TensorRT的Python API来构建引擎。下面是一个简化的示例，演示了关键步骤：

import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) def build_engine(onnx_file_path, engine_file_path): builder = trt.Builder(TRT_LOGGER) network = builder.create_network(EXPLICIT_BATCH) parser = trt.OnnxParser(network, TRT_LOGGER) 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() # 对于Jetson，可以设置工作空间大小，例如1GB config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB # 启用FP16精度，可以显著提升速度（Jetson支持良好） if builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) print("已启用FP16精度。") print('正在构建引擎，这可能需要几分钟...') serialized_engine = builder.build_serialized_network(network, config) with open(engine_file_path, 'wb') as f: f.write(serialized_engine) print(f'引擎已保存至: {engine_file_path}') return serialized_engine # 使用简化后的ONNX模型构建引擎 onnx_path = "damo-yolo-tiny_sim.onnx" engine_path = "damo-yolo-tiny.engine" build_engine(onnx_path, engine_path)

构建引擎的过程会比较慢，可能需要几分钟，因为TensorRT在尝试各种内核实现来寻找最快的那个。构建一次之后，保存下来的.engine文件就可以反复加载使用，非常方便。

4.2 编写推理脚本

引擎建好了，接下来就是写个脚本让它干活。推理过程主要包括：加载引擎、准备输入数据、执行推理、解析输出。

import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np import cv2 import time class DAMOYOLO_TRT: def __init__(self, engine_path): self.TRT_LOGGER = trt.Logger(trt.Logger.WARNING) self.engine = self.load_engine(engine_path) self.context = self.engine.create_execution_context() self.inputs, self.outputs, self.bindings, self.stream = self.allocate_buffers() def load_engine(self, engine_path): with open(engine_path, 'rb') as f, trt.Runtime(self.TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read()) def allocate_buffers(self): inputs, outputs, bindings = [], [], [] stream = cuda.Stream() 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) bindings.append(int(device_mem)) if self.engine.binding_is_input(binding): inputs.append({'host': host_mem, 'device': device_mem}) else: outputs.append({'host': host_mem, 'device': device_mem}) return inputs, outputs, bindings, stream def infer(self, input_image): # 预处理图像：调整大小、归一化、转换通道顺序 (HWC -> CHW) preprocessed = self.preprocess(input_image) np.copyto(self.inputs[0]['host'], preprocessed.ravel()) # 将数据从CPU内存拷贝到GPU内存 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) # 将结果从GPU内存拷贝回CPU内存 cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream) self.stream.synchronize() # 后处理：将输出数据解析为检测框、置信度和类别 detections = self.postprocess(self.outputs[0]['host']) return detections def preprocess(self, img, input_shape=(640, 640)): # 简化的预处理：保持长宽比resize，填充到正方形，归一化 # 实际应用可能需要更精细的处理以匹配训练时的预处理 img_resized = cv2.resize(img, input_shape) img_normalized = img_resized.astype(np.float32) / 255.0 # CHW转换 img_chw = np.transpose(img_normalized, (2, 0, 1)) return np.ascontiguousarray(img_chw) def postprocess(self, output, conf_threshold=0.5): # 这里需要根据DAMO-YOLO具体的输出格式来解析 # 假设output是平铺的一维数组，需要reshape成 [num_boxes, 6] (x1, y1, x2, y2, conf, cls) # 实际解析逻辑需参考模型定义 print("输出数据形状（原始）:", output.shape) # 此处应实现非极大值抑制(NMS)等后处理步骤 # 为示例，直接返回原始输出（实际使用时必须替换） return output # 使用示例 if __name__ == "__main__": detector = DAMOYOLO_TRT("damo-yolo-tiny.engine") test_img = cv2.imread("test.jpg") if test_img is not None: start = time.time() results = detector.infer(test_img) end = time.time() print(f"推理耗时: {(end-start)*1000:.2f} ms") # 处理results，并在图像上画框... else: print("未找到测试图片 test.jpg")

这个类封装了TensorRT推理的核心流程。你需要根据DAMO-YOLO模型实际的输出结构，完善postprocess函数，实现框的解码、置信度过滤和非极大值抑制。

5. 性能调优与实际问题解决

模型跑起来只是第一步，让它跑得又好又快，还需要一些调优技巧。

1. 精度与速度的权衡（FP16 vs FP32）在构建引擎时，我们启用了FP16。这能带来显著的速度提升，但可能会带来微小的精度损失。对于大部分检测任务，FP16的精度完全够用。如果对精度要求极端苛刻，可以只用FP32，或者尝试TensorRT的FP16+精度校准功能。

2. 批处理大小（Batch Size）上面的例子批处理大小是1。如果可以一次性处理多张图片（比如从视频流中缓存几帧），增大批处理大小能更充分地利用GPU，提高吞吐量。在builder.create_network时，可以定义动态批次维度，或者在预处理时堆叠多张图片。

3. Jetson设备本身的优化

运行模式：Jetson设备有几种功耗模式。对于持续推理任务，可以设置为最大性能模式（MAXN）。
```
sudo nvpmodel -m 0 # 对于Jetson Nano，0通常是最大性能模式 sudo jetson_clocks # 锁定时钟频率
```
散热：长时间高负载运行，尤其是Jetson Nano，务必保证良好散热，过热会导致降频，性能骤降。
内存交换：如果遇到内存不足的错误，可以适当增加交换空间（swap），但这会影响速度，是下策。最好还是选用内存更大的设备或更小的模型。

4. 常见错误与解决