TensorRT基础入门（一）

Darchan

已于 2025-03-31 11:31:04 修改

阅读量997

点赞数 24

分类专栏：工程化技巧文章标签：人工智能 python 开源

于 2025-03-31 11:26:40 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_36354875/article/details/146508840

版权

工程化技巧专栏收录该内容

9 篇文章

订阅专栏

TensorRT基础入门（一）

本文属于学习笔记，在重点章节或代码位置加入个人理解，欢迎批评指正！

参考

CUDA与TensorRT部署部署实战第三章

一. TensorRT介绍

学习目标

TensorRT是什么，理解TensorRT的工作流程，TensorRT的极限，以及其他DNN编译器(其他厂商)

1.1 TensorRT是什么

针对NVIDIA GPU的一种优化编译器，同时也是开发工具。可以自动优化模型，寻找最优调度和并行策略。

1.2 常用工作流程

pytorch模型 -> onnx -> TensorRT

1.3 限制

算子不支持
不同TensorRT版本
优化不及预期：预期使用TensorCore, 实际因为内部一些额外处理操作，分配了CUDA core
天生并行差的layer：如1*1 conv

二. TensorRT应用场景

学习目标

TensorRT意义，自动驾驶部署中需要关注的内容

2.1 优化意义

模型训练：精度越高越好
模型部署：压缩模型

2.2 自动驾驶部署中关注的内容

实时性（realtime）
电力消耗：小于100w
远距离

三. TensorRT模块

学习目标

简单理解Tensor中的优化策略，为什么可以做层融合，理解量化的意义

3.1 Tensor中的优化策略

层融合：减少kernel开销与memory操作，提高效率。
1）水平层融合，水平方向存在同类layer。
2）垂直层融合（conv+bn），公式
Tips：模型中新的激活函数，替换为relu尝试，relu内部没有计算
kernel auto-tuning：针对不同kernel, 硬件会自己采用不同优化策略
Quantization：量化，将单精度类型(FP32)训练权重转变为半精度(FP16)或者整型(INT8, INT4)
fp16浮点数表示法：

四. 导出并分析onnx

学习目标

学习PyTorch导出ONNX的方法，以及onnxsimplify的使用。

4.1 导出onnx方法

使用torch.onnx.export函数，注意其中一些参数的用法

   torch.onnx.export(
        model         = model, 
        args          = (input,),
        f             = "../models/example_two_head.onnx",
        input_names   = ["input0"],
        output_names  = ["output0", "output1"],   # 双输出参数
        dynamic_axes  = {                                   # 动态shape
            'input0':  {0: 'batch'},
            'output0': {0: 'batch'}，
            'output1': {0: 'batch'}，
        },
        opset_version = 12)

4.2 导出后检查

 # 检查导入的onnx model
 onnx.checker.check_model(model_onnx)

4.3 常见导出onnx有变换情况

4.3.1 conv+batchnorm

导出onnx模型时，原模型中conv + bn，进行了层融合，详细可以参考公式推导。故使用netron进行查看，便看不到br节点。

4.3.2 torch.flatten导出一系列节点

flatten算子导出时，完成了维度变换 B, C, H, W -> B, C, L 。这一个过程产生了shape->slice->concat->reshape这一系列计算节点, 为什么？

原因

‌shape 节点‌：获取输入张量的完整形状信息（如 [B, C, H, W]）。
‌slice 节点‌：提取前两个维度（B 和 C）。
‌concat 节点‌：将 [B, C] 与 H*W（通过 -1 自动计算）拼接成目标形状 [B, C, -1]。
‌reshape 节点‌：根据拼接后的形状执行实际维度合并。

x = torch.flatten(x, 2, 3)

flatten算子简化前后对比
在这里插入图片描述

4.4 使用onnx-simplifier来进行onnx的简化

import onnxsim
## 简化onnx模型
    model_onnx, check = onnxsim.simplify(model_onnx)
    assert check, "assert check failed"
    onnx.save(model_onnx, file)

五. 剖析onnx架构，理解ProtoBuf

参考：ONNX模型格式的结构和原理

学习目标

学习ONNX的Proto架构，使用onnx.helper创建onnx，修改onnx

5.1 onnx是什么

onnx是一种神经网络的格式，采用Protobuf二进制形式进行序列化模型，根据用于定义的数据结构来进行序列化存储。

5.2 onnx中的proto结构

详细位置：onnx proto 数据结构源码
ONNX模型分析与使用

关键性数据结构

- ModelProto ：整个模型信息，最顶层结构
--- GraphProto ：描述模型的计算逻辑，由一系列的节点组成。这些节点通过它们的输入和输出形成一个‌有向无环图（DAG）‌。
----- NodeProto（repeated ）：节点。每个计算节点当中还包含了一个 AttributeProto 数组，用于描述该节点的属性，例如 Conv 层的属性包含 group，pads 和strides 等等
----- TensorProto（repeated ）： 权重， initializer， raw_data是以字节byte形式存储的权重信息。
----- ValueInfoProto（repeated ）：输入/输出信息
----- TensorAnnotation（repeated ）：量化信息

Tips. onnx 的IR指的是什么？
ONNX 中的 ‌IR（Intermediate Representation，中间表示）‌ 是模型在不同深度学习框架和工具之间转换时的标准化中间格式，核心目的是实现跨框架兼容和高效部署。其具体含义和结构可总结如下：

IR 的定义与核心组成‌
1‌）计算图（Graph）‌
以有向无环图（DAG）形式表示模型结构，包含输入/输出张量（Tensor）、计算节点（Node）和依赖关系。每个节点对应一个算子（Operator），节点间通过张量传递数据‌。
2）张量（Tensor）‌
定义数据的类型、形状和传输路径，用于描述输入、输出及中间计算结果‌。
3‌）算子集合（Operator Set）‌
覆盖常见深度学习操作（如卷积、池化），通过版本控制支持不同框架的算子兼容性‌。

5.3 创建onnx

onnx提供了一些很方便的api来创建onnx

onnx.helper.make_tensor
onnx.helper.make_tensor_value_info
onnx.helper.make_attribute
onnx.helper.make_node
onnx.helper.make_graph
onnx.helper.make_model

Eg. 创建conv节点

    # 创建conv节点，注意conv节点的输入有3个: input, w, b
    conv1_node = onnx.helper.make_node(
        name         = "conv2d_1",
        op_type      = "Conv",
        inputs       = [
            model_input_name, 
            conv1_weight_name,
            conv1_bias_name
        ],
        outputs      = [conv1_output_name],
        kernel_shape = [conv1_kernel, conv1_kernel],
        pads         = [conv1_pads, conv1_pads, conv1_pads, conv1_pads],
    )

5.4 解析权重

initializer中的raw_data （权重）是以byte形式存储，需要转换再读取。

# 注意，因为weight是以字节的形式存储的，所以要想读，需要转变为float类型
def read_weight(initializer: onnx.TensorProto):
    shape = initializer.dims
    data  = np.frombuffer(initializer.raw_data, dtype=np.float32).reshape(shape)
    print("\n**************parse weight data******************")
    print("initializer info: \
            \n\tname:      {} \
            \n\tdata:    \n{}".format(initializer.name, data))