23、边缘AI应用实战:YOLOv8部署、实时目标检测Pipeline、模型服务化(Triton)
说实话,到了这一章,才真正到了大家最兴奋的部分。前面我们折腾了那么多环境、驱动、库,说白了都是为了这一刻——让模型在Jetson上真正跑起来,而且跑得漂亮。
YOLOv8是目前最主流的实时目标检测模型之一。我最早接触YOLO系列还是v3版本,那时候在Jetson TX2上跑,帧率惨不忍睹。现在到了v8,配合TensorRT优化,在Orin上跑个几百帧都不是梦。今天我们就来完整走一遍:从模型导出、TensorRT加速、实时Pipeline搭建,再到用Triton Inference Server做模型服务化。
23.1 YOLOv8模型导出与TensorRT优化
模型部署的第一步,是把PyTorch模型转成TensorRT能吃的格式。YOLOv8官方已经提供了非常方便的导出工具,但我个人建议,别直接用默认参数,得针对Jetson平台做一些调优。
23.1.1 安装Ultralytics与导出ONNX
首先,确保你的Jetson上已经安装了Ultralytics库。如果还没装,用pip搞定:
pip install ultralytics
然后,下载一个预训练的YOLOv8模型,比如nano版本(适合边缘设备):
from ultralytics import YOLO
# 加载预训练模型
model = YOLO('yolov8n.pt')
# 导出为ONNX格式
model.export(format='onnx', imgsz=640, half=True)
这里有个细节:half=True表示导出FP16精度的模型。Jetson Orin的Tensor Core对FP16有硬件加速,推理速度能提升将近一倍。我在项目中遇到过,如果模型精度要求极高,可以保留FP32,但大多数边缘场景下FP16完全够用。
23.1.2 ONNX转TensorRT引擎
ONNX模型拿到手之后,我们需要用trtexec工具把它编译成TensorRT的引擎文件(.engine)。这个工具在JetPack SDK里已经自带了。
trtexec --onnx=yolov8n.onnx \
--saveEngine=yolov8n.engine \
--fp16 \
--workspace=4096 \
--minShapes=input:1x3x640x640 \
--optShapes=input:1x3x640x640 \
--maxShapes=input:1x3x640x640
嗯,这里要注意:--workspace参数设置的是显存工作空间,单位是MB。Jetson Orin的显存有限,我一般设成4096(4GB),如果你用的是Orin NX 16GB版本,可以适当调大。
23.2 实时目标检测Pipeline搭建
引擎文件有了,接下来就是搭建实时推理Pipeline。我习惯用Python写一个轻量级的推理类,把预处理、推理、后处理都封装起来。
23.2.1 推理引擎封装
import tensorrt as trt
import numpy as np
import cv2
class YOLOv8TRT:
def __init__(self, engine_path):
self.logger = trt.Logger(trt.Logger.WARNING)
self.runtime = trt.Runtime(self.logger)
with open(engine_path, 'rb') as f:
self.engine = self.runtime.deserialize_cuda_engine(f.read())
self.context = self.engine.create_execution_context()
self.inputs, self.outputs, self.bindings = [], [], []
self.stream = None
self._allocate_buffers()
def _allocate_buffers(self):
# 分配输入输出缓冲区
for i in range(self.engine.num_bindings):
name = self.engine.get_binding_name(i)
shape = self.engine.get_binding_shape(i)
size = trt.volume(shape)
dtype = trt.nptype(self.engine.get_binding_dtype(i))
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(i):
self.inputs.append({'name': name, 'host': host_mem, 'device': device_mem, 'shape': shape})
else:
self.outputs.append({'name': name, 'host': host_mem, 'device': device_mem, 'shape': shape})
def infer(self, image):
# 预处理
input_tensor = self._preprocess(image)
np.copyto(self.inputs[0]['host'], input_tensor.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)
cuda.synchronize()
# 后处理
output = self.outputs[0]['host'].reshape(self.outputs[0]['shape'])
return self._postprocess(output)
这段代码看起来有点长,但核心逻辑其实很简单:把图像数据拷贝到GPU,执行推理,再把结果拷回来。我封装的时候特意把预处理和后处理也放进了类里,这样调用方只需要传一张图片就行。
23.2.2 实时视频流处理
对于实时目标检测,我们通常需要处理摄像头或视频文件。下面是一个完整的Pipeline示例:
import cv2
def run_realtime_detection(engine_path, video_source=0):
model = YOLOv8TRT(engine_path)
cap = cv2.VideoCapture(video_source)
# 设置低延迟模式
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
while True:
ret, frame = cap.read()
if not ret:
break
# 推理
detections = model.infer(frame)
# 绘制检测框
for det in detections:
x1, y1, x2, y2, conf, cls_id = det
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
cv2.putText(frame, f'{cls_id}: {conf:.2f}', (int(x1), int(y1)-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
cv2.imshow('YOLOv8 Detection', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
你想想看,这个Pipeline跑在Jetson Orin上,YOLOv8n模型能做到200+FPS。我实际测试过,在Orin NX 16GB上,输入640x640分辨率,推理延迟只有3-4毫秒。嗯,这里要注意,如果使用USB摄像头,帧率可能会被摄像头本身限制,建议用CSI摄像头或者RTSP流。
CAP_PROP_BUFFERSIZE,结果摄像头缓冲了多帧,导致推理结果和显示画面不同步。另外,如果使用多线程处理,记得加锁保护共享资源。
23.3 模型服务化:Triton Inference Server
单机推理虽然爽,但实际项目中往往需要把模型能力暴露给其他服务。这时候就需要Triton Inference Server出马了。Triton是NVIDIA推出的高性能推理服务器,支持多种后端,包括TensorRT、ONNX Runtime、PyTorch等。
23.3.1 Triton的安装与配置
在Jetson上安装Triton,最简单的方式是用Docker。JetPack SDK里已经包含了Triton的容器镜像:
docker pull nvcr.io/nvidia/tritonserver:23.10-py3
然后,我们需要准备模型仓库。Triton要求模型按特定目录结构存放:
model_repository/
└── yolov8n/
├── 1/
│ └── model.plan # TensorRT引擎文件
└── config.pbtxt # 模型配置文件
config.pbtxt文件定义了模型的输入输出格式:
name: "yolov8n"
platform: "tensorrt_plan"
max_batch_size: 8
input [
{
name: "input"
data_type: TYPE_FP16
dims: [3, 640, 640]
}
]
output [
{
name: "output"
data_type: TYPE_FP16
dims: [1, 84, 8400]
}
]
这里有个坑:YOLOv8的输出维度是[1, 84, 8400],其中84表示4个坐标+80个类别概率,8400是检测框的数量。我第一次配置的时候写错了维度,结果Triton一直报错,折腾了半天才发现。
23.3.2 启动Triton并调用服务
模型仓库准备好之后,启动Triton容器:
docker run --gpus all -it --rm \
-v /path/to/model_repository:/models \
-p 8000:8000 -p 8001:8001 -p 8002:8002 \
nvcr.io/nvidia/tritonserver:23.10-py3 \
tritonserver --model-repository=/models
启动成功后,我们可以用Python客户端调用推理服务:
import tritonclient.http as httpclient
import numpy as np
client = httpclient.InferenceServerClient(url='localhost:8000')
# 准备输入数据
input_data = np.random.randn(1, 3, 640, 640).astype(np.float16)
input_tensor = httpclient.InferInput('input', input_data.shape, 'FP16')
input_tensor.set_data_from_numpy(input_data)
# 执行推理
output_tensor = httpclient.InferRequestedOutput('output')
result = client.infer('yolov8n', [input_tensor], [output_tensor])
# 获取结果
output = result.as_numpy('output')
print(f'Detection results shape: {output.shape}')
说白了,Triton把复杂的推理过程封装成了RESTful API或gRPC接口,其他服务只需要发HTTP请求就能用上模型能力。我在一个智慧安防项目里就是这么干的:摄像头采集的画面通过RTSP流送到Jetson,Jetson上跑Triton做推理,结果再推送到云端。整个Pipeline非常稳定,连续跑了几个月没出过问题。
23.4 性能调优与最佳实践
最后,分享几个我在实际项目中总结的性能调优技巧:
| 优化项 | 说明 | 效果 |
|---|---|---|
| FP16推理 | 使用半精度浮点数 | 速度提升1.5-2倍 |
| 批量推理 | 将多帧合并为一个batch | 吞吐量提升2-3倍 |
| 异步推理 | 使用CUDA Stream并行处理 | 延迟降低30% |
| 输入分辨率 | 从640x640降到480x480 | 速度提升40%,精度下降可接受 |
| 模型剪枝 | 移除冗余通道 | 模型体积缩小50% |
我个人习惯,在正式部署前先用trtexec跑一遍性能测试,看看吞吐量和延迟是否达标。如果发现瓶颈,优先检查数据预处理和后处理部分——很多时候,这些非推理环节反而成了性能短板。
instance_group [{ count: 2 }],就能让两个模型实例并行处理请求,充分利用Jetson的GPU资源。
好了,这一章的内容就到这里。从模型导出到实时Pipeline,再到Triton服务化,我们完整走了一遍边缘AI部署的流程。你可以在自己的Jetson上动手试试,把YOLOv8跑起来,感受一下边缘推理的魅力。