torch2trt代码走读

torch2trt code reading

项目地址： https://github.com/NVIDIA-AI-IOT/torch2trt

commit id:170e44954252342557af2f15a0a665d0afd1572b

简介

torch2trt 是 nvidia tensort 组 推出的，将torch 模型（torch.nn.Module）整图转换为CudaEngine 并用 TRTModule进行封装（TRTModule 继承于torch.nn.Module，用于forword）这样就实现了以下几个特点：

1.独立于torch提供了一套基于trt且兼容torch的量化/转换/推理

2.从用户角度，只需简单调用convert 接口，推理部分不用做任何改变（因为继承自torch.nn.module）

核心类/函数

Flattener (torch2trt/flattener.py:31)

as we all konw that, torch.nn.module forward 接口的入参是非常灵活多变的，可以支持dict/list/nn.tensor多种，然而trt 只支持以list形式输入，所以，Flattener 左右就是将输入变成list，也可以从list 中还原原始输入。

torch2trt函数 （torch2trt/torch2trt.py:654）

整个工程的入口，经过这个函数生成 trtModule 实例

tensorrt_converter (torch2trt/torch2trt.py:855)

一个python 装饰器，装饰所有的 convert函数（在torch2trt/converters 下 py文件中定义了转换函数）其作用就是建立torch中的函数/类（带转换的）和torch2trt中的转换函数之间的联系，并把其注册到 CONVERTERS 中

ConversionHook（torch2trt/torch2trt.py:326）

转换钩子函数，其handle convert函数，具体做法是：

1.将torch 中的待转换函数（原函数）替换成 wrapper（torch2trt/torch2trt.py:291）

2.在wrapper 函数中首先调用 原函数进行forward 然后调用转换函数为network增加对应的 layer

ConversionContext （torch2trt/torch2trt.py:417）

ConversionContext 是一个十分重要的类，主要完成了如下工作：

1.收集所有的 转换函数，并用conversionhook 封装

2. 遍历 待转换得 torch.nn.module 里的所有 sub module，给他们设置回调函数（通过torch的 接 口：module.register_forward_pre_hook/module.register_forward_hook）

2.1 .register_forward_pre_hook 设置的是module.forward 前的回调函数，他会把当前转换的module 入栈（由list 实现，放在ConversionContext 内），同时检查在 传入的参数中有无针对此module 跑在什么device上（DLA/GPU）的信息。

2.2 .register_forward_hook 注册的是module.forward 后的回调函数，他会把当前转换的module 出栈

TRTModule

TRTModule 继承 torch.nn.module 实现了使用ICudaEngine进行推理，其依照nn。module的接口实现了序列化/反序列化/推理。

代码走读

模型转换

入口：torch2trt/torch2trt.py:654

首先在 line:679 将入参都加入到 kwargs中。然后判断输入（inputs ）的类型，创建dataset，torch2trt 提供了如下几种dataset:

Dataset 用来handle input数据，统计出这组数据的 max shape/min shape（用于optimize profiile trt对dynamic shape 支持）。

line 697 ～ 724： 判断入参（关于shape的）合法性，不合法就从dataset中取。

line 735 行开始 进行模型转换们首先是 line 735 ～ 769：如果使用onnx，则:

1.构造一个当前转换模型的 warpper，其功能是：对于输入，使用设置的 flatterner 将 flat的 输入转会原始的输入，塞入原始模型中，对于输出，使用设置的flattener 将输出flattern，这个warpper的定义在：torch2trt/flatten_module.py:23.

2.用 torch.onnx.export 将 module warpper/inputs/outputs 塞入，转换生成onnx model。然后使用onnx_graphsurgeon（nvidia 自己出的 用来 onnx model 修改的 python package：https://github.com/NVIDIA/TensorRT/tree/master/tools/onnx-graphsurgeon/onnx_graphsurgeon）

导入onnx 模型，并作一点点优化（常量折叠）然后丢给 trt自己的onnxParser 将其转成 INetwork。

line 770 开始，是不使用onnx，直接对torch 模型进行转换的代码：

首先使用 with 语法创建 ConversionContext对象（以下简称ctx），在python中 with语句会创建一个作用域，在进入这个作用域后 会调用 with 实例化出的对象的 __enter__ 方法，在离开该作用域前调用此对象的 __exit__ 方法，因此我们需要先搞清楚 ctx的 __init**/** ​enter__/ __exit__ 函数做了什么：

__init__（torch2trt/torch2trt.py:419）:构造函数主要功能:

1.初始化了一堆成员变量，其中 INetwork对象使用了一个 warpper（NetworkWarpper）

1.1 NetworkWarpper的定义在：torch2trt/torch2trt.py:362，其主要作用是在INetWork 的 add XXXLayer 增加回调函数（torch2trt/torch2trt.py:396）针对其返回的结果（ILayer）调用 _configure_layer,在此函数中首先创建了SizeWrapper （用于决定是否给 torch.Tensor的size函数增加warpper这里由于是选择不替换所以在这里先不讲）实例以及一个作用域。此作用域的一大坨代码主要做了两个事情：1.根据当前处理submodule的name 为ILayer 设置name，2. 根据device_type 为当前layer 设置device type

2.使用ConversionHook 手机所有的转换函数，并将其包装成hook

3. 对于 module 内部的 submodule建立一个map（key 为module ，value为module name）

__enter__（torch2trt/torch2trt.py:480）：

1.在line 483-485 行，调用所有hook的 __enter__(定义在：torch2trt/torch2trt.py:338) 此时完成了 对 待转换函数（torch.nn.fucntional.XXX/torch.nn.ConvXd…）的warpper。

2.在487-493 行给每个submodule设置了两个回调函数：

3.给torch.Tensor.size/torch.Tensor.__getattribute__ 设置 warpper

3.1 size warpper 首先判断tensor是否经过了trt转换（增加了属性`_trt`(ITensor)）,如果没有直接返回原始的size，如果有，则会通过给network增加 shape op → slice op →connection op 为返回的 size 也增加 _trt 的 Itensor

3.2 get_attribute warpper 也就是针对调用 shape 属性时判断是否使用 size warpper

__exit__（torch2trt/torch2trt.py:501）：

此函数主要做了还原的工作：

1.调用hook的 __exit__ 取消warpper

2.删除submodule的 回调函数

3.还原 torch.Tensor的两个warpper

line 776：

调用 ctx 的 add_inputs 函数，此函数的定义在torch2trt/torch2trt.py:516，主要作用是给 INetwork添加input（当然需要考虑到 tensor name 和 tensor位置），然后给原始的torch tensor 增加属性： *trt, 然后调用 module的前向传播，由于在 * __enter__ 是 给torch相关的 api 做了warpper 因此此时就会触发 转换，为 INetwork增加layer。

最后（line 781）为network makr output

至此 INetwork 组图过程结束

line 783-828 行的代码是 配置 buildConfig 然后build engine 的接口调用

line 830-835 实例化 TRTModule 并返回。

converter function

以 conv2d （trt version≥7.0）为例讲解convert 函数的编写： torch2trt/converters/Conv.py：7

首先获取torch 原类型对象（module）以及输入，然后调用函数add_missing_trt_tensors（torch2trt/torch2trt.py：142）此函数对于没有转换 trt的 torchtensor（即没有 _trt 属性的）会直接以常量形式添加到INetwork中。

然后line 13：35 congmodule中获取attributes （拷贝到cpu），然后为network add conv layer且为输出增加_trt 属性

TRTModule 使用

torch2trt/torch2trt.py：558

1.在构造函数中创建 excutionContext 用于推理，还有 注册了一个回调函数 给state_dict(可以理解为序列化成 dict时)，此回调函数主要用于将 engine/input name/output name 等

2.重载了_load_from_state_dict 用于从dict反序列化出TRTModule

3.forward 函数，注意输入的tensor要在 gpu上