tkDNN 使用与源码浏览

0. 前言

• 参考资料：

  • Github
  • tkDNN代码的使用和理解

• tkDNN 是什么

  • 简单说，就是利用 cuDNN 和 TensorRT 实现了一些网络，目标部署在 Jetson 设备上。
  • 官方介绍如下文所示

tkDNN is a Deep Neural Network library built with cuDNN and tensorRT primitives, specifically thought to work on NVIDIA Jetson Boards. It has been tested on TK1(branch cudnn2), TX1, TX2, AGX Xavier, Nano and several discrete GPUs. The main goal of this project is to exploit NVIDIA boards as much as possible to obtain the best inference performance. It does not allow training.

• 为什么要学 tkDNN
  • 这是 darknet-yolov4 tensorrt 部署方案之一，也是 darknet 上应该推荐的是这个。
  • 其他还有 yolov4 tensorrt 的方案，但要么是使用了 Python API，要么是使用了别的版本的 yolov4。

1. 使用

• 第一步：安装依赖（所有安装过程都可以参考repo提供的Dockerfile）
  • cmake 必须在 3.15 及以上
  • tkDNN 的依赖库包括 CUDA/CUDNN/TensorRT/Eigen3/OpenCV/yaml-cpp
  • Eigen3/yaml-cpp 直接使用 apt 安装即可：sudo apt install libyaml-cpp-dev libeigen3-dev
  • OpenCV 导入：好像对于 libopencv_dnn/libopencv_videoio 的导入有问题，需要修改配置文件。
  • 关于 CUDA/CUDNN/TensorRT：猜测默认 CMake 配置文件导入的是 apt 安装的，但我一般都使用离线安装的，需要修改一些内容。

# CMakeLists.txt 文件需要修改的内容如下

# CUDA 相关修改
include_directories(/usr/local/cuda/include) # 新增
link_directories(/usr/local/cuda/lib64) # 新增
find_package(CUDA 9.0 REQUIRED)

# opencv 相关修改
# https://github.com/ceccocats/tkDNN/issues/210
# 方案一
# find_package(OpenCV REQUIRED)
# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DOPENCV")
include_directories("/usr/include/opencv4/")
link_directories("/usr/lib/aarch64-linux-gnu/")
link_libraries(
    opencv_calib3d
    opencv_core
    opencv_dnn
    opencv_features2d
    opencv_flann
    opencv_gapi
    opencv_highgui
    opencv_imgcodecs
    opencv_imgproc
    opencv_ml
    opencv_objdetect
    opencv_photo
    opencv_stitching
    opencv_videoio
    opencv_video)
# 方案二
set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} ${CUDNN_LIBRARIES} ${OpenCV_LIBS} yaml-cpp -lopencv_dnn -lopencv_videoio)

# cmake/FindCUDNN.cmake 文件修改内容如下

# 主要就是导入 TensorRT 的形式修改，如果是用 apt 安装的tensorrt那就不用修改，如果是离线版安装的就需要修改
# find_cuda_helper_libs(nvinfer)
# set(NVINFER_LIBRARY ${CUDA_nvinfer_LIBRARY} CACHE FILEPATH "location of the nvinfer library")
# unset(CUDA_nvinfer_LIBRARY CACHE)
set(NVINFER_LIBRARY /home/ubuntu/TensorRT/lib/libnvinfer.so)

• 第二步：编译。其实没啥好说的，就是 mkdir build && cd build && cmake .. && make
  • 编译好后有两类程序 test_xxx 和 demo。
  • 前者用于生成 TensorRT Engine 文件，后者通过 Engine 文件对视频文件进行推理。
• 第三步：测试DarkNet

./test_yolo4tiny # 生成 yolo4tiny_fp32.rt 文件
./demo yolo4tiny_fp32.rt ../demo/yolo_test.mp4 y # 推挤结果通过 cv2.imshow 展示

2. 源码浏览

2.1 文件夹结构

• C++/Cuda源码位于四个位置
  • include：所有头文件
  • demo：demo 代码，使用 engine 文件进行推理，测速或计算mAP
  • src：主要代码都在这里
  • tests：其实就是使用 tkDNN 源码生成 TensorRT Engine 文件
• Cuda 源码全部位于 src/kernels
• cmake 配置文件
  • 根目录下有 CMakeLists.txt，该文件中配置了 set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/cmake)，所以会读取 cmake/*.cmake 文件。
  • cmake/ 文件夹中有 FindCUDNN.cmake 和 tkDNNConfig.cmake
  • 前者在 find_package(CUDNN) 时调用，会配置CUDNN信息以及TensorRT信息
  • 后者不知道什么时候调用，可能是将 tkDNN 作为第三方包，其他库要用的时候调用？

2.2 项目总体结构

• include 中的头文件与 src 中的原文件配合使用，总体功能包括：
  • 模型结构的抽象，如 Network、Layer 等
  • 所有Layer的具体实现，实现方法包括调用cudnn、cuda代码、cuBLAS等
  • TensorRT相关，主要就是将Network通过TensorRT API构建，所有Layer基本上都改写为 TensorRT Plugin
  • 其中 src/kernels 中的所有代码都是 .cu 文件，即cuda文件，需要使用 nvcc 进行编译
  • 所有 TensorRT Plugin 都定义在 include/tkDNN/pluginsRT 中，基本上继承了 IPlugin 对象，通过 kernels 或其他 src 中的文件实现模型前向推理。
  • 也定义了一些工具类，比如模型量化、画图、评估模型结果等。
• demo 和 tests 文件夹中的内容差不多，可以看做是调用 tkDNN 中的 API 进行的一些应用
  • tests 中的内容就是通过训练好的模型构建 TensorRT Engine
  • demo 中的内容就是通过构建好的 TensorRT Engine 文件对视频进行推理，或在测试数据集进行测试、计算 mAP