jetson nano yolo tensorrt

### 如何在 Jetson Nano 上用 TensorRT 优化 YOLO 模型部署 #### 准备工作 为了成功地在Jetson Nano上使用TensorRT优化YOLO模型，需先准备好必要的开发环境。这包括安装适用于Jetson Nano的操作系统以及设置用于转换和运行模型所需的软件包。 #### 训练自定义YOLOv5模型 可以在主机上训练自己的YOLOv5模型，并将其保存以便后续转换为适合Jetson Nano使用的格式[^1]。此阶段涉及数据集准备、调整超参数直至获得满意的性能指标为止。 #### 将PyTorch模型转化为ONNX格式 完成训练之后，下一步就是把由PyTorch框架构建出来的YOLOv5网络结构及其权重信息导出成通用中间表示形式——即ONNX文件。这一操作可以通过调用`torch.onnx.export()`函数实现： ```python import torch from models.experimental import attempt_load model = attempt_load('path/to/best.pt', map_location=torch.device('cpu')) # 加载.pth或.pt模型 dummy_input = torch.randn(1, 3, imgsz, imgsz).to(torch.device('cpu')) input_names = ["image"] output_names = ['output'] dynamic_axes = {'image': {0: 'batch_size'},'output': {0: 'batch_size'}} torch.onnx.export(model, dummy_input, "best.onnx", input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, opset_version=12) ``` #### ONNX至TensorRT的转换 有了上述生成好的`.onnx`文件后，在Jetson Nano设备上通过命令行工具`trtexec`可以方便快捷地把它编译成为专有的TensorRT引擎文件（通常以`.engine`结尾），从而充分利用GPU硬件特性提升推理速度[^2]。 ```bash sudo /usr/src/tensorrt/bin/trtexec --onnx=/home/user/best.onnx --saveEngine=/home/user/best.engine ``` #### 编写Python脚本集成CSI摄像头与TensorRT推断流程 最后一步则是编写一段能够读取来自CSI接口连接之相机图像流的同时调用之前创建完毕的TensorRT引擎来进行目标识别任务的小程序。这里给出一个简单的例子作为参考[^3]: ```python import cv2 import numpy as np import tensorrt as trt import pycuda.driver as cuda import pycamera.array as carray def load_engine(trt_runtime, engine_path): with open(engine_path, 'rb') as f: engine_data = f.read() engine = trt_runtime.deserialize_cuda_engine(engine_data) return engine TRT_LOGGER = trt.Logger() runtime = trt.Runtime(TRT_LOGGER) with load_engine(runtime, "/home/user/best.engine") as engine: context = engine.create_execution_context() cap = cv2.VideoCapture(gstreamer_pipeline(), cv2.CAP_GSTREAMER) while True: ret_val, frame = cap.read() if not ret_val: break # Preprocess image here... inputs, outputs, bindings, stream = allocate_buffers(engine) do_inference(context, bindings=bindings, inputs=[frame], outputs=outputs, stream=stream) detections = post_processing(outputs) draw_boxes(frame, detections) cv2.imshow("CSI Camera", frame) keyCode = cv2.waitKey(30) & 0xFF if keyCode == ord('q'): break cap.release() cv2.destroyAllWindows() ```