Jetson Orin NX 刷机、配置与使用_orin nx刷机

本文从Jetson Orin NX 的刷机开始，介绍该设备的具体配置与代码开发方法。

目录

1 设备信息与准备工作

1.1 固态格式化与分区

1.2 准备ubuntu虚拟机

2 刷机

2.1 安装SDK manager

2.2 Jetson Orin NX设置

2.3 使用SDK manager刷机

3 NX的其他环境配置

3.1 jetson-stats

3.2 opencv

3.3 python环境相关

3.3.1 pip换源

3.3.2 miniconda

3.3.3 torch + torchvision

3.3.4 yolo

3.4 jetson-inference + jetson-utils

3.4.1 简介

3.4.2 安装

3.5 ffmpeg

3.6 JetsonGPIO

4 NX的使用

4.1 普通代码的编写与构建

4.1.1 用cmake管理项目

4.1.2 一份代码示例

4.2 TensorRT深度学习模型加速

4.2.1 借助jetson-inference部署nvidia指定的网络

4.2.2 借助TensorRTx部署当前最流行的网络

4.2.3 从零开始部署你自己的网络

4.2.3.1 在服务器上的模型开发

4.2.3.2 engine的构建

4.2.3.3 engine的使用

4.3 OpenCV的使用

4.4 CUDA函数的编写

4.5 其他

5. 手册&文档

---

1 设备信息与准备工作

买的是亚博智能的板子，附带有一条128G的nvme固态

亚博智能官方教程链接：Jetson Orin NX教程，密码为ntgy

1.1 固态格式化与分区

如果拿到的固态不是空的，则首先需要自己格式化，否则烧写系统会失败。用到的工具是DiskGenius。需要自行准备一个额外的设备：固态盒子。

将固态放入固态盒子插到电脑上，在DiskGenius里右键选中，删除所有分区，建立新分区。

出现下图时，新分区大小填最大值，文件系统类型选择ext4。

在界面左上角点击保存更改，并进行格式化。

1.2 准备ubuntu虚拟机

前往Ubuntu 20.04.6 LTS (Focal Fossa)下载ubuntu2004镜像(经验证，低版本的ubuntu可能只能装jetpack5，而更高版本的ubuntu只能装jetpack6，2004是两个都能装)。

在vmware中创建虚拟机，这里最重要的就是虚拟机的磁盘大小，最好给个80G以上。踩坑的过程中曾经试过40G，下载时不小心多勾选了一些不必要的内容，空间就不够了。

这是最后虚拟机磁盘占用的情况。

2 刷机

参考文章：Jetson Xavier NX刷机安装

2.1 安装SDK manager

前往英伟达官网下载SDK manager SDK Manager | NVIDIA Developer

’

选择下载.deb文件（可能需要注册英伟达账号）

将.deb文件传到虚拟机里，在终端使用sudo apt install 命令安装，完成后在ubuntu的开始菜单内可以打开。

2.2 Jetson Orin NX设置

将固态硬盘插到板子上，在板子的背面(天线下方)找到一排插针，使用杜邦线连接第三个引脚FC_REC与第四个引脚GND。

找到正面的type-c接口，将其连接到电脑上(官方给的线是两头type-c的，可能需要用到一个type-c转type-a的转接器)。

为nx接上屏幕，刷机到一定进度的时候nx会开机亮屏，接上屏幕以免错过这个过程，误以为安装过程卡死。

接上19V电源适配器开机，在虚拟机弹出的界面中选择将该设备连接到虚拟机。

2.3 使用SDK manager刷机

使用英伟达账号登陆后来到step01-step04界面

step01

注意项如图

第一，Host Machine的勾要去掉，这个是用来在ubuntu虚拟机上下载CUDA环境的，我们并不需要，取消后可以少下载十几个G的资源。第二，Target Hardware 可以点击刷新按钮自动匹配，Orin有8G和16G两种，注意选择正确的内存版本。第三，JetPack选择5.1.3，虽然JetPack6.0提前预览版已经发布(写于2024.4.29)，但经过我的尝试，选择6.0会在最后一步写入固态的时候报错，导致系统烧写不正常。

step02

勾选同意协议并开始下载，下完了会自动开始烧写，全程保持NX的usb线和主机连上。

step03

进度条到中间时会跳一个类似下面的窗口，四个需要选择的框，前两个选你板子的型号，第三个选Runtime，第四个选择NVMe。然后点击flash。

等待一段时间后又会跳出一个新的窗口，同时NX的ubuntu系统会跑起来，显示屏亮起。这时给NX接上鼠标键盘，和主机连上同一个WIFI。用ifconfig命令查看NX的IP地址，并且填到虚拟机跳的这个框里。点击install，会为NX安装剩余组件。

如果刚才这一步NX没有正常进系统，而是显示屏亮了一下就灭了，则可以尝试以下重新给NX上电，如果还是没法正常进系统，则说明烧写失败了，回到step01重新开始。(为了防止烧写出错，我实际上在重新开始前，还重新格式化了固态)

step04

全部烧写完毕后SDK manager会自动进入step04，此时可以断开NX和主机的usb连接，拔掉NX的跳线，重新启动NX，开始正常使用。

3 NX的其他环境配置

首先换源，去清华源网站清华大学开源软件镜像站

备份/etc/apt/sources.list文件后

sudo gedit /etc/apt/sources.list

将清华源的内容添加到文件末尾

3.1 jetson-stats

jetson-stats是监控jetson设备运行状态的工具，可在命令行用jtop命令打开

sudo -H pip3 install jetson-stats
jtop

翻到第7页INFO，可以查看系统各依赖库的版本，刚安装时OpenCV是没有和CUDA联合编译的，那一行应该是false，还需重新对opencv进行重装。

3.2 opencv

参考此博客，先卸载原本的版本

sudo apt-get purge libopencv*

然后去github下载opencv和opencv-contrib的源码，我个人喜欢用3.4.15版本的(跟不上时代了)，习惯用opencv4的话可以装4。

分别解压后将opencv-contrib文件夹移动到opencv文件夹下

在官方文档里找到需要预先安装的包，官网要求的如下：

sudo apt-get install build-essential
sudo apt-get install cmake git libgtk2.0-dev pkg-config libavcodec-dev libavformat-dev libswscale-dev
sudo apt-get install python-dev python-numpy libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff-dev libjasper-dev libdc1394-22-dev

在opencv文件夹下执行

mkdir build
cd build
 
cmake -D WITH_CUDA=ON -D CUDA_ARCH_BIN="7.2" -D WITH_cuDNN=ON -D OPENCV_DNN_CUDA=ON  -D cuDNN_VERSION='8.6' -D cuDNN_INCLUDE_DIR='/usr/include/' -D CUDA_ARCH_PTX="" -D OPENCV_EXTRA_MODULES_PATH=../opencv_contrib-3.4.15/modules -D WITH_GSTREAMER=ON -D WITH_LIBV4L=ON -D BUILD_opencv_python3=ON -D BUILD_TESTS=OFF -D BUILD_PERF_TESTS=OFF -D BUILD_EXAMPLES=OFF -D CMAKE_BUILD_TYPE=RELEASE -D CMAKE_INSTALL_PREFIX=/usr/local .. 
 
sudo make -j10
sudo make install

等待很久才会完成，最后配置系统变量

sudo gedit /etc/ld.so.conf.d/opencv.conf
# 写入下面这行
/usr/local/lib
# 保存退出
sudo ldconfig

重启后jtop就能看到OpenCV compiled with CUDA: yes了

3.3 python环境相关

这一步包含的全是python相关包的安装。由于我们还希望NX上能跑一些python代码，所以还需要在板子上装torch、yolo等相关的包。

3.3.1 pip换源

执行以下命令

cd ~
mkdir ~/.pip
sudo gedit ~/.pip/pip.conf
# 写入以下内容
[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple
[install]
trusted-host =
    pypi.tuna.tsinghua.edu.cn
# 保存退出

3.3.2 miniconda

之前在Jetson Xavier NX(jetpack 5.1.1)上我是习惯装一个miniconda虚拟环境的，也没出过错，但是在Jetson Orin NX(jetpack 5.1.3)上我在miniconda的虚拟环境里编译torchvision时总是遇到报错，最后也没解决。nvdia论坛上有人提到torchvision安装失败和虚拟环境有关，所以暂时不清楚是miniconda导致的，还是jetpack5.1.3+orin NX本身带来的bug。这里还是写一下minconda的安装过程，读者可自行选择是否安装。

首先在miniconda官网上下载安装包，选择aarch64的版本

然后执行下面命令即可

chmod +x ./Minconda3-xxxx.sh
sudo ./Minconda3-xxxx.sh

中间会问你要不要在终端启动时就激活conda环境，根据自己情况选择一下，其余的选项要么是直接回车要么是yes。

conda创建的环境最好是python3.8的，后续的torch限定了这个python版本。

conda换源：

sudo gedit ~/.condarc
# 写入以下内容
channels:
  - defaults
show_channel_urls: true
default_channels:
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
# 保存退出

3.3.3 torch + torchvision

英伟达为jetson设备提供了专门的torch包，需要在论坛上按照操作安装。如果直接用pip安装的话就是cpu版，没法GPU加速的(如果是pip安装其他的包，里面依赖torch，也会给你自动安装一个cpu版的torch，很头疼)。

whl文件的下载链接，torchvision从源码编译的方法，以及torch和torchvision的对应关系都在这个论坛页里有，这里就不重复了。 中间出了什么问题也可以在论坛里搜索下，我这边按照这个流程只能装好gpu版本的torch，torchvision安装没成功，原因我上面分析了，可能是使用了miniconda导致的。

3.3.4 yolo

我这里先创了一个新的conda环境，避免前面装好的torch和yolo依赖的包冲突。

最新的yolov8用起来已经很方便了，根据ultralytics的readme，一行代码就能安装好

pip install ultralytics

从release页面下载yolov8s.pt到本地

测试安装是否成功

yolo predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg'

3.4 jetson-inference + jetson-utils

3.4.1 简介

jetson-utils是英伟达为jetson设备提供的C++开发SDK，使用C++和CUDA实现了许多功能，方便开发者直接调用，具体包括以下功能：

其中camera、cuda、display、image、threads下的函数都是软件开发中经常需要使用的。从图像流的获取到图像的预处理再到图像渲染到屏幕，整个过程中包含的各种函数这里都提供了经过CUDA加速的版本。

jetson-inference是基于jetson-utils进一步开发的深度学习应用，可以使用它提供的sample编译出能直接运行深度学习网络的可执行文件，具体功能包括：

3.4.2 安装

介绍完背景，就进入这两个库的安装。

jetson-inference从源码编译的教程在这，按照它说的执行即可，jetson-utils作为它的一个submodule，会自动一块编译

sudo apt-get update
sudo apt-get install git cmake
 
git clone https://github.com/dusty-nv/jetson-inference
cd jetson-inference
git submodule update --init
 
sudo apt-get install libpython3-dev python3-numpy
 
cd jetson-inference
mkdir build
cd build
cmake ../
make -j$(nproc)
sudo make install
sudo ldconfig

网络好的时候上述操作是没问题的。但由于一些原因，板子是没法科学上网的，所以我常常遇到git clone这一步就失败的情况。就算git clone成功了，也会卡死在git submodule update --init这步。如果跳过git submodule update --init，则会因为缺失组件，cmake都通不过。之前我被这个小问题卡了很久，很难受。最后选择了一个最笨的办法：用科学上网的windows主机访问jetson-inference，然后download zip，这样下载下来的项目内部各submodule的文件夹里都是空的，然后再打开.gitmodules文件，挨个访问这些url，download zip，再把其中的代码解压出来放到指定的path里。

手动组装后把完整的代码传输到板子上，最后打开jetson-inference的CMakeLists.txt，把submodule warning下面的语句都注释掉，就能成功编译了。

最后在make install时会跳一个error

make install error: jetson-inference/utils/python/python/jetson does't exist

这是从windows拷贝代码到linux导致的一个bug，在jetson-inference/utils/python/python/路径下创建一个名为jetson的空文件夹即可解决问题。详见这个issue

因为库都是由源码编译的，所以我们可以自由修改其中的代码，改成你想要的功能。比如我曾经想用gstBufferManager里的一个成员变量，但是它原本被设为protected，外部没权限访问，我就直接改了源码，把它放到public里，然后重新编译了这两个库。再比如glDisplay的默认窗口最上方会有一个白边，导致其实没法全屏显示(类似于游戏里窗口化和无边窗口化)，这里其实是可以通过改源码把这个白边去掉的。改动后重新编译的命令如下：

cmake ../
make -j$(nproc)
sudo make install
sudo ldconfig

3.5 ffmpeg

这个需求比较小众，在jetson设备上大部分情况下还是会使用gstreamer来做多媒体编解码。我是因为需要用到ffplay这个工具而系统自带的版本过低才想到重装ffmpeg。仅作个人记录，可跳过。

前往ffmpeg官网下载source code，解压后执行

cd ffmpeg-7.0
sudo apt-get update -qq && sudo apt-get -y install \
  autoconf \
  automake \
  build-essential \
  cmake \
  git-core \
  libass-dev \
  libfreetype6-dev \
  libgnutls28-dev \
  libmp3lame-dev \
  libsdl2-dev \
  libtool \
  libva-dev \
  libvdpau-dev \
  libvorbis-dev \
  libxcb1-dev \
  libxcb-shm0-dev \
  libxcb-xfixes0-dev \
  meson \
  ninja-build \
  pkg-config \
  texinfo \
  wget \
  yasm \
  zlib1g-dev
 
sudo apt install libunistring-dev libaom-dev libdav1d-dev
sudo apt-get install libsdl2-2.0 libsdl2-dev
 
./configure  --prefix=/usr/local/ffmpeg --enable-shared --enable-ffplay

打开ffbuild/config.mak文件，搜索一下ffplay，找到一行写着CONFIG_FFPLAY=yes，确保前面没有感叹号!。有感叹号说明不会生成ffplay。

make -j6
sudo make install

配置环境变量

sudo gedit ~/.bashrc
# 添加内容
export PATH="/usr/local/ffmpeg/bin:${PATH}"
export LD_LIBRARY_PATH="/usr/local/ffmpeg/lib:{LD_LIBRARY_PATH}"
# 保存退出
source ~/.bashrc
 
ffmpeg -version
ffplay -version

3.6 JetsonGPIO

jetson设备都包含一些插针形式的I/O接口，40pin的管脚对照表可以在jetsonhacks上查到，内部集成了SPI、I2C、UART等接口。官方提供了一套用python控制GPIO口的工具jetson-gpio，但如果项目需要我们用C++控制GPIO口，则需要使用github上大神编写的开源库JetsonGPIO。这两个库其实就是使用的编程语言不一样，其他函数用法基本一致。

官方python版在刷机时就已经安装好了，如果没有，就用pip安装。

C++版的根据其安装指引，执行以下命令：

git clone https://github.com/pjueon/JetsonGPIO
cd JetsonGPIO
mkdir build
cd build
cmake .. -DBUILD_EXAMPLES=ON
cmake --build . --target examples

4 NX的使用

4.1 普通代码的编写与构建

4.1.1 用cmake管理项目

如前文所述，在jetson设备上的软件开发离不开jetson-utils提供的接口，所以开发中必须把jetson-utils包含到项目里。其次，图像处理过程中常用到OpenCV，也需要包含到项目里。综合考虑，肯定是要用cmake来管理项目，而项目又会分很多功能，最好是用子文件夹的方式把各功能区分开，假如目录结构如下：

project
    -camera
        -myCamera.cpp
        -myCamera.h
        -cameraTest.cpp
        -CMakeLists.txt
    -main
        -main.cpp
        -CMakeLists.txt
    -CMakeLists.txt

则最顶层CMakeLists.txt的示例如下：

cmake_minimum_required(VERSION 3.10)
project(project)
 
# opencv & cuda
find_package(OpenCV REQUIRED)
find_package(CUDA REQUIRED)
 
# jetson-utils
find_library(/home/xxx/jetson-inference/build/aarch64/lib NAMES jetson-utils REQUIRED)
 
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/bin)
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib)
set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib)
file(MAKE_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY})
file(MAKE_DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
file(MAKE_DIRECTORY ${CMAKE_ARCHIVE_OUTPUT_DIRECTORY})
 
# gstreamer
# include_directories(/usr/include/gstreamer-1.0 /usr/lib/aarch64-linux-gnu/gstreamer-1.0/include /usr/include/glib-2.0 /usr/include/libxml2 /usr/lib/aarch64-linux-gnu/glib-2.0/include/)
 
include_directories( ${PROJECT_BINARY_DIR} ${PROJECT_SOURCE_DIR})
 
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
add_subdirectory(camera)
add_subdirectory(main)

这份cmake文档可以把CUDA、jetson-utils和OpenCV的依赖库都加到项目的路径里来，并指定了两个subdirectory，我们需要创建这两个子文件夹并在里面编写对应功能的.h和.cpp。子文件夹的代码也同样用cmake构建，下面给出一个子文件夹下的CMakeLists.txt示例

get_directory_property(hasParent PARENT_DIRECTORY)
if(hasParent)
    cuda_add_library(l_camera myCamera.cpp)
    target_link_libraries(l_camera jetson-utils)
else()
    cmake_minimum_required(VERSION 3.10)
    project(cameraTest)
    find_library(/home/xxx/jetson-inference/build/aarch64/lib NAMES jetson-utils REQUIRED)
    find_package(CUDA REQUIRED)
    find_package(OpenCV REQUIRED)
 
    cuda_add_executable(cameraTest cameraTest.cpp myCamera.cpp)
    target_link_libraries(cameraTest ${OpenCV_LIBS} jetson-utils)
endif()

此处用get_directory_property(hasParent PARENT_DIRECTORY)来判断是否作为子模块进行cmake，并执行不同的行动。若作为子模块，则生成一个名为l_camera的静态库供主模块使用；若作为单独模块，则指定可执行文件名为cameraTest，并为其链接上opencv和jetson-utils库。

接下来要做的就是编写.cpp和.h文件了，根据需要的功能调用opencv和jetson-utils的函数，这里就不展开了。

4.1.2 一份代码示例

这里给出一个示例代码及其CMakeLists.txt帮助初学者了解代码结构，这个示例可以运行，用于读取并显示一张图片。

目录结构：

sample

        -bird_0.jpg

        -CMakeLists.txt

        -test.cpp

CMakeLists.txt内容

cmake_minimum_required(VERSION 3.10)
project(test)
find_package(OpenCV REQUIRED)
find_package(CUDA REQUIRED)
find_library(/home/orin/jetson-inference/build/aarch64/lib NAMES jetson-utils REQUIRED)
 
cuda_add_executable(test test.cpp)
target_link_libraries(test ${OpenCV_LIBS} jetson-utils)

test.cpp内容

#include <jetson-utils/cudaUtility.h>
#include <jetson-utils/glDisplay.h>
#include <jetson-utils/cudaMappedMemory.h>
#include <jetson-utils/imageIO.h>
#include <unistd.h>
 
int main( int argc, char** argv )
{
    int width = 368;
    int height = 500;
    void* img = nullptr;
    cudaAllocMapped((void**)&img,  width * height * sizeof(uchar3));
    loadImage("../bird_0.jpg", &img, &width, &height, IMAGE_RGB8);
 
	glDisplay* dis = glDisplay::Create(NULL, 600, 600);
 
    for(int i = 0; i < 100; i++)
    {
        dis->BeginRender();
        dis->RenderImage(img, width, height, IMAGE_RGB8, 0, 0);
        dis->EndRender();
    }
	SAFE_DELETE(dis);
    cudaFreeHost(img);
	printf("shutdown complete.\n");
	return 0;
}

这里有几个重点，第一是在使用jetson-utils做图像处理时，首先要使用cudaAllocMapped为图像分配零拷贝内存，这个内存是CPU-GPU之间共享的，无需我们自己进行cpu内存到gpu显存的拷贝，回收时使用cudaFreeHost。第二就是使用glDisplay显示图像，可以去读jetson-utils/glDisplay.h的源码，了解各函数的使用方法。

运行这份代码的命令如下：

mkdir build
cd build
cmake ..
make 
./test 

4.2 TensorRT深度学习模型加速

虽然前面也介绍了python环境的搭建，可以直接用pytorch跑一些深度学习代码，但是为了提高运行效率，方便和项目其他功能配合，最终还是得用C++跑TensorRT模型。这里的过程就比较麻烦，分成两个大块。第一，我们要把训练好的模型转换成TensorRT的engine，第二，要用C++代码实现读图、跑模型、输出结果的功能。

在第一步转换模型时，就分为三种情况。

第一种情况是你要做的任务包含在jetson-inference里，要用的模型也正好包含在jetson-inference里。比如要做Image Recognition，想用的模型是googlenet，去官网一看，发现都被包含在jetson-inference了，那就省事很多，因为jetson-inference已经帮你把一整套东西写好了，imageNet已经被包装成一个类，在<jetson-inference/imageNet.h>里可以查看源码，照着官方教程给出的demo修改即可。

第二种情况是你要转换的模型不是jetson-inference里包含的模型，但它是当前很流行的模型(比如yolo)，那么在github社区里有大神已经帮你写好了转化的代码，项目叫做TensorRTx，这个项目里实现的模型很多，从alexnet到yolo系列，再到swin-Transformer都有，大多数都是目标检测的网络。

第三种情况就是你要转换的模型完全是自己开发的，那就需要自己从头开始做转换了。

4.2.1 借助jetson-inference部署nvidia指定的网络

jetson-inference编译好以后，在jetson-inference/build/aarch64/bin/目录下会出现几个可执行文件，和深度学习模型有关的有imagenet、detectnet、segnet、posenet、actionnet、backgroundnet、depthnet，此外还有camera-capture和video-viewer两个小工具。

每种任务都由对应的ReadMe教程，这里以imagenet为例讲解。根据文档，直接使用这个可执行文件进行推理的方法是如下：

./imagenet --network=<network-name> <input-path> <output-path>
# 例
./imagenet --network=resnet-18 images/jellyfish.jpg images/test/output_jellyfish.jpg
./imagenet --network=resnet-18 jellyfish.mkv images/test/jellyfish_resnet18.mkv

其中network需要从支持的表格里选择

如果执行命令前模型文件还不存在的话，就会自动去外网下载资源，经常断线。所以比较推荐的做法是提前在windows主机上下载网络模型文件，然后拷贝到NX上。从model mirror的地址下载对应的tar.gz压缩包，进入jetson-inference/data/network/目录，创建一个和压缩包同名的文件夹，将压缩包内的文件拷贝到同名文件夹里。如图所示：

input-path和output_path就是两个文件的路径，可以支持图片或视频。

第一次执行命令时，终端会运行很久，这是因为程序正在由caffemodel生成TensorRT的engine文件。等它跑完后，再次执行一样的命令，就能跑得很快了。同时，在对应的文件夹里会出现一个新的engine文件。

虽然用jetson-inference编译好的可执行文件做识别很简单，但我们总得把这个功能集成到自己的项目里，自己写C++控制整个流程。这块官方也给出了自己编写代码的教程。示例代码如下：

#include <jetson-inference/imageNet.h>
#include <jetson-utils/loadImage.h>
 
int main( int argc, char** argv )
{
	// a command line argument containing the image filename is expected,
	// so make sure we have at least 2 args (the first arg is the program)
	if( argc < 2 )
	{
		printf("my-recognition:  expected image filename as argument\n");
		printf("example usage:   ./my-recognition my_image.jpg\n");
		return 0;
	}
 
	// retrieve the image filename from the array of command line args
	const char* imgFilename = argv[1];
 
	// these variables will store the image data pointer and dimensions
	uchar3* imgPtr = NULL;   // shared CPU/GPU pointer to image
	int imgWidth   = 0;      // width of the image (in pixels)
	int imgHeight  = 0;      // height of the image (in pixels)
		
	// load the image from disk as uchar3 RGB (24 bits per pixel)
	if( !loadImage(imgFilename, &imgPtr, &imgWidth, &imgHeight) )
	{
		printf("failed to load image '%s'\n", imgFilename);
		return 0;
	}
 
	// load the GoogleNet image recognition network with TensorRT
	// you can use "resnet-18" to load ResNet-18 model instead
	imageNet* net = imageNet::Create("googlenet");
 
	// check to make sure that the network model loaded properly
	if( !net )
	{
		printf("failed to load image recognition network\n");
		return 0;
	}
 
	// this variable will store the confidence of the classification (between 0 and 1)
	float confidence = 0.0;
 
	// classify the image, return the object class index (or -1 on error)
	const int classIndex = net->Classify(imgPtr, imgWidth, imgHeight, &confidence);
 
	// make sure a valid classification result was returned	
	if( classIndex >= 0 )
	{
		// retrieve the name/description of the object class index
		const char* classDescription = net->GetClassDesc(classIndex);
 
		// print out the classification results
		printf("image is recognized as '%s' (class #%i) with %f%% confidence\n", 
			  classDescription, classIndex, confidence * 100.0f);
	}
	else
	{
		// if Classify() returned < 0, an error occurred
		printf("failed to classify image\n");
	}
	
	// free the network's resources before shutting down
	delete net;
	return 0;
}

可以看到使用起来非常方便，只需要创建一个imageNet类的对象，然后调用Classify()、GetClassDesc()这两个接口即可得到结果。如前面所介绍的那样，C++代码都用cmake来构建，这里则需要在cmake里加入链接jetson-inference的代码

target_link_libraries(my-recognition jetson-inference)

4.2.2 借助TensorRTx部署当前最流行的网络

前文提到，如果你要部署的网络没有被jetson-inference包含，但被TensorRTx包含，则可以在这个项目的基础上进行模型的部署。首先下载该项目

git clone https://github.com/wang-xinyu/tensorrtx
cd tensorrtx

想部署什么网络就去找对应的ReadMe按照流程操作就行，这里以tensorrtx-yolov8为例，走一遍流程。

首先去ultralytics的model zoo里下载模型的.pt文件，YOLO系列的模型都是分好几种大小的，从小到大分别为n/s/m/l/x，这里我下载的是yolov8s.pt。值得注意的是这一步只需要下载.pt文件，并不需要yolov8的源码。

把yolov8s.pt文件放到tensorrt/yolov8/文件夹下，执行权重格式的转换

python gen_wts.py -w yolov8s.pt -o yolov8s.wts -t detect

这里如果报syntax error，说明你调起来的是系统自带的python2.x，需要用python3.x运行。结束后会出现一个新的文件yolov8s.wts。

接着构建项目：

mkdir build
cd build
cp ../yolov8s.wts .
cmake ..
make

会生成一个可执行文件yolov8_det，这个可执行文件有两种运行模式，-s是序列化模型，生成TensorRT的.engine文件，-d是反序列化模型进行推理。

# 序列化，生成engine
sudo ./yolov8_det -s yolov8s.wts yolov8s.engine s 
# 最后一个参数用于区分模型的大小，用的是yolov8s就得填s

序列化会需要比较久的时间，如果报错，则检查一下最后一个参数是否填错了。engine文件正确生成后即可使用-d功能进行推理。

# 反序列化
sudo ./yolov8_det -d yolov8s.engine ../images c # cpu postprocess
sudo ./yolov8_det -d yolov8s.engine ../images g # gpu postprocess

这里倒数第二个参数必须是一个包含图片的文件夹，最后一个参数控制nms筛除重复框的后处理用gpu做还是cpu做，实测花的时间差的不是特别多。

到这里，就已经跑通了TensorRTx的yolov8，但还是前面的那个问题，我们不能永远用别人编译好的可执行文件，总得学会自己写C++控制整个流程，这个就得需要仔细地读tensorrtx的源码，在其基础上裁剪出所需要的功能。为了能看懂这些代码，还要对TensorRT进行系统的学习，看大量的英文文档。我曾经基于tensorrtx，在Jetson Xavier NX上写过一个yolov5的功能，但时间比较久了，不知道在新的设备、新的JetPack环境下还能否跑通。感兴趣的可以在我的github项目上查看源码。对于TensorRT的engine生成后自己写C++把模型跑起来的问题，下文的4.2.3.3节还将额外介绍。

4.2.3 从零开始部署你自己的网络

从零开始部署自己的网络是一个非常麻烦的事情。整个流程可大致分为三块，在服务器上的模型开发、在NX上构建engine、在NX上使用engine跑起来。现在开发深度学习模型大多是在服务器上用Torch编写模型然后进行训练，但设计网络时基本不会考虑用到的Torch算子是否被TensorRT支持，算子不支持或者功能不匹配是进行模型转化的第一个障碍。第二步在NX上构建engine时可能会遇到各种各样的报错，有的报错甚至很难搜到相关的信息，不知道自己究竟卡在什么地方。好不容易找到一个相关的github issue，最后发现是你用的TensorRT版本问题，而TensrRT版本是你给NX刷机时jetpack自带的，根本没法升级，遇上这种情况真的会很崩溃。第三步使用engine就是考验使用jetson-utils的熟练程度，倒是没有特别麻烦。

我的项目需求是在NX上部署一个图像融合的深度学习网络，这个方向不像目标检测、跟踪一样热门，想部署就只能自己动手，下文我都会以这个需求进行举例。先简单介绍下我这个领域，红外与可见光图像融合就是读入分辨率相同的红外和可见光图像各一张，输出一张融合的图像，好处是融合图像包含的信息更丰富，后续传输或显示时只用处理这一张图，避免了红外可见光模式来回切换。

4.2.3.1 在服务器上的模型开发

在服务器上首先第一步要做的是安装一个TensorRT，最好和板子上的版本一致或相近。在NX上可以用jtop查看板子上TensorRT的版本，如果是jetpack5.1.3，那么自动安装的是8.5.2.2。服务器上能安装什么版本也不是能自由选择的，得看服务器的CUDA环境，下载链接在此。安装的目的是能在服务器上用TensorRT预先尝试转换一下，如果在服务器上都转换失败，那说明算法使用的某些算子还尚未被TensorRT支持。这里就有两条路了，如果你是模型的开发者，那么可以修改torch代码，把不支持的算子替换成别的，比如我记得早期TensorRT对SiLU激活函数的支持就不是很好，最方便的办法就是替换成ReLU，重新训练一下。第二条路就是像TensorRTx的作者一样，自己写新的TensorRT plugin，实现你需要的算子，但这个就非常难了，我是不会的。

正式开始模型转换，第一步是从torch模型转化成onnx，用python完成，示例代码如下

import torch
import torch.onnx
from models.model_cfg import net_for_test
import cfg
 
cfg.num_parallel = 3
cfg.predictor_num_parallel = 2
cfg.use_exchange = True
 
dummy_input = torch.randn(2, 1, 3, 480, 640)
Net = net_for_test()
state_dict = torch.load('./ckpt/model/checkpoint-119.pkl')
 
Net.load_state_dict(state_dict, strict=True)
 
input_names = ["Imgs"]
output_names = ["Fused"]
torch.onnx.export(Net, dummy_input, "tokenfusion_model_op11.onnx", verbose=True,
                  opset_version=11, input_names=input_names, output_names=output_names)

这里要注意的是如果你的模型有多个输入，那最好像我一样给他stack起来，做成一个tensor。我输入图像BCHW是(1,3,480,640)，但有两张一样的，自然的想法是做一个input1，一个input2，这样在转onnx时不会有什么问题，但是最终转TensorRT的时候可能会遇到一些奇怪的错误。然后就是要规定input_names和output_names，这两个名字后续会用到。最后，opset_version是会影响到转化过程的，可能会遇到opset_version填的太高，导致一些很常见的算子TensorRT居然也不支持的情况，可以多改变几个version进行尝试。

第二步，使用trtexec在服务器上尝试转换TensorRT engine。这里的trtexec是安装TensorRT时自带的一个命令行工具，用法如下

cd /xxx/TensorRT-8.x.x.x/targets/x86_64-linux-gnu/bin
 
./trtexec --onnx=/xxx/model.onnx --saveEngine=/xxx/model_fp32.engine
 
./trtexec --onnx=/xxx/model.onnx --saveEngine=/xxx/model_fp16.engine --fp16

把路径填对就行了，加了最后的--fp16就会转化成fp16的TensorRT engine，不加就是fp32的。如果这里用命令行工具转换时报错，就需要根据提示信息一点点调通，遇到算子不支持的情况时则要考虑调整原始模型的设计。

第三步，测试一下生成的engine是否正常，也是写python代码。这里需要提前pip安装一下TensorRT，在之前下载的TensorRT-8.x.x.x/python/路径下有whl文件，选择合适的版本安装即可。测试代码如下

import os
import sys
import cv2
import copy
import torch
import numpy as np
import time
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
# import trt_common
from PIL import Image
import glob
 
os.environ['CUDA_VISIBLE_DEVICES']='2'
TRT_LOGGER = trt.Logger()
 
# EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
EXPLICIT_BATCH = 1
 
if sys.getdefaultencoding() != 'utf-8':
    # reload(sys)
    sys.setdefaultencoding('utf-8')
 
# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem
    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
    def __repr__(self):
        return self.__str__()
 
def get_engine(engine_file_path):
    with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())
 
# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream
 
# This function is generalized for multiple inputs/outputs for full dimension networks.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference_v2(context, bindings, inputs, outputs, stream):
    # Transfer input data to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]
 
if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES']='2'
    trt_name = '/xxx/model_fp16.engine'
 
    engine = get_engine(trt_name)
    context = engine.create_execution_context()
    inputs, outputs, bindings, stream = allocate_buffers(engine)
 
    data_dir = "/xxx/"
    img_type = 'png'
    data_ir = sorted(glob.glob(data_dir + "/IR/*." + img_type), key=lambda name: int(name[49:-5]))
    data_vi = sorted(glob.glob(data_dir + "/VI/*." + img_type), key=lambda name: int(name[49:-5]))
    
    time_list = []
    for i in range(len(data_ir)):
    
        ir = cv2.imread(data_ir[i])
        vi = cv2.imread(data_vi[i])
    
        ir = cv2.cvtColor(ir, cv2.COLOR_BGR2RGB)
        vi = cv2.cvtColor(vi, cv2.COLOR_BGR2RGB)
    
        ir = np.transpose(ir, (2, 0, 1)).astype(np.float32)
        vi = np.transpose(vi, (2, 0, 1)).astype(np.float32)
    
        ir = np.expand_dims(ir, axis=0)
        vi = np.expand_dims(vi, axis=0)
    
        ir /= 255.0
        vi /= 255.0
        ir = (ir - 0.5) / 0.5
        vi = (vi - 0.5) / 0.5
        
        stacked = [ir, vi]
        inputs[0].host = np.ascontiguousarray(stacked)
        
        t1 = time.time()
        trt_outputs = do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
        t2 = time.time()
        print(" trt: ", t2-t1)
        
        gen_Fused = copy.deepcopy(trt_outputs[0])
        gen_Fused = gen_Fused.reshape((3, 480, 640))
        img = (((gen_Fused - gen_Fused.min()) * 255) / (gen_Fused.max() - gen_Fused.min())).transpose(1, 2, 0).clip(0, 255).astype(np.uint8)
        img = Image.fromarray(img)
        img.save('./results_fp16/' + data_ir[i][49:-4] + '.png')

前面的函数是官方给的模板，我们要管的只有main里面的内容，把输入数据用numpy数组构建好，调用do_inference_v2函数得到输出，再自己按照需求解码成图像。这里的主要任务是对比Torch和TensorRT模型生成的结果是否一致，评估精度丢失。如果这里没问题了，就可以把onnx模型拷贝到NX开始部署了。

4.2.3.2 engine的构建

首先明确，在NX上是可以跑python的，部署中遇到问题，查报错信息又查不到的时候，可以使用python一步步排查。比如，用C++跑engine不正常，那就可以用python版的tensorrt跑一下，如果跑通了，就说明是C++代码写的不对。如果python版的tensorrt跑engine跑不通，那就说明engine转化失败，退一步检查onnx模型是否正常，用python写一个onnruntime的代码检查一下。如果onnx也跑不通，还可以继续往回退，在NX上装个torch跑一下原模型试试。

在NX上构建engine有很多种方法，从简单到复杂，分别是：(1) 基于onnx，用polygraphy命令行工具转换; (2) 基于onnx，调用TensorRT官方API，写C++或python代码转换; (3) 从零开始，调用TensorRT官方API，一层层的手动复现网络。其中第一种非常容易，不出意外的话一行命令就能搞定；第二种稍微有点麻烦，但有官方的example可以抄；第三种非常麻烦，但TensorRTx里头部署网络都是这么干的。我一个个来解释。

(1) 基于onnx，用polygraphy命令行工具转换

这里需要用到的是英伟达开发的一个新工具polygraphy，和trtexec命令行工具类似，可以一行命令启动转换。那为什么不像在服务器上那样用trtexec，一定要换一个工具呢？是因为在服务器上装TensorRT时会给你安装trtexec这个工具，但NX上没有附带这个工具。

安装方法可以在github的readme里找到，直接pip install就行

python -m pip install colored polygraphy --extra-index-url https://pypi.ngc.nvidia.com

因为我们要用的是polygraphy的命令行工具，所以只关注它的CLI怎么使用即可，教程在此。polygraphy内置了几个有用的子功能，通过不同的关键词来调用。polygraphy inspect 可以列出模型各个层的信息。polygraphy convert用于onnx格式到TensorRT格式的转换。polygraphy check用于检查一个onnx模型是否有无法运行的层。polygraphy surgeon用于转化失败时的诊断。polygraphy run可以运行多种格式的模型并提供评估。

我们主要要用的功能就是polygraphy convert，在convert的教程中给出了使用方法：

# 转化为fp16
polygraphy convert --fp-to-fp16 -o identity_fp16.onnx identity.onnx
# 转化为int8  需要编写data_loader.py，提供读入calibration数据的函数
polygraphy convert identity.onnx --int8 --data-loader-script ./data_loader.py --calibration-cache identity_calib.cache -o identity.engine

如果需要量化为INT8精度，则需要自己编写data_loader.py中读入calibration数据的函数，教程中也有示例，使用时把“x”改成转onnx时指定的变量名，把np.ones改成真实数据组成的numpy数组即可

import numpy as np
 
INPUT_SHAPE = (1, 1, 2, 2)
 
def load_data():
    for _ in range(5):
        yield {
            "x": np.ones(shape=INPUT_SHAPE, dtype=np.float32)
        }

转换过程会比较慢，但最终没报错并出现新的engine文件即为成功。

(2) 基于onnx，调用TensorRT官方API，写C++或python代码转换

TensorRT的API有python版本的和C++版本的，用哪种都可以。

用python调用TensorRT的API需要import tensorrt，这里如果你用了conda虚拟环境，则可能会出现import不成功的问题。这是因为jetpack装机的时候会把tensorrt装在系统的python环境里，我们需要手动链接到虚拟环境内。

sudo ln -s /usr/lib/python3.8/dist-packages/tensorrt* /home/xxx/miniconda3/envs/xxx/lib/python3.8/site-packages/

python的模板就看TensorRT官方的，以efficientnet为例，核心代码如下，思路就是先调用OnnxParser这个函数把onnx模型加载进来，然后通过config设置一系列的转换参数，最后调用builder.build_serialized_network，执行模型的转换。如果要量化为INT8，则还要额外对calibrator进行设置，这里就不提了

self.network = self.builder.create_network(0)
self.parser = trt.OnnxParser(self.network, self.trt_logger)
 
onnx_path = os.path.realpath(onnx_path)
with open(onnx_path, "rb") as f:
    if not self.parser.parse(f.read()):
        log.error("Failed to load ONNX file: {}".format(onnx_path))
 
......
 
self.config.set_flag(trt.BuilderFlag.FP16)
......
 
engine_bytes = self.builder.build_serialized_network(self.network, self.config)

如果要用C++做，这里也给出核心代码

IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
 
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = build_engine(builder, config, onnx_name);
assert(engine != nullptr);
 
// Serialize the engine
IHostMemory* modelStream{nullptr};
modelStream = engine->serialize();
assert(modelStream != nullptr);
 
// release resources
SAFE_DELETE(engine);
SAFE_DELETE(builder);
SAFE_DELETE(config);
 
std::ofstream p(engine_name, std::ios::binary);
if (!p) 
{
    std::cerr << "could not open plan output file" << std::endl;
    return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
SAFE_DELETE(modelStream);
std::cout << "engine created." << std::endl;

上面这一段创建builder和config，调用build_engine()创建engine，然后serialze()成二进制文件，最后写磁盘。核心的build_engine()函数内容如下

ICudaEngine* build_engine(IBuilder* builder, IBuilderConfig* config, std::string& onnx_name) 
{
    // create builder, explicit batch mode
    INetworkDefinition* network = builder->createNetworkV2(1U);
 
    auto parser = nvonnxparser::createParser(*network, gLogger);
 
	if (!parser->parseFromFile(onnx_name.c_str(), 0))
	{
        std::cout << "parse error!" << std::endl;
		return nullptr;
	}
 
    // builder settings 
    builder->setMaxBatchSize(1);
    config->setMaxWorkspaceSize(1ULL << 32); // 4GB
 
    // using FP16 precision
    config->setFlag(BuilderFlag::kFP16);
 
    // build engine
    std::cout << "Building engine, please wait for a while..." << std::endl;
    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    std::cout << "Build engine successfully!" << std::endl;
 
    // release resources
    SAFE_DELETE(network);
 
    return engine;
}

可以看到，核心流程都是一样的，先初始化builder和config，然后用onnxparser把模型读入，设置config后buildEngineWithConfig转化模型。同样的，如果需要量化为INT8，则要额外对calibrator进行设置，还涉及calibrator加载数据的代码，可以参考TensorRTx的实现。

(3) 从零开始，调用TensorRT官方API，一层层的手动复现网络

这条路就比较麻烦了，TensorRTx的代码都是这么干的，是很好的学习资料。有空可以找个简单网络看看它代码怎么写的，比如alexnet。

// 定义network
INetworkDefinition* network = builder->createNetworkV2(0U);
// 在网络入口处addInput，为输入数据绑定INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
assert(data);
// 从wts文件里获取每一层的权重，存到map里
std::map<std::string, Weights> weightMap = loadWeights("../alexnet.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
// 实现一个Conv层
IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 64, DimsHW{11, 11}, weightMap["features.0.weight"], weightMap["features.0.bias"]);
assert(conv1);
// 如果这层有参数，就要如下设置
conv1->setStrideNd(DimsHW{4, 4});
conv1->setPaddingNd(DimsHW{2, 2});
// 实现一个ReLU层
IActivationLayer* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
assert(relu1);
 
...... // 一层层地实现
 
// 在网络出口标记输出数据，绑定OUTPUT_BLOB_NAME
fc3->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*fc3->getOutput(0));
 
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(1 << 20);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

如果你自定义的网络里存在TensorRT原本不支持的算子，那么就只能用这种方式构建网络。过程中还需要自己写插件，用CUDA实现特定计算。可以参考TensorRTx里yolov8的plugin。

4.2.3.3 engine的使用

得到没有bug的engine之后，也有很多种方法使用它。比如用python的TensorRT库跑，或者用polyrgraphy run工具跑，这些都只适用于简单测试。这里着重介绍在项目里如何写C++来使用engine。

这里我通常把inference的功能封装成一个类，然后在类外只需要实例化这个类并调用接口即可。下面展示源码

PTET.h文件：

#ifndef __PTET_H__
#define __PTET_H__
 
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <chrono>
#include <jetson-utils/cudaUtility.h>
#include <jetson-utils/cudaMappedMemory.h>
#include "NvInfer.h"
 
using namespace nvinfer1;
using namespace std;
 
class PTET
{
public:
 
    PTET(string engine_path);
    ~PTET();
    bool Init();
    bool doInference(void* ir, void* vi, void* fused_image);
 
    class Logger : public ILogger           
    {
        void log(Severity severity, AsciiChar const * msg) noexcept override
        {
            // suppress info-level messages
            if (severity <= Severity::kWARNING)
                printf("[PTET] %s\n", msg);
        }
    } gLogger;
 
    string engine_name;
    const int INPUT_H = 480;
    const int INPUT_W = 640;
    const int INPUT_C = 3;
 
    const char* INPUT_NAME = "Imgs";
    const char* OUTPUT_NAME = "Fused";
 
    nvinfer1::IRuntime* runtime = nullptr;
    nvinfer1::ICudaEngine* engine = nullptr;
    nvinfer1::IExecutionContext* context = nullptr;
    cudaStream_t stream;
    
    float* buffers[2];
 
};
 
#endif

还记得onnx转换时指定了input_name和output_name吗，这里就需要用到这两个名字。这里对外开放一个函数doInference()，方便调用。

PTET.cpp文件：

#include "PTET.h"
#include <jetson-utils/cudaResize.h>
#include "myCudaConvert.h"
 
#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr)\
{\
    cudaError_t error_code = callstr;\
    if (error_code != cudaSuccess) {\
        std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\
    }\
}
#endif  // CUDA_CHECK
 
PTET::PTET(string engine_path) : engine_name(engine_path)
{
    if(!Init()) 
    {
        printf("[PTET] Init failed.\n");
    }
}
 
PTET::~PTET()
{
    SAFE_DELETE(context);
    SAFE_DELETE(engine);
    SAFE_DELETE(runtime);
    cudaStreamDestroy(stream);
 
    cudaFreeHost(buffers[0]);
    cudaFreeHost(buffers[1]);
}
 
bool PTET::Init()
{
    printf("[PTET] Loading image fusion engine ...\n");
    std::ifstream file(engine_name, std::ios::binary);
    if (!file.good()) 
    {
        printf("[PTET] Read engine file failed.\n");
        return false;
    }
    char *trtModelStream = nullptr;
    size_t size = 0;
    file.seekg(0, file.end);
    size = file.tellg();
    file.seekg(0, file.beg);
    trtModelStream = new char[size];
    file.read(trtModelStream, size);
    file.close();
 
    runtime = createInferRuntime(gLogger);
    if(!runtime)
    {
        printf("[PTET] Create runtime falied.\n");
        return false;
    }
    engine = runtime->deserializeCudaEngine(trtModelStream, size);
    if(!engine)
    {
        printf("[PTET] Deserialize engine falied.\n");
        return false;
    }
    context = engine->createExecutionContext();
    if(!context)
    {
        printf("[PTET] Create context falied.\n");
        return false;
    }
    delete[] trtModelStream;
 
    cudaStreamCreate(&stream);
 
    cudaMalloc((void**)&buffers[0], INPUT_W * INPUT_H * sizeof(float3) * 2);
    cudaMalloc((void**)&buffers[1], INPUT_W * INPUT_H * sizeof(float3));
 
    return true;
}
 
bool PTET::doInference(void* ir, void* vi, void* fused_image)
{
    cudaPacked2Planner((float3*) ir, INPUT_W, INPUT_H, (float*)buffers[0]);
    cudaPacked2Planner((float3*) vi, INPUT_W, INPUT_H, (float*)buffers[0] + INPUT_W * INPUT_H * sizeof(float3));
 
    auto start = std::chrono::high_resolution_clock::now();
 
    context->setTensorAddress(INPUT_NAME, buffers[0]);
    context->setTensorAddress(OUTPUT_NAME, buffers[1]);
    context->enqueueV3(stream);
    
    auto end = std::chrono::high_resolution_clock::now();
    int time_elapse = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    std::cout << time_elapse << "ms" << std::endl;
    
    cudaStreamSynchronize(stream);
 
    cudaPlanner2Packed((float*)buffers[1], INPUT_W, INPUT_H, (float3*)fused_image);
    cudaStreamSynchronize(stream);
}
 
/*  inference阶段：
    std::ifstream file(engine_name, std::ios::binary);
    ... 把file中数据流读入trtModelStream ...
    IRuntime* runtime = createInferRuntime(gLogger); // 创建runtime
    ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); // 对.engine文件读出的流反序列化生成engine
    IExecutionContext* context = engine->createExecutionContext(); // 从engine得到context
    float* buffers[2]; // 网络输入输出参数的缓冲区列表
    cudaMalloc((void**)&buffers[inputIndex], INPUT_SIZE); // 为输入/输出数据分配空间
    cudaMalloc((void**)&buffers[outputIndex], OUTPUT_SIZE);
    ... 准备输入数据，存入buffer[inputIndex] ...
    context.enqueue(batchSize, buffers, stream, nullptr); // 调用enqueue执行推理， 输入buffers[inputIndex]， 获得buffers[outputIndex]
    ... 结果拷回cpu ...
    */

这里的流程比较固定，一步步来就行。首先读入engine文件，反序列化，得到执行模型所需的context指针。在准备数据之前，需要先开辟固定大小的共享内存。执行模型前，首先要将输入数据拷到共享内存缓冲区buffer[0]里。调用context->enqueueV3(stream);即执行模型，调用cudaStreamSynchronize(stream);做一个同步，等待GPU执行完毕。输出结果存放在你指定的缓冲区buffer[1]里。最后把结果拷贝出来使用即可。这里涉及的两个格式转换函cudaPacked2Planner()和cudaPlanner2Packed()是我自己写的CUDA函数，将在4.4节进行介绍。

PTET_demo.cpp文件：

#include "PTET.h"
#include <jetson-utils/imageIO.h>
#include <jetson-utils/cudaResize.h>
#include <jetson-utils/glDisplay.h>
#include "myCudaConvert.h"
#include <string>
#include <vector>
#include <chrono>
 
using namespace std;
 
int main()
{
    const char* ir_path = "../../build_network/calibration/IR/00004N.png";
    const char* vi_path = "../../build_network/calibration/VI/00004N.png";
    string engine_path = "../../build_network/tokenfusion_fp16.engine";
    const char* save_path = "../result_fp16.png";
 
    void* ir = NULL;
    void* vi = NULL;
    int width = 0;
    int height = 0;
    PTET* ptet = new PTET(engine_path);
 
    void* fused_image = NULL;
    cudaMalloc((void**)&fused_image, ptet->INPUT_W * ptet->INPUT_H * sizeof(float3));
 
    void* ir_resized = NULL;
    void* vi_resized = NULL;
    cudaMalloc((void**)&ir_resized, ptet->INPUT_W * ptet->INPUT_H * sizeof(float3));
    cudaMalloc((void**)&vi_resized, ptet->INPUT_W * ptet->INPUT_H * sizeof(float3));
 
    for(int i = 0; i < 20; i++)
    {
        loadImage(ir_path, &ir, &width, &height, IMAGE_RGB32F);
        loadImage(vi_path, &vi, &width, &height, IMAGE_RGB32F);
        cudaResize((float3*) ir, width, height, (float3*) ir_resized, ptet->INPUT_W, ptet->INPUT_H);
        cudaResize((float3*) vi, width, height, (float3*) vi_resized, ptet->INPUT_W, ptet->INPUT_H);
        // auto start = std::chrono::high_resolution_clock::now();
        ptet->doInference(ir_resized, vi_resized, fused_image);
        // auto end = std::chrono::high_resolution_clock::now();
        // int time_elapse = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
        // std::cout << time_elapse << "ms" << std::endl;
    }
 
    saveImage(save_path, (float3*)fused_image, ptet->INPUT_W, ptet->INPUT_H, 100, make_float2(-1.0, 1.0));
 
    cudaFreeHost(fused_image);
    cudaFreeHost(ir_resized);
    cudaFreeHost(vi_resized);
 
    SAFE_DELETE(ptet);
 
    return 0;
}

在外部测试所写类的功能，调用jetson-uitls提供的loadImage加载图像，最后将结果存储在本地。

CMakeLists.txt文件：

cmake_minimum_required(VERSION 2.8)
project( ptet_demo )
 
find_package(CUDA REQUIRED)
find_package(OpenCV REQUIRED)
find_library(/xxx/jetson-inference/build/aarch64/lib NAMES jetson-utils REQUIRED)
 
link_directories(/usr/local/cuda/lib64)
link_directories(/usr/lib/aarch64-linux-gnu)
cuda_add_library(l_ptet PTET.cpp myCudaConvert.cu)
cuda_add_executable(ptet_demo PTET_demo.cpp)
target_link_libraries(ptet_demo l_ptet jetson-utils nvinfer cudart ${OpenCV_LIBS})

用cmake构建，不多赘述。

4.3 OpenCV的使用

这块其实没什么内容要写，需要什么功能就去opencv doc里寻找相关功能的C++函数即可。这里指出几个在jetson设备上使用OpenCV和其他设备上不一样的地方。

其一是OpenCV在compile with CUDA以后，可以使用它的GpuMat类。但是由于大多数时候在jetson设备上开辟内存都是用cudaAllocMapped函数去开辟CPU-GPU之间的共享内存，所以我自己开发中的感受是使用GpuMat类对算法加速的帮助不大，也可能是我太菜了没用对。

其二是如果用cudaAllocMapped去开辟CPU-GPU之间的共享内存，那么是可以直接用共享内存的指针去初始化一个看起来在cpu上的Mat的，这个很常用，我举个具体的例子

void* frame_IR;
cudaAllocMapped(&frameIR, cameraIR_W*cameraIR_H*sizeof(uchar3));
 
cv::Mat imIR;
imIR = cv::Mat(cameraIR_H, cameraIR_W, CV_8UC3, frameIR);

这里调用的cv:Mat构造函数如下，可以在doc里找到这个重载

Mat (int rows, int cols, int type, void *data, size_t step=AUTO_STEP)

然后我void* data传的是frame_IR，其指向一块共享内存。这样做是不会有什么问题的，如果你担心两个指针指向同一份数据导致数据的一致性出问题，则可以使用

cv::Mat imIR;
imIR = cv::Mat(cameraIR_H, cameraIR_W, CV_8UC3, frameIR).clone();

这样就避免了使用cudaMemcpy函数拷贝数据，写的时候会更方便一点。

4.4 CUDA函数的编写

在jetson设备上的图像算法开发流程基本是这样：需要一个功能->在jetson-utils里找有没有现成的函数->没有则用OpenCV实现。但是有时数据已经在GPU上了，该操作的前一步和后一步都是借助jetson-utils的函数通过CUDA实现的，这时候中间插入一个OpenCV，数据就要拷贝来拷贝去，降低运行效率。最好的办法就是自己写一个CUDA函数。

开始之前推荐阅读这篇文章，了解CUDA架构的基本知识，知道host、device、thread、grid、block的概念。然后去jetson-utils/cuda/下找几个.h .cu文件读一下看看别人怎么写的。

这里拿我曾经写过的一个CUDA函数为例讲解一下，在这个例子中，我想实现的功能是改变图像数据的存储方式。对于一幅RGB图像来说，每个像素点都具备R,G,B三个数据，但在内存里只能线性地存储，这就出现了两种存储方式。packed方式采用(R1G1B1),(R2G2B2),...,(RnGnBn)的顺序，以一个像素地内容作为一个包，写完它的RGB三个值以后才移动到下一个像素。而planner方式采用(R1R2...Rn),(G1G2...Gn),(B1B2...Bn)的顺序，把某一个颜色通道的数据全写完再进入下一个颜色通道。所以需要进行两种存储方式的转换。由于我正好还需要做归一化，所以把归一化操作也合并在这个函数里。

.h文件中：

#include <jetson-utils/cudaUtility.h>
 
// Normalize IMAGE_RGB32F(0, 255) to (0.0f, 1.0f)
// Then change the data form Packed mode to Planner mode (rgbrgbrgbrgb --> rrrrggggbbbb)
cudaError_t cudaPacked2Planner(float3* input, size_t width, size_t height, float* output);

按照模板，声明一个返回值类型为cudaError_t的函数，数据以指针形式传入，也以指针形式传出。

.cu文件中：

__global__ void gpuPacked2Planner(float3* input, size_t width, size_t height, float* output)
{
	const int x = blockIdx.x * blockDim.x + threadIdx.x;
	const int y = blockIdx.y * blockDim.y + threadIdx.y;
 
	if( x >= width || y >= height )
		return;
 
    const int pixel = y * width + x;
 
    output[pixel] = (float)input[pixel].x / 255.0f; // red
    output[pixel + width * height] = (float)input[pixel].y / 255.0f; // green
    output[pixel + width * height * 2] = (float)input[pixel].z / 255.0f; // blue
}
 
cudaError_t cudaPacked2Planner(float3* input, size_t width, size_t height, float* output)
{
    if( !input || !output )
		return cudaErrorInvalidDevicePointer;
 
    if( width == 0 || height == 0)
		return cudaErrorInvalidValue;
 
    // launch kernel
	const dim3 blockDim(8, 8);
	const dim3 gridDim(iDivUp(width,blockDim.x), iDivUp(height,blockDim.y));
    
    gpuPacked2Planner<<<gridDim, blockDim>>>(input, width, height, output);
 
    return CUDA(cudaGetLastError());
}

下方的cudaPacked2Planner函数就是一个模板，blockDim(8, 8);的设置是我从jetson-utils里依葫芦画瓢抄来的。唯一需要改的部分就是核函数调用这一行gpuPacked2Planner<<<gridDim, blockDim>>>(input, width, height, output);

上面的gpuPacked2Planner()函数即为CUDA核函数，__global__指定了这个函数由CPU调用，GPU执行，算法的处理逻辑都写在核函数中。首先计算x和y，得到的x就是图像矩阵中的列号，y就是图像矩阵的行号。pixel = y * width + x，就得到了从0开始计算，像素位置的index。由于GPU是并行执行的，所以这里可以想象已经为pixel开了for循环来遍历它的每个位置。

output[pixel] = (float)input[pixel].x / 255.0f; // red
output[pixel + width * height] = (float)input[pixel].y / 255.0f; // green
output[pixel + width * height * 2] = (float)input[pixel].z / 255.0f; // blue

由于输入是rgbrgbrgb，输出是rrrgggbbb，所以原本的R分量需要放在[0,width*height)的范围内，原本的G分量需要放在[width*height, width*height*2)的范围内，原本的B分量需要放在[width*height*2, width*height*3)的范围内。顺便做一个归一化，数值除以255.0。至此就完成了这个CUDA核函数的编写。

代码的构建仍然是使用cmake，CMakeLists.txt里可能会用到的两行如下：

cuda_add_library(xxx xxx.cu)
target_link_libraries(xxx nvinfer cudart)

4.5 其他

包含两个小知识点。

第一，系统右上角状态栏可以选板子的功率，最高25W。提高功率可以开启更多的CPU一起工作。

第二，sudo jetson_clocks命令可以使处理器暂时运行在最高性能。

5. 手册&文档

Archived Documentation For Jetson Software

Jetson 下载中心 | NVIDIA 开发者

Jetson Linux Developer Guide

TensorRT Documentation

TensorRT: Class List

TensorRT Operators

PyTorch for Jetson - Announcements - NVIDIA Developer Forums

OpenCV doc

Jetson Zoo - eLinux.org

Jetson Orin NX教程 ntgy

C++参考手册