Depth AnythingTensorRT Jetson orin C++版

最新推荐文章于 2024-08-28 09:35:38 发布

CCChester

最新推荐文章于 2024-08-28 09:35:38 发布

阅读量1.1k

点赞数 25

文章标签： c++ 开发语言

本文链接：https://blog.csdn.net/ccchester/article/details/138010704

版权

最新的深度估计方法，源码是python和桌面级GPU运行，如果实现在嵌入式设备jetson Orin上用C++实时运行一定很有意思。

1、需要什么准备工作？

先下depth anything的代码：
https://github.com/LiheYoung/Depth-Anything.git
再下一个实现C++ tensorRT部署的范例代码–参考用，主要还是自己写一个：
https://github.com/spacewalk01/depth-anything-tensorrt.git。

2、获得onnx模型

要实现在C++上做深度模型推理，目前推荐的方法是先为onnx模型，再用C++调用：

2.1 先下载depth anything 的预训练模型：https://huggingface.co/spaces/LiheYoung/Depth-Anything/tree/main/checkpoints

2.2 配置depth anything 的环境，针对jetson orin平台应该有所不同

看看都需要安装啥
gradio_imageslider #这个不用变
gradio==4.14.0 #这个不用变 Requires: Python >=3.8 注意，安装的时候会出现torchvision 0.15.1需要 torch但其未安装的错误，但是不打紧马上就安装了，应该是自动安装了torchvision 0.15.1。
torch #这个安装jetson对应的版本 PyTorch for Jetson - Jetson & Embedded Systems / Jetson Nano - NVIDIA Developer Forums 我的是jet版本是5.1-b147 ，对应torchvision的torch版本是2.0.0。点击下面的红框便可自动下载。
在这里插入图片描述
下载完成后，在下载位置进入虚拟环境，执行 pip install torch包的名字.whl
torchvision #自动就装上了
opencv-python #这个不用变
huggingface_hub #这个不用变
onnx# 这个是新加的，用来导出onnx模型的。
开始安装：（假设anaconda已经安装到了jetson orin上）。1. 创建环境 conda create -n your_env_name python=3.8 , your_env_name 是你的环境名称，自己起一个吧，我的是 Depth。2. 按顺序pip install 安装吧

2.3 文件调整：

将depth anythingtensorRT中的 dpt.py 放到Depth Anything文件所在子文件depth anything中，替换原来的 dpt.py ；然后将export.py 复制到Depth Anything文件;depth anythingtensorRT的作者移除了dpt.py 中前向传播的最后的序列操作避免和tensorRT冲突；使用export.py 将模型导出为onnx，名字为depth_anything_vit{}14.onnx。执行命令为：
python export.py --encoder vitb --load_from depth_anything_vitb14.pth --image_shape 3 518 518
看一下export.py是怎么个操作：主要有两个函数， main和export_model。前者主要前者从命令行中获取必要的参数：如编码器，预训练模型的路径，图像的维度（三通道，高，宽）。注意，编码器和模型要对应上，vitb对应depth_anything_vitb14.pth。最后将参数都传给export_model函数；export_model返回一个ONNX模型：获得模型的名称后：如vitb，调用depth_anything的类DPT_DINOv2对模型进行初始化。total_params为获得所加载模型的参数大小；从预训练模型中加载模型的权重，并创建一个dummy输入获得范例输出（为什么要给一个dummy input等等，其实是这样的，因为我们知道pytorch模型和tf模型是不一样的，前者是动态计算图，后者是静态的，所以就会导致一个问题，对于pytorch模型，在运行之前是不知道各结点的状态和计算关系的，所以我们在转onnx的时候，需要实际地跑一下网络，得到每层和每个节点的计算图，这样才能转换模型，所以我们需要随便指定一个满足要求的输入即可。由于这种实际运行才能得到转换结果的情况）。指定onnx的路径 xxx.onnx。使用torch.onnx.export(网络模型，dummy输入，onnx路径，导出的onnx版本，模型输入名称，模型输出名称，verbose=true)，有一个问题，onnx版本跟torch有关，torch2.0.0版本的默认onnx是14。尝试方法是先用11版本生成一个，再用14版本的生成一个。
最后将opencv-dll文件拷贝到Depth Anything/build/Release文件夹下。

2.4 安装tensorRT

这个在jetson中已经安装好了的，但是安装位置在/usr/lib/python3.8/dist-packages/中，不能被虚拟环境中定位使用。因此我们需要软链接一下，运行如下命令：sudo ln -s /usr/lib/python3.8/dist-packages/tensorrt* /home/nvidia/anaconda3/envs/orin/lib/python3.8/site-packages/
测试一下，运行如下指令：

python -c "import tensorrt;print(tensorrt.__version__)"

若出现版本号，则成功

3，在C++环境中使用onnx进行推理

3.1修改Cmakelist文件。

假设各位已经初始化了一个功能包。
需要在CmakeList中添加OPENCV TENSORRT三个库的地址
jetson orin应该如何在Cmakelists中添加上述三个功能？
opencv 可以直接使用 find_package(OpenCV REQUIRED) ；TENSORRT需要先设置路径，但是一直显示找不到？先不管tensort。

3.2修改代码

修改主函数，包含对应的头文件 #include <NvInfer.h> 出现报错，找不到cuda_runtime.h 在CmakeList添加

find_package(CUDA REQUIRED)  include_directories(${CUDA_INCLUDE_DIRS})

自己写一个C++读取图片的历程，获得cv::Mat 形式的图像。
调用depth_model.predict进行图像推理即可

cv::Mat result_d = depth_model.predict(frame);

如何写depth_model的预测函数？
模仿depth anythingtensorRT创建两个文件。然后直接赋值粘贴即可。
同时参考其Cmakelist文件，将多个文件编译为一个可执行文件。
并且设置
set(TENSORRT_LIBS nvinfer nvinfer_plugin nvparsers nvonnxparser)

 target_link_libraries(${PROJECT_NAME}
    ${OpenCV_LIBS}
    ${CUDA_LIBRARIES}
    ${TENSORRT_LIBS}
)

编译，运行

3.3运行测试

使用10帧图像预热模型，出现错误：ONNX以INT64的权重生成，tensorrt不支持，转为INT32位的。这些警告都可以忽视。但是需要注意的是 32G内存的只能跑最小的模型，其他的模型会出现内存不足的问题。
代码已开源至github