yolo转tensorrt nano
目录
需要onnxruntime-gpu，弃用
不需要onnxruntime-gpu
安装cmake
安装onnxsim ok
Windows 上导出 ONNX、GPU 测试 到 Jetson Nano 上生成 TensorRT engine 并推理
一、Windows 上导出 ONNX
二、拷贝 ONNX 到 Jetson Nano
三、在 Nano 上生成 TensorRT Engine
四、Nano 上 TensorRT 推理示例（Python）
pip install onnx>=1.12.0
pip install onnx onnxsim
需要onnxruntime-gpu，弃用
from ultralytics import YOLO# 加载模型
model = YOLO("models/best_0814.pt")# 导出 TensorRT FP16 engine
model.export(format="engine", device=0, half=True)不需要onnxruntime-gpu
model = YOLO(r"/home/jszh/project/jump3d/RopeCount-main/models/best_0814.pt") # YOLO‑12 检测 + 跟踪# model.export(format="engine", device=0, half=True, dynamic=False, simplify=True)export_path = model.export(format="engine",device=0,half=True,dynamic=False,simplify=True,
)print('export_path',export_path)安装cmake
博客地址：
https://blog.csdn.net/jacke121/article/details/151937445
安装onnxsim ok
# 设置必要的环境变量
export CMAKE_POLICY_VERSION_MINIMUM=3.5
export ONNX_OPT_USE_SYSTEM_PROTOBUF=ON# 清除之前的构建缓存
pip cache purge #这句如果报错就不要了。
rm -rf ~/.cache/pip# 重新安装
pip install onnxsim --no-cache-dir --force-reinstallWindows 上导出 ONNX、GPU 测试 到 Jetson Nano 上生成 TensorRT engine 并推理
，附带示例代码
一、Windows 上导出 ONNX
-
安装依赖：
pip install ultralytics onnx onnxruntime-gpu numpy
-
导出 ONNX：
from ultralytics import YOLO # 加载训练好的模型 model = YOLO(r"C:\path\to\best_0814.pt") # 导出为 ONNX export_path = model.export( format="onnx", opset=17, # ONNX 版本 simplify=True, # 可选，简化模型 ) print('ONNX 模型路径:', export_path)
-
测试 ONNX 模型 GPU 推理：
import onnxruntime as ort import numpy as np # 加载 ONNX 模型，使用 GPU sess = ort.InferenceSession(export_path, providers=['CUDAExecutionProvider']) input_name = sess.get_inputs()[0].name # 假设输入是 640x640 RGB x = np.random.randn(1, 3, 640, 640).astype(np.float32) # 推理 outputs = sess.run(None, {input_name: x}) print("输出数量:", len(outputs)) print("输出形状:", [o.shape for o in outputs])
二、拷贝 ONNX 到 Jetson Nano
-
将
best_0814.onnx拷贝到 Nano（例如/home/jetson/models/）：
scp best_0814.onnx jetson@nano_ip:/home/jetson/models/
-
Nano 上安装 TensorRT：
sudo apt update sudo apt install python3-pip pip3 install nvidia-pip # 或直接使用 JetPack 自带的 TensorRT
三、在 Nano 上生成 TensorRT Engine
使用官方 trtexec 工具：
cd /home/jetson/models/ # FP16 精度生成 engine trtexec --onnx=best_0814.onnx --saveEngine=best_0814_fp16.trt --fp16 --workspace=1024
-
--fp16：使用半精度加速 -
--workspace：显存临时缓冲区大小（MB） -
如果 Nano 内存有限，可以适当调整
生成的
best_0814_fp16.trt就是针对 Nano GPU 优化的 TensorRT engine，可以直接推理。
四、Nano 上 TensorRT 推理示例（Python）
import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # 加载 engine with open("best_0814_fp16.trt", "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() # 假设输入是 1x3x640x640 input_shape = (1, 3, 640, 640) x = np.random.randn(*input_shape).astype(np.float32) # 分配 GPU 内存 d_input = cuda.mem_alloc(x.nbytes) d_output = cuda.mem_alloc(1 * 1000 * 7 * 4) # 示例，需根据实际输出调整 bindings = [int(d_input), int(d_output)] # 将输入拷贝到 GPU cuda.memcpy_htod(d_input, x) # 推理 context.execute_v2(bindings) # 从 GPU 拷贝输出 output = np.empty([1, 1000, 7, 4], dtype=np.float32) # 示例形状 cuda.memcpy_dtoh(output, d_output) print("TensorRT 输出 shape:", output.shape)
注意：
d_output和output的大小需要根据你的 YOLO 模型输出调整，一般是[batch, num_boxes, 7]或[batch, num_classes, h, w]，可以用engine.get_binding_shape(i)获取。