TensoeRT加速

1. YOLOv8-3D运行

代码地址如下：

https://github.com/bharath5673/YOLOv8-3D/tree/23264ab5803f6bac8a423175f0083143c39e13d5/assets

1.1 环境配置

先从上述网址clone下代码放入文件夹，在进行环境配置。

所用库版本如下所示：

CUDA 11.4.315 cuDNN 8.6.0.166 TensorRT 8.5.2.2 Jetpack 5.1.1

tensorflow 2.11.0+nv23.01   torch 1.14.0a0

安装方法与小车开发板环境配置过程相同，可参考小车环境配置过程。

1.2 代码运行

环境配置完成后，先从训练部分入手，训练选用kitti数据集。

下载数据集，已下载好的数据集放置在orin nano 8G板子上，自取。

下载完成后，将train代码中数据集地址改为自己放数据集的地址。

代码给出多种网络选择，推荐使用mobilenetv2，后续工作都是基于此网络。

select_model = 'resnet50'
select_model ='resnet101'
select_model = 'resnet152'
select_model = 'vgg11'
select_model = 'vgg16'
select_model = 'vgg19'
select_model = 'efficientnetb0'
select_model = 'efficientnetb5'
select_model = 'mobilenetv2'

数据集及网络选择完成后，根据算力自定义训练参数。 

## select model  and input size
# select_model = 'resnet18'
# select_model = 'resnet50'
# select_model ='resnet101'
# select_model = 'resnet152'
# select_model = 'vgg11'
# select_model = 'vgg16'
# select_model = 'vgg19'
# select_model = 'efficientnetb0'
# select_model = 'efficientnetb5'
select_model = 'mobilenetv2'

label_dir = '/home/jetson/YOLOv8-3D/YOLOv8-3D-main/training/label_2/'
image_dir = '/home/jetson/YOLOv8-3D/YOLOv8-3D-main/training/image_2/'

设置完成后就可以进行训练。

可能有如下报错：

ImportError: /lib/aarch64-linux-gnu/libGLdispatch.so.0: cannot allocate memory in static TLS block

 错误大概是在导入cv2时报错，无法正确分配内存，原因是找不到相应的.so文件，解决方法如下：

export LD_PRELOAD=/lib/aarch64-linux-gnu/libGLdispatch.so.0

ValueError: Unknown loss function: custom_loss_function. Please ensure this object is passed to the `custom_objects` argument

这是因为Keras在保存模型时，会保存训练时候的一切东西，包括优化器状态、权重、损失函数等，但由于这个损失函数函数是我们自定义的，并不在系统自带的损失函数函数中，我们下次重新加载的时候，自然就找不到这个损失函数，故引发了ValueError: Unknown loss function: focal_loss。
Keras早就想到了这点，因此可以在load_model中通过custom_objects参数传入自定义的类或者函数：

model = tf.keras.models.load_model('/home/jetson/YOLOv8-3D/YOLOv8-3D-main/mobilenetv2/mobilenetv2_weights.h5', custom_objects={"orientation_loss": orientation_loss}, compile = False)

tf.keras.models.load_model() 加载模型时，有一个参数compile，默认是True，会自动compile。但用户自定义的loss或者metric无法被识别。这时候需要把自定义的函数通过custom_objects传进去。

但是，使用custom_objects参数传入自定义函数，可以解决'.h5'格式的模型。不能解决savedmodel格式的模型。设置compile = False 

2.TensorRT应用

TensorRT 的核心是一个 C++ 库，可促进对 NVIDIA 图形处理单元 (GPU) 的高性能推理。它旨在与 TensorFlow、PyTorch、MXNet 等训练框架以互补方式工作。

TensorRT 通过组合层和优化内核选择来优化网络，以改善延迟、吞吐量、电源效率和内存消耗。通过程序指定，它会优化网络以较低的精度运行，从而进一步提高性能并减少内存需求。

Engine（引擎）：在深度学习中，"Engine" 通常指的是深度学习推理引擎，即一个经过优化和编译的模型，以在特定硬件上进行高效的推理。例如，TensorRT（TensorRT是NVIDIA的深度学习推理库）可以将训练好的模型转换为 Engine 文件，以便在 GPU 上进行快速推理。

2.1 具体实现

对于YOLO系列代码和pytorch框架，过程较简单

可以直接将.pt权重转为.engine

from ultralytics import YOLO

# Load a model
model = YOLO('/home/jetson/YOLOv8-3D/YOLOv8-3D-main/yolov8n.pt')  # load a custom trained

# Export the model
model.export(format='engine',half=True,simplify=True)

yolo下使用相对较简单，推理时只用将权重改为engine权重地址就可以 

bbox2d_model = YOLO('/home/jetson/YOLOv8-3D/YOLOv8-3D-main/yolov8n.engine', task='detect')

对于tensorflow框架下，相对复杂

先将.h5权重转换为onnx格式

import tensorflow as tf
from tf2onnx import convert

def orientation_loss(y_true, y_pred):
    # Find number of anchors
    anchors = tf.reduce_sum(tf.square(y_true), axis=2)
    anchors = tf.greater(anchors, tf.constant(0.5))
    anchors = tf.reduce_sum(tf.cast(anchors, tf.float32), 1)
    
    # Define the loss
    loss = -(y_true[:,:,0]*y_pred[:,:,0] + y_true[:,:,1]*y_pred[:,:,1])
    loss = tf.reduce_sum(loss, axis=1)
    epsilon = 1e-5  ##small epsilon value to prevent division by zero.
    anchors = anchors + epsilon
    loss = loss / anchors
    loss = tf.reduce_mean(loss)
    loss = 2 - 2 * loss 

    return loss

# with tf.keras.utils.custom_object_scope({'orientation_loss': orientation_loss}):
#     model = tf.keras.models.load_model('/home/jetson/YOLOv8-3D/YOLOv8-3D-main/mobilenetv2/mobilenetv2_weights.h5')
# 加载 TensorFlow 模型
model = tf.keras.models.load_model('/home/jetson/YOLOv8-3D/YOLOv8-3D-main/mobilenetv2/mobilenetv2_weights.h5', custom_objects={"orientation_loss": orientation_loss}, compile = False)
model.compile(loss=orientation_loss, optimizer='Adam', metrics=['accuracy']) # 将 'your_optimizer' 替换为实际的优化器
# 将 TensorFlow 模型转换为 ONNX 格式
onnx_model, _ = convert.from_keras(model)

# 保存 ONNX 模型
with open('../onnx_model.onnx', 'wb') as f:
    f.write(onnx_model.SerializeToString())

在终端打开，将onnx格式转为engine格式

/usr/src/tensorrt/bin/trtexec --onnx=/home/jetson/YOLOv8-3D/YOLOv8-3D-main/onnx_model.onnx --saveEngine=your_model.engine

 后续优化工作可以从将3D部分的.engine应用起来入手，初步的从.h5权重转换为.engine已经完成。

3.文件夹介绍

mobilenetv2 会保存训练后权重

demo 运行代码

train 训练代码

h52onnx tensorflow中.h5转为.onnx代码

zhijian 剪枝权重