AI智能摘要
本文提供了一份从零开始的保姆级TensorRT量化部署教程。你将学会如何配置环境、拉取模型、导出ONNX并进行FP16量化编译，最终在Orin NX上实现高效的实时推理，附带完整代码和避坑指南。
— AI 生成的文章内容摘要

本机配置

  • Jetson Orin NX 16G
  • Ubuntu 22.04
  • Jetpack 6.2.1

基本要求：请确保你的 JetPack > 6，SWAP > 20G (很重要，不然会 OOM)

调整最佳性能模式

接下来的操作都需要较高的性能，所以我们通过 jtop 来调整释放最高性能

sudo jtop

按下 6 键或直接点击 CTRL 进入控制面板

1768028418 image

将风扇配置调为 cool，Jetson Clocks 开启，能耗模式调整为 MAXN

1768028481 image

输入以下命令暂时关闭桌面节省内存

sudo init 3 # 3 改为 5 或重启恢复桌面

构建并运行容器环境

首先先拉取我的项目，然后构建基于 pytorch 的深度学习环境

git clone -b modelscope https://github.com/Shattered217/SAM3-TensorRT.git
cd SAM3-TensorRT/
sudo docker build -t sam3-trt-aarch64 -f docker/Dockerfile.aarch64 .

通过 modelscope 拉取 SAM3 的模型（由于 meta 对 CN 用户不友好，所以如果要去 huggingface 申请仓库访问权限务必使用美国身份，最好主页也跟着改一下，不然大概率被拒

sudo apt-get install -y git-lfs # git 拉取大文件必备
git clone https://www.modelscope.cn/facebook/sam3.git

运行容器

sudo docker run -it --rm \
  --network=host \
  --gpus all \
  --ipc=host \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  --runtime=nvidia \
  --env HF_TOKEN \
  -v "$PWD":/workspace \
  -w /workspace \
  sam3-trt-aarch64 bash

导出 ONNX

python3 python/onnxexport.py

注意在导出的脚本中预定义了 prompt 为 “dog”，后期如需修改则要重新导出

trt 量化编译

/usr/src/tensorrt/bin/trtexec --onnx=onnx_weights/sam3_static.onnx --saveEngine=sam3_fp16.plan --fp16 --verbose
1768030114 image

出现 PASSED 即测试通过，如果 FAILED 也没事，一般是内存不够导致测试失败了，只要 .plan 生成成功即可

1768030191 image

编译并运行测试程序

编译

mkdir cpp/build && cd cpp/build
cmake ..
make

由于我们预训练的提示词是 狗 ，所以我们需要准备狗的测试图片，第一个参数的测试集文件夹，第二参数是模型路径，导出图片位于 ./results

root@jetson-orin-nx:/workspace/cpp/build# ./sam3_pcs_app test/ ../../sam3_fp16.plan

效果展示

1768030771 image
1768030787 dog
1768030838 n02091831 573
1768030822 n02091831 573

特别感谢以下项目

https://github.com/dataplayer12/SAM3-TensorRT