求助：TensorRT与PyCUDA（或其他CUDA Python接口库）学习资源及Jetson Orin Nano上PyTorch模型TensorRT优化部署问题

兄弟我太懂你这种火烧眉毛的感觉了！从舒适的PyTorch训练圈跳去搞TensorRT，简直像进了个完全陌生的新世界，官方文档又散又绕，AI生成的代码还一堆C++ runtime错误，搜半天也没找到能直接落地的教程，急得头大对吧？结合我自己踩过的坑，给你整理点能快速上手的资源和实操步骤，争取在老板度假回来前搞定！

一、快速上手的学习资源（无外链纯干货） #

• NVIDIA官方TensorRT Python示例仓库：里面有现成的PyTorch转TensorRT的代码片段，都是针对Jetson平台优化过的，直接改改模型路径、输入维度就能跑，比啃官方文档高效10倍
• Jetson官方AI实验室教程：专门针对Orin/Nano这类边缘设备，从PyTorch导出ONNX到TensorRT构建引擎全流程一步步走，还附了常见错误的排查方法，比如动态batch适配、FP16精度优化这些高频坑
• PyCUDA官方入门教程：如果需要手写简单的CUDA算子配合TensorRT用，这个教程是纯Python的，不用碰复杂的C++编译流程，适合快速上手内存管理、数据拷贝这些基础操作
• 国内技术博客实战笔记：搜“Jetson Orin TensorRT PyTorch部署”能找到不少从业者的踩坑记录，比如ONNX算子不兼容、版本匹配问题、推理时内存溢出这些实际场景的解决方案，都是实打实的经验

二、Jetson Orin Nano上PyTorch转TensorRT实操避坑指南 #

1. 先把PyTorch模型导出成ONNX（避坑第一关） #

• 导出前必须把模型切到评估模式：model.eval()，关掉dropout、BN这些训练专属层，不然TensorRT解析会直接报错
• 如果需要动态batch，导出时一定要指定动态维度：

  import torch
  dummy_input = torch.randn(1, 3, 224, 224).cuda()  # 对应你的模型输入维度
  torch.onnx.export(
      model,
      dummy_input,
      "model.onnx",
      dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},  # 动态batch配置
      opset_version=12  # 选TensorRT支持的opset版本，Jetson 5.1对应opset12没问题
  )
  

• 用onnxruntime先验证ONNX模型是否正常：

  import onnxruntime as ort
  sess = ort.InferenceSession("model.onnx")
  output = sess.run(None, {"input": dummy_input.cpu().numpy()})
  

  这一步能提前排除导出的模型本身的问题，避免后面白忙活

2. 用TensorRT构建推理引擎（Jetson专属优化） #

• 先拿trtexec工具做预测试，这个工具是排查问题的神器，命令：

  trtexec --onnx=model.onnx --saveEngine=model.trt --fp16 --workspace=4096
  

  --fp16一定要开，Jetson Orin Nano对FP16支持很好，性能能提好几倍；--workspace给够显存，避免构建时内存不足
• 如果trtexec能成功生成引擎，再用Python API加载推理：

  import tensorrt as trt
  import pycuda.driver as cuda
  import pycuda.autoinit
  import numpy as np
  
  TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
  # 加载引擎
  with open("model.trt", "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
      engine = runtime.deserialize_cuda_engine(f.read())
  # 创建推理上下文
  with engine.create_execution_context() as context:
      # 分配内存
      input_idx = engine.get_binding_index("input")
      output_idx = engine.get_binding_index("output")
      input_mem = cuda.mem_alloc(dummy_input.numpy().nbytes)
      output_mem = cuda.mem_alloc(output[0].nbytes)
      bindings = [int(input_mem), int(output_mem)]
      # 拷贝输入数据到CUDA内存
      cuda.memcpy_htod(input_mem, dummy_input.numpy())
      # 推理
      context.execute_v2(bindings)
      # 拷贝输出数据到CPU
      output_data = np.empty_like(output[0])
      cuda.memcpy_dtoh(output_data, output_mem)
  

3. 紧急错误排查（老板回来前救急） #

• 遇到C++ runtime错误：先看trtexec的报错日志，大部分是ONNX算子不兼容，比如自定义算子或者某些高版本PyTorch算子，换成torch.nn.functional里的标准算子再导出试试
• 版本不匹配问题：Jetson的JetPack版本要和TensorRT、PyTorch、CUDA严格对应，比如JetPack 5.1对应TensorRT 8.5、PyTorch 2.0，别乱装最新版
• 动态batch崩溃：导出ONNX和构建引擎时都要指定动态维度，推理前还要用context.set_binding_shape(input_idx, (batch_size, 3, 224, 224))设置当前batch的形状

按照这个流程走，两天时间绝对够跑通！先拿最小的测试模型验证流程，再替换成你的业务模型，遇到问题先查trtexec的报错，比瞎调Python代码有用多了！

备注：内容来源于stack exchange，提问作者Trung