资讯

历史

科技

环境与自然

成长

游戏

财经

文学与艺术

美食

健康

家居

文化

情感

汽车

三农

军事

旅行

运动

教育

生活

星座命理

TensorRT优化原理及模型转换的三种方式

创作时间:

作者:

@小白创作中心

TensorRT优化原理及模型转换的三种方式

引用

来源

https://315386775.github.io/posts/2023/12/model-tensorrt/

TensorRT 是 NVIDIA 开发的一个高性能深度学习推理框架，主要用于优化和加速在 NVIDIA GPU 上运行的深度学习模型。通过将训练好的模型转换为 TensorRT 格式，可以显著提升模型在 GPU 上的运行速度。本文将详细介绍 TensorRT 的优化原理及其模型转换的多种方式。

TensorRT 优化原理

TensorRT 通过以下五种类型的优化来提高深度学习模型的吞吐量：

算子融合（层与张量融合）：通过融合一些计算 op 或去掉一些多余 op 来减少数据流通次数以及显存的频繁使用。例如，卷积、偏置和激活层可以合并成一个 CBR 结构，只占用一个 CUDA 核心。对于 Conv + ReLU 这样的结构，TensorRT 通常会将其合并成一个 Conv 进行运算。
Concat 层的消除：对于 channel 维度的 concat 层，TensorRT 通过非拷贝方式将层输出定向到正确的内存地址来消除 concat 层，从而减少内存访存次数。

量化：TensorRT 支持 IN8 量化、FP16、TF32 等不同精度的使用，这些精度可以显著提升模型执行速度且不会显著降低模型精度。在转换为 FP16 时，由于 FP16 的动态范围低于 FP32，一些权重可能会因为溢出而缩小。使用缩放和偏置项以 INT8 精度映射这些权重。为了在 INT8 TensorRT 中表示 FP32 分布，使用 KL 散度来测量差异并将其最小化。TensorRT 使用迭代搜索而不是基于梯度下降的优化来寻找阈值。
内核自动调整：根据不同的显卡架构、SM 数量、内核频率等，选择不同的优化策略和计算方式，寻找最适合当前架构的计算方式。例如，有多种方法可以执行卷积操作，但在这个选定的平台上哪一种是最优化的，TRT 会自动选择。
动态张量显存：通过调整一些策略可以减少模型中显存开辟和释放的次数，从而减少模型运行的时间。
多流执行：使用 CUDA 中的 stream 技术，最大化实现并行操作。

TensorRT 模型构建

使用 TensorRT 生成模型主要有两种方式：

直接通过 TensorRT 的 API 逐层搭建网络：这一过程类似使用一般的训练框架，如使用 Pytorch 搭建网络。
将中间表示的模型转换成 TensorRT 的模型：比如将 ONNX 模型转换成 TensorRT 模型。

直接构建

利用 TensorRT 的 API 逐层搭建网络，这一过程类似使用一般的训练框架，如使用 Pytorch 搭建网络。

IR 转换模型

直接使用 trtexec 工具，trtexec 工具有许多选项用于指定输入和输出、性能计时的迭代、允许的精度和其他选项。

构建阶段的常用参数

--onnx=<model> ：指定输入 ONNX 模型。
--minShapes=<shapes> , --optShapes=<shapes> , --maxShapes=<shapes> ：指定用于构建引擎的输入形状的范围。仅当输入模型为 ONNX 格式时才需要。
--memPoolSize=<pool_spec> ：指定策略允许使用的工作空间的最大大小，以及 DLA 将分配的每个可加载的内存池的大小。
--saveEngine=<file> ：指定保存引擎的路径。
--fp16 、 --int8 、 --noTF32 、 --best ：指定网络级精度。
--sparsity=[disable|enable|force] ：指定是否使用支持结构化稀疏的策略。
--best : 启用所有精度以达到最佳性能（默认 = 禁用）
--calib= : 读取INT8校准缓存文件

针对 BS=1 的模型进行转换

TensorRT-7.2.3.4/bin/trtexec \
        --onnx=model/detail_1_3_64_64.onnx \
        --workspace=4096 \
        --saveEngine=detail_1_3_64_64.trt \
        --device=0 \
        --verbose

动态尺度支持 NCHW 中的 N、H 以及 W，我们在转换模型的时候需要额外指定三个维度信息即可(最小、最优、最大)

TensorRT-7.2.3.4/bin/trtexec  --onnx=yolov4_-1_3_416_416_dynamic.onnx \
                              --minShapes=input:1x3x416x416 \
                              --optShapes=input:8x3x416x416 \
                              --maxShapes=input:8x3x416x416 \
                              --workspace=4096 \
                              --saveEngine=M-dynamic_b8_fp16.engine \
                              --fp16

针对精度转换为 INT8 的模型，需要校准数据

batch_size = 8
batched_input = np.zeros((batch_size, 224, 224, 3))
for i in range(batch_size):
    img_path = './data/img%d.JPG' % (i % 4)
    img = image.load_img(img_path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    batched_input[i, :] = x
    batched_input = tf.constant(batched_input)
print('batched_input shape: ', batched_input.shape)
conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
    precision_mode=trt.TrtPrecisionMode.INT8,
    max_workspace_size_bytes=8000000000,
    use_calibration=True)
converter = trt.TrtGraphConverterV2(
    input_saved_model_dir='resnet50_saved_model',
    conversion_params=conversion_params)
def calibration_input_fn():
    yield (batched_input, )
converter.convert(calibration_input_fn=calibration_input_fn)
converter.save(output_saved_model_dir='resnet50_saved_model_TFTRT_INT8')
print('Done Converting to TF-TRT INT8')

IR 转换模型：代码转换

从头定义网络创建 engine 的九个步骤：

创建logger
创建builder
创建network
向network中添加网络层
设置并标记输出
创建config并设置最大batchsize和最大工作空间
创建engine
序列化保存engine
释放资源

使用 Pytorch 实现一个只对输入做一次池化并输出的模型，其中需要使用 TensorRT 的 OnnxParser 功能，它可以将 ONNX 模型解析到 TensorRT 的网络中。实现代码如下：

import torch
import onnx
import tensorrt as trt
# network defined
class NaiveModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.pool = torch.nn.MaxPool2d(2, 2)
    def forward(self, x):
        return self.pool(x)
# load ONNX model
onnx_model = onnx.load('model.onnx')
# create builder and network
logger = trt.Logger(trt.Logger.ERROR)
builder = trt.Builder(logger)
EXPLICIT_BATCH = 1 << (int)(
    trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(EXPLICIT_BATCH)
# parse onnx
parser = trt.OnnxParser(network, logger)
parser.parse(onnx_model.SerializeToString())
config = builder.create_builder_config()
config.max_workspace_size = 1<<20
profile = builder.create_optimization_profile()
profile.set_shape('input', [1,3 ,224 ,224], [1,3,224, 224], [1,3,224 ,224])
config.add_optimization_profile(profile)
# create engine
with torch.cuda.device(device):
    engine = builder.build_engine(network, config)
with open('model.engine', mode='wb') as f:
    f.write(bytearray(engine.serialize()))
    print("generating file done!")

IR 转换时，如果有多 Batch、多输入、动态 shape 的需求，都可以通过多次调用 set_shape 函数进行设置。set_shape 函数接受的传参分别是：输入节点名称，可接受的最小输入尺寸，最优的输入尺寸，可接受的最大输入尺寸。一般要求这三个尺寸的大小关系为单调递增。

使用 onnx2trt 工具转 engine

#可以使用 onnx2trt 可执行文件将 ONNX 模型转换为序列化的 TensorRT 引擎
onnx2trt my_model.onnx -o my_engine.trt
#ONNX 模型也可以转换为人类可读的文本
onnx2trt my_model.onnx -t my_model.onnx.txt
#查看更多命令行参数
onnx2trt -h
#通过运行查看更多所有可用的优化过程
onnx2trt -p

TensorRT 周边配套

ONNX GraphSurgeon：可以修改我们导出的 ONNX 模型，增加或者剪掉某些节点，修改名字或者维度等；
Polygraphy：各种小工具的集合，例如比较 ONNX 和 trt 模型的精度，观察 trt 模型每层的输出等等，主要用来 debug 一些模型的信息，还是比较有用的。
TensorRT Plugin：目前只支持 C++ 实现，Plugin 是我们针对某个需要定制化的层或目前 TensorRT 还不支持的层进行实现、封装。TensorRT Plugin 是对网络层功能的扩展，TensorRT 官方已经预构建了一些在目标检测中经常使用的 Plugin，如：NMS、PriorBOX 等，我们可以在 TensorRT 直接使用，其他的 Plugin 则需要我们自己创建。