【详细完整部署流程】在Jetson NX上使用PyTorch转TensorRT搭建并部署自己的模型

概述

本文内容包括：在主机上使用PyTorch搭建网络，使用torch.onnx导出ONNX模型，上传到Jetson NX开发板上后使用trtexec将ONNX模型转为TensorRT模型，再通过C++ TensorRT实现模型推理。本文推理代码参考YOLOX的tools/export_onnx.py，模型参考某单目测距模型（UDepth）。

一、软硬件需求

硬件版本

Jetson Xavier NX官方开发板（后称Jetson NX）

软件版本

主机

Ubuntu 20.04
PyTorch：1.11.0/cu115
CUDA：11.6
ONNX：1.13.0
opencv-python：4.7.0（模型转换不需要cv2）

Jetson

CUDA：11.4
TensorRT：8.4.1.5 （重要）
OpenCV：4.5.4（不支持CUDA加速）
CMake等C++开发编译环境：正常版本即可，问题不大

二、PyTorch模型转ONNX格式

模型概况

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image

from model.udepth import UDepth
from model.guided_filter import FastGuidedFilter, GuidedFilter, ConvGuidedFilter
from CPD.CPD_ResNet_models import CPD_ResNet

class End2EndUDepth(nn.Module):
    """Combines the UDepth backbone and affiliated networks"""
    def __init__(self, udepth_model_path, cpd_model_path, mode='eval'):
        super(End2EndUDepth, self).__init__()
        self.mode = mode

        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        self.udepth_backbone = UDepth.build(n_bins=80, min_val=0.001, max_val=1, norm="linear")
        
        self.load_model(self.udepth_backbone, udepth_model_path)

        self.mask_net = CPD_ResNet()

        self.load_model(self.mask_net, cpd_model_path)

        self.transform = transforms.Compose([
                transforms.Resize((352, 352)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        self.guided_filter = GuidedFilter(r=4, eps=0.2)

    def forward(self, img):
        """Inputs should be type of torch.tensor"""
        # Get SOD mask, default shape = (240,320).
        mask_inp = self.transform(
            # transforms.ToPILImage()(img)
            transforms.ToPILImage()(img.view(img.shape[1:]))
            ).unsqueeze(0).to(self.device)
        _, mask = self.mask_net(mask_inp)
        mask = F.interpolate(mask, size=(240,320), mode='bilinear', align_corners=False)
        mask = mask.sigmoid()

        # UDepth backbone
        x = img.div(255.0)
        x = F.interpolate(x, size=(480,640), mode='bilinear', align_corners=False)
        _, depth = self.udepth_backbone(x)

        # GuidedFilter
        depth = depth / torch.max(depth) # normalization
        out = self.guided_filter(mask, depth) * 255 # different from cv2.GuidedFilter
        out = F.interpolate(out, size=(img.size()[2], img.size()[3]), mode='bilinear', align_corners=False)

        # # Without GuidedFilter
        # depth = depth / torch.max(depth) # normalization
        # out = depth * 255
        # out = F.interpolate(out, size=(img.size()[2], img.size()[3]), mode='bilinear', align_corners=False)

        # avoid using squeeze() because it brings if-else in onnx model and causes error in TensorRT
        out = out.view(out.shape[2], out.shape[3]).contiguous()

        return out.to(torch.float32) # OK

    def load_model(self, model, model_path):
        if model_path is not None:
            model.load_state_dict(torch.load(model_path))
        model.to(self.device)

        if self.mode == 'train':
            model.train()
        else:
            model.eval()
      

以上为模型定义。作者以原模型推理流程为样本，将骨干网络和CPD_ResNet（用于显著目标检测的级联部分解码器，CVPR 2019）结合，并使用一种快速可训练引导滤波模块（CVPR 2018）的PyTorch实现替换原推理流程中的cv2.ximgproc.GuidedFilter，避免TensorRT推理后还需要调用不支持CUDA的OpenCV。此模型只能用来推理，训练请参考原模型及论文。

模型转换并部署到平台，实际上会产生constant fold即常量折叠，原有的.pth的模型文件通过model.load_state_dict(torch.load(model_path))调用将模型参数读取并存储在torch.nn.Module模型中，当导出为ONNX文件时，这些模型参数就会随着模型结构被合并到输出的文件里。

需要关注的地方在于：

1. 输入维度：作者曾调试了整整一个星期才发现，如果使用torch.onnx导出中间模型，再使用trtexec转换模型，输入维度必须为(batch_size, channels, height, width)，否则会出现以下诡异的输出：
   
   这个结果作者浪费了大量时间才找到问题所在。这个图片是在水边拍的，输入通道如果为RGB，红通道显然亮度值会偏低，第一排图像偏暗，符合直觉。感兴趣的读者可以自己推导一下：在输入维度为(1080,1920,3)时推理结果为什么会出现这样的情况（正确输入维度应该为(1,3,1080,1920)）。
2. 输出维度：按需求即可，本文中是(1080,1920)的单通道深度图。
3. 一些不能使用的函数及推荐使用的函数：首先推荐一篇知乎文章Pytorch转onnx, TensorRT踩坑实录 - 小葱蘸大酱的文章 - 知乎，对不能使用的函数讲的比较细。从实践中来看，作者遇到了以下会出问题的函数：
   • torch.nn.functional.pad：上面的文章中指出，该函数不被TensorRT支持

     [E] Error[4]: [shuffleNode.cpp::symbolicExecute::387] Error Code 4: Internal Error (Reshape_ 68: IShuffleLayer applied to shape tensor must have 0 or1 reshape dimensions: dimensions were [-1,2]) 
     

     可以尝试使用其他函数替代。
   • squeeze：很常用的函数，用来压缩输出维度非常好用，不知道为什么TensorRT竟然不适配这个函数。出问题的原因在于squeeze会在ONNX模型中引入If层导致模型可能会输出动态维度，torch转TensorRT时squeeze中的If问题 - JasonZhu的文章这篇文章讲的很好。这里提出一个很方便的解决方法，即使用view替换squeeze。对于(1, 1, 1080, 1920)的深度图输出（深度估计模型）

     # avoid using squeeze() because it brings if-else in onnx model and causes error in TensorRT
     out = out.view(out.shape[2], out.shape[3]).contiguous()
     

     等同于

     out = out.squeeze(0).squeeze(0).contiguous()
     

     印象中是否调用contiguous并不影响结果。
   • 小结：尽量用一些常用，常见，基础的函数！！
     在前文提到的“快速可训练引导滤波模块”的原始实现中，

     # Original implementation
     N = self.boxfilter(Variable(x.data.new().resize_((1, 1, h_x, w_x)).fill_(1.0)))
     
     # My implementation
     N = self.boxfilter(x.clone().view(1, 1, h_x, w_x).contiguous().fill_(1.0))
     

     作者使用的是一个比较老的模块，nn.autograd.Variable，现在官网连原始文档都没了。
     UDepth的原始代码中也使用了nn.functional.upsample这个被弃用的函数，本文用interpolate替换掉了。
     作者在此还是建议养成在搭建模型的时候尽量用一些没有warning的，基础的接口的习惯，以免出现乱七八糟难以解决的问题。

按照上述原则，读者可以用PyTorch自行搭建和训练自己的模型。

模型转换

在此直接提供PyTorch模型到ONNX的转换代码，是从YOLOX的tools/export_onnx.py里抽出来修改的，很好用。

from loguru import logger

import torch
from torch import nn
import numpy as np

from model import UDepth
from model.end2end_udepth import End2EndUDepth

def main():
    model_path = 'Your_model.pth'
    cpd_path = 'Your_model.pth'

    onnx_file = 'output_onnx_model.onnx'

    end2end_model = End2EndUDepth(model_path, cpd_path) # Build model for eval

    logger.info("loading checkpoint done.")

    dummy_input = torch.randn(1, 3, 1080, 1920).cuda() # UDepth input size

    logger.info("Input size as: ({},{},{},{})".format(dummy_input.size(0),
                                                      dummy_input.size(1),
                                                      dummy_input.size(2),
                                                      dummy_input.size(3)))

    torch.onnx.export(
        end2end_model,
        dummy_input,
        onnx_file,
        input_names=["End2End-input"],
        output_names=["End2End-output"],
        opset_version=11
    )

    logger.info("generated onnx model named {}".format(onnx_file))

if __name__ == "__main__":
    main()

这里需要修改的内容有：

1. 模型end2end_model改成要转换的模型；
2. dummy_input改成你的输入维度，一定要是(1,c,h,w)的形式；
3. onnx_file输出文件；
4. input_names和output_names以列表形式给出，要记下来，部署要用；
5. opset_version这个参数是有用的、有区别的，可以先填个11，不太高不太低，出了问题再说。

模型简化（可选）

如果装了onnx-simplifier（作者是0.4.10），可以运行

python3 -m input_onnx.onnx simp_onnx.onnx

来简化模型，本文通过简化，模型大小缩小到原来的1/4，单次推理的运行速度也从原来的260ms降低到145ms。

三、ONNX（.onnx）转TensorRT模型（.trt）

转换方式

1. 首先明确，.trt或者.engine在实践中是一样的，都是序列化的模型文件，且该序列化模型文件和硬件平台绑定，也就是说如果想要在Jetson NX上部署某种模型，就必须要在该平台上进行ONNX到TensorRT模型的转换，不能在主机上单独下一个TensorRT，转换为.trt再导入到部署平台上，该流程也可以理解为某种意义上的本地编译。
2. 如果只是想在主机上部署TensorRT加速，这里提供一个流程：
   • 按照Nvidia的官方教程，通过tar包安装；
   • 下载某个版本的TensorRT包（需要注册Nvidia Developer的账号），鼠标移到链接上看看浏览器右下角的链接，看看自己下的到底是哪个版本的包；
   • 主机平台上解压到某个文件夹，按照官方教程所示安装几个whl包即可；
   • cd到bin目录，目录下有个trtexec的可执行文件，接下来参考下文Jetson NX部署的部分。
3. 在Jetson NX平台上部署TensorRT加速，TensorRT应该已经随Nv的Ubuntu预装到/usr/src/tensorrt/bin目录下，cd到该目录，执行

   ./trtexec --onnx=abs_path_to_onnx.onnx --saveEngine=abs_path_to_trt.trt
   

   即可调用TensorRT默认的转换器完成ONNX模型到TensorRT模型的转换。详细的参数可以--help查看，作者尝试修改过一些参数，实际上只要你的推理流程是对的，模型是正确的，并不需要调一些什么workspace之类的参数，当然了，还是有两个参数需要关注一下：
   • --int8：INT8量化，和--fp16是互斥的；
   • --fp16：如果不设置该选项或者--int8，TensorRT默认是采用32位浮点推理，当然就会慢一些，但是也需要考虑模型输出的数据类型，如果设置为不兼容的数据类型，比如模型固定输出torch.float32但是设置一个fp16推理，也是会报错的。

四、通用的TensorRT推理代码（C++）

简介

这部分的代码是参考YOLOX的推理框架，做了一些修改，总的来说是一个非常简单朴素的推理流程，如果有一些后处理的需要，读者可以自己看情况添加。示例代码输入RGB图像，输出单通道深度图。
题外话：如果是主机上想用OpenCV做后处理，其实没必要全都装到根目录下，可能会有一些环境污染之类的问题，作者更喜欢隔离的比较好的环境。实际上只要下载好OpenCV源码，编译完成以后，在你的推理代码CMakeLists.txt里添加上OpenCV的链接库就行了，添加include路径和链接路径：

include_directories(xxx/OpenCV/opencv-4.5.4/include)
link_directories(xxx/OpenCV/opencv-4.5.4/build/lib)

添加可执行文件的动态链接库：

target_link_libraries(your_executable opencv_core opencv_video opencv_videoio opencv_imgcodecs)

具体要添加哪些.so可以通过报错来判断，cv::Mat这样的基础模块当然应该是在opencv_core.so里；VideoCapture之类的应该在包含有video的库文件中；fourcc编码在opencv_imgcodecs.so里。

推理代码

代码按序合并即可运行，以下分模块介绍。

include和宏定义

#include <dirent.h>
#include <chrono>
#include <fstream>
#include <iostream>
#include <numeric>
#include <opencv2/opencv.hpp>
#include <sstream>
#include <vector>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"

#define CHECK(status)                                    \
  do {                                                   \
    auto ret = (status);                                 \
    if (ret != 0) {                                      \
      std::cerr << "Cuda failure: " << ret << std::endl; \
      abort();                                           \
    }                                                    \
  } while (0)

#define DEVICE 0  // GPU id

using namespace nvinfer1;

const char* INPUT_BLOB_NAME = "End2End-input";
const char* OUTPUT_BLOB_NAME = "End2End-output";
const char* OUTPUT_VIDEO = "depth_estimate.MP4";

#define OUTPUT_TYPE float32_t  // nvinfer1::DataType::kHALF

static Logger gLogger;

1. 以上头文件包含CUDA库，C++标准库，OpenCV库和日志模块（代码太长就不提供了，应该可以避开，如果避不开可以在YOLOX的GitHub中找到该文件）等。
2. CHECK宏定义是TensorRT推理常见的一种写法，用来判断某次CUDA调用是否正常运行，用宏定义实现避免函数入栈操作，增强代码可读性。
3. INPUT_BLOB_NAME和OUTPUT_BLOB_NAME后面会会用到，这里替换成第二节的模型转换中提到的以Python list形式给出的内容。
4. OUTPUT_TYPE宏定义是为了便于做模型不同精度的量化，基本数据类型和nvinfer1::DataType的对应关系为：

   float32_t => nvinfer1::DataType::kFLOAT
   float16_t => nvinfer1::DataType::kHALF
   int8_t => nvinfer1::DataType::kINT8
   

   具体的输出数据怎么解释，可以在输出结果后再reinterpret_cast或者static_cast来进行转换。这里不建议用非跨平台的数据类型比如float等，跨平台下可能会有一些想不到的坑，用float32_t等比较直观。

推理函数

void doInference(IExecutionContext& context, float* input, uint32_t input_size,
                 OUTPUT_TYPE* output, uint32_t output_size) {
  const ICudaEngine& engine = context.getEngine();

  // Pointers to input and output device buffers to pass to engine.
  // Engine requires exactly IEngine::getNbBindings() number of buffers.
  assert(engine.getNbBindings() == 2);
  void* buffers[2];

  // In order to bind the buffers, we need to know the names of the input and output tensors.
  // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
  assert(engine.getBindingDataType(inputIndex) == nvinfer1::DataType::kFLOAT);
  const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
  // assert(engine.getBindingDataType(outputIndex) == nvinfer1::DataType::kINT8); // For int8
  // assert(engine.getBindingDataType(outputIndex) == nvinfer1::DataType::kHALF); // For fp16
  assert(engine.getBindingDataType(outputIndex) == nvinfer1::DataType::kFLOAT);  // For fp32
  // int mBatchSize = engine.getMaxBatchSize();

  // Create GPU buffers on device
  CHECK(cudaMalloc(&buffers[inputIndex],
                   input_size * sizeof(float)));  // 3-channel input
  CHECK(
      cudaMalloc(&buffers[outputIndex],
                 output_size * sizeof(OUTPUT_TYPE)));  // 1-channel depth output

  // Create stream
  cudaStream_t stream;
  CHECK(cudaStreamCreate(&stream));

  // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
  CHECK(cudaMemcpyAsync(buffers[inputIndex], input,
                        input_size * sizeof(float), cudaMemcpyHostToDevice,
                        stream));
  // context.enqueue(1, buffers, stream, nullptr);
  context.enqueueV2(buffers, stream, nullptr);
  CHECK(cudaMemcpyAsync(output, buffers[outputIndex],
                        output_size * sizeof(OUTPUT_TYPE),
                        cudaMemcpyDeviceToHost, stream));
  cudaStreamSynchronize(stream);

  // Release stream and buffers
  cudaStreamDestroy(stream);
  CHECK(cudaFree(buffers[inputIndex]));
  CHECK(cudaFree(buffers[outputIndex]));
}

本文提供的推理函数可以直接用，只要改动input和output的数据类型和对应的input_size、output_size即可。该推理函数将以input开头的大小为input_size的一维数组缓冲区数据作为输入，输出推理结果到以output开头的大小为output_size的一维数组缓冲区中。
这里简要介绍一下几个关键函数：

1. engine.getBindingDataType(Index)：engine即CUDA引擎对象指针，在主函数中通过反序列化.trt文件可以获取一个CUDA引擎对象；getBindingDataType方法返回绑定在数据位置Index上的数据类型，而该参数又由getBindingIndex方法获取，其中参数为上一小节include和宏定义中第三点的内容。
   举例来说，getBindingIndex("xxxInputName")应该返回对应输入所在的维度index，即0（后文主函数会提到一个和该函数相似的getBindingDimension()函数，都是获取TensorRT模型上的一些属性信息）；获得index == 0以后再调用getBindingDataType(index)，就会返回输入维度的数据类型。
2. cudaMalloc：CUDA设备上的内存（显存）分配，返回值是一个表示分配状态是否成功的int，指向内存的指针以第一个参数（void **）给出。
3. cudaMemcpyAsync：CUDA设备间的异步内存拷贝，该函数未必表现出完全的异步行为，参考CSDN博主Zhninu的文章、CUDA设备的固定内存以及官方对于CUDA设备异步行为的描述。该函数参数分别为：目的地址，源地址，拷贝的块大小，传输类型（从什么设备传输到什么设备），CUDA流标识符。 详细内容参考官方文档。
4. context.enqueueV2()：将一次推理入队到流上，对于简单的模型推理这一部分就不用改了，参考官方文档。
5. cudaStreamSynchronize(stream)：CUDA流同步，异步推理需要在此处阻塞直到流推理完成，此时输出缓冲区就可以获取到推理结果了。

预处理和后处理

void MemcpyCVImgToInputArray(const cv::Mat& mat, float* des) {
  auto channel = mat.channels();
  auto width = mat.cols;
  auto height = mat.rows;
  for (int c = 0; c < channel; c++) {
    for (int h = 0; h < height; h++) {
      for (int w = 0; w < width; w++) {
        des[c * width * height + h * width + w] =
            static_cast<float>(mat.at<cv::Vec3b>(h, w)[c]);
      }
    }
  }
}

cv::Mat DecodeOutput(OUTPUT_TYPE* dist, int img_w, int img_h) {
  cv::Mat mat(img_h, img_w, CV_8UC1);
  // mat.data = dist; // For uint8_t

  // Define c = 1
  for (int h = 0; h < img_h; ++h) {
    for (int w = 0; w < img_w; ++w) {
      mat.at<uint8_t>(h, w) = static_cast<uint8_t>(dist[h * img_w + w]);
    }
  }  // For float

  return mat;
}

1. 预处理，将uint8_t表示的像素值转为输入数据类型float的一维数组，即展开。
2. 后处理，CUDA推理的输入输出都是一维数组，得到输出以后同样地，将展开的数据拼起来。

主函数

int main(int argc, char** argv) {
  cudaSetDevice(DEVICE);
  // create a model using the API directly and serialize it to a stream
  char* trtModelStream{nullptr};
  size_t size{0};

  if (argc == 4 && std::string(argv[2]) == "-v") {
    const std::string engine_file_path{argv[1]};
    std::ifstream file(engine_file_path, std::ios::binary);
    if (file.good()) {
      file.seekg(0, file.end);
      size = file.tellg();
      file.seekg(0, file.beg);
      trtModelStream = new char[size];
      assert(trtModelStream);
      file.read(trtModelStream, size);
      file.close();
    }
  } else {
    std::cerr << "arguments not right!" << std::endl;
    return -1;
  }

  // Build CUDA inference pipeline
  IRuntime* runtime = createInferRuntime(gLogger);
  assert(runtime != nullptr);
  ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
  assert(engine != nullptr);
  IExecutionContext* context = engine->createExecutionContext();
  assert(context != nullptr);
  delete[] trtModelStream;

  // In UDepth, get(0) for input dimension, get(1) for output dimension
  // e.g. for (1080,1920,3) input, output dim is (1080,1920)
  auto out_dims = engine->getBindingDimensions(1);
  auto output_size = 1;
  for (int j = 0; j < out_dims.nbDims; j++) {
    output_size *= out_dims.d[j];
  }
  static OUTPUT_TYPE* dist = new OUTPUT_TYPE[output_size];

  cv::Mat img;
  auto video_cap = cv::VideoCapture(argv[3]);

  int channels = 3;
  int img_w = video_cap.get(cv::CAP_PROP_FRAME_WIDTH);
  int img_h = video_cap.get(cv::CAP_PROP_FRAME_HEIGHT);
  int input_size = img_w * img_h * channels;
  float* input_buffer = new float[input_size];

  int frame_count = video_cap.get(cv::CAP_PROP_FRAME_COUNT);
  int total_infer_time_in_ms = 0;

  int fourcc = video_cap.get(cv::CAP_PROP_FOURCC);

  auto video_writer =
      cv::VideoWriter(OUTPUT_VIDEO, fourcc, video_cap.get(cv::CAP_PROP_FPS),
                      cv::Size(img_w, img_h), false);

  while (video_cap.read(img)) {
    MemcpyCVImgToInputArray(img, input_buffer);

    // run inference
    auto start = std::chrono::system_clock::now();
    doInference(*context, input_buffer, input_size, dist, output_size);
    auto end = std::chrono::system_clock::now();
    std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end -
                                                                       start)
                     .count()
              << "ms" << std::endl;
    total_infer_time_in_ms +=
        std::chrono::duration_cast<std::chrono::milliseconds>(end - start)
            .count();

    auto frame = DecodeOutput(dist, img_w, img_h);
    cv::imwrite("/home/tc_nx/workdir/deploy/udepth_test.jpg", frame);
    video_writer.write(frame);
  }
  delete[] input_buffer;
  delete[] dist;

  video_writer.release();  // Called by deconstructor

  std::cout << "Average inference time: "
            << total_infer_time_in_ms * 1.0 / frame_count << "ms\n";

  // // destroy the engine, deprecated
  // context->destroy();
  // engine->destroy();
  // runtime->destroy();
  delete context;
  delete engine;
  delete runtime;

  return 0;
}

1. if (argc == 4 && std::string(argv[2]) == "-v")这块作者偷懒copy了，实际上做的工作就是用标准文件输入流读取序列化的.trt文件，判断这么多自己写的时候其实没必要。
2. 创建CUDA推理管道，包括CUDA运行时，CUDA引擎和执行上下文，内存里的引擎文件（也就是.trt）反序列化到ICudaEngine对象里以后就可以delete掉了。
3. engine->getBindingDimensions(1)会得到一个对象，对于简单的单输入单输出模型，0对应输入，1对应输出，对象的成员d是一个数组，d[0]就是输入/输出的第一维size，以此类推。
4. 其他流程不复杂，参见代码即可。

CMakeLists.txt

cmake_minimum_required(VERSION 2.6)

project(udepth)

add_definitions(-std=c++11)

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)

find_package(CUDA REQUIRED)

include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include) # Necessary
include_directories(/data/cuda/cuda-10.2/cuda/include)
link_directories(/data/cuda/cuda-10.2/cuda/lib64)
link_directories(/usr/local/cuda-11.4/targets/aarch64-linux/lib) # Necessary
# cudnn
include_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/include)
link_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/lib64)
# tensorrt
include_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/include)
link_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/lib)

find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})

# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -O0 -Wfatal-errors -fopenmp -D_MWAITXINTRIN_H_INCLUDED -g") # For debug
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -O2 -Wfatal-errors -fopenmp -D_MWAITXINTRIN_H_INCLUDED") # For release

add_executable(udepth ${PROJECT_SOURCE_DIR}/udepth.cpp)
target_link_libraries(udepth nvinfer)
target_link_libraries(udepth cudart)
target_link_libraries(udepth ${OpenCV_LIBS})

CMakeLists的部分代码来源于YOLOX，作者的环境里并没有/data/cuda的路径，但这些内容还是保留了。关键的是要添加CUDA的include路径和链接库，链接库位置应该在/usr/local/cuda-11.4/targets/aarch64-linux/lib类似的目录下，读者部署前需要自行检查。

总结

多看官方文档，包括PyTorch和NVIDIA的文档，如果想要扎实掌握部署的流程，这部分是绕不开的。