降低阶段¶
降低阶段由多个通道组成,这些通道是将图从高级表示形式映射到低级表示形式的操作。每个通道都做一些特定的事情,例如内联方法调用。其目的是显著减少转换阶段在实际映射到 TensorRT 时需要处理的内容。我们的目标是更接近 1->1 的操作转换,而不是寻找适用的子图,限制转换器的数量并缩小每个转换器的范围。
您可以通过将日志级别设置为 Level::kGraph
来查看每个通道的效果
使用的通道¶
EliminateCommonSubexpression¶
删除图中的公共子表达式
消除死代码¶
死代码消除将检查节点是否具有副作用,如果有则不会删除它。
消除异常或传递模式¶
脚本模块中的常见模式是维度保护,如果输入维度不是预期的维度,则会抛出异常。
%1013 : bool = aten::ne(%1012, %24) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py:248:11
= prim::If(%1013) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py:248:8
block0():
= prim::RaiseException(%23) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py:249:12
-> ()
block1():
-> ()
由于我们正在编译时解析所有这些,并且 TensorRT 图中没有异常,因此我们只需将其删除。
消除冗余保护¶
为输出完全由其输入决定的操作消除冗余保护,即如果此类操作的输入受到保护,则我们允许删除操作输出上的保护
冻结模块¶
冻结属性并内联常量和模块。传播图中的常量。
融合 AddMM 分支¶
脚本模块中的常见模式是不同维度的张量使用不同的构造来实现线性层。我们将这些不同的变体融合到单个变体中,该变体将被 Unpack AddMM 通道捕获。
%ret : Tensor = prim::If(%622)
block0():
%ret.1 : Tensor = aten::addmm(%self.fc.bias, %x9.1, %3677, %3, %3)
-> (%ret.1)
block1():
%output.1 : Tensor = aten::matmul(%x9.1, %3677)
%output0.1 : Tensor = aten::add_(%output.1, %self.fc.bias, %3)
-> (%output0.1)
我们将这组块融合到这样的图中
%ret : Tensor = aten::addmm(%self.fc.bias, %x9.1, %3677, %3, %3)
融合线性¶
匹配 aten::linear
模式并将其融合到单个 aten::linear
中。此通道将 JIT 生成的 addmm 或 matmul + add 融合回线性
融合展平线性¶
当输入层高于 1D 时,TensorRT 隐式地将输入层展平为全连接层。因此,当存在 aten::flatten
-> aten::linear
模式时,我们删除 aten::flatten
。
降低图¶
给定一个方法的图,其第一个参数是 %self,将其降低为所有属性访问都替换为图的显式输入的图(而不是在 %self 上执行的 prim::GetAttr 的结果)。返回一个元组 (graph, parameters),其中图的最后 module.parameters.size() 个输入是此方法中使用的可训练参数。其余输入是函数的真实输入。
降低元组¶
LowerSimpleTuples
:
删除 TupleConstruct 和 TupleUnpack 匹配的元组,但在 if 语句、循环中以及作为输入/输出时保留元组
LowerAllTuples
:
删除 _所有_ 元组,如果某些元组无法删除则引发错误,ONNX 使用它来确保转换前没有元组,但它不适用于输入包含元组的图。
模块回退¶
模块回退由必须成对运行的两个降低通道组成。第一个通道在冻结之前运行,以在图中应在 PyTorch 中运行的模块周围放置分隔符。第二个通道在冻结后标记这些分隔符之间的节点,以表示它们应在 PyTorch 中运行。
NotateModuleForFallback
在模块调用周围预冻结放置分隔节点,以表示图中节点应在 PyTorch 中运行的位置
MarkNodesForFallback
查找分隔符,然后标记分隔符之间的所有节点,以告知分区在 PyTorch 中运行它们
窥孔优化¶
此优化通道的目的是捕获您可能感兴趣的所有小的、易于捕获的窥孔优化。
- 目前,它执行以下操作
消除无操作“expand”节点
将 x.t().t() 简化为 x
删除 Contiguous¶
删除 contiguous 运算符,因为我们正在执行 TensorRT 内存已经是连续的。
删除 Dropout¶
由于我们正在执行推理,因此删除 dropout 运算符。
删除 To¶
删除执行类型转换的 aten::to
运算符,因为 TensorRT 自身会管理它。重要的是,这是最后运行的通道之一,以便其他通道有机会将所需的类型转换运算符移出主命名空间。
解包 AddMM¶
将 aten::addmm
解包为 aten::matmul
和 aten::add_
(以及一个额外的 trt::const
操作来冻结 TensorRT 图中的偏置)。这使我们可以重用 aten::matmul
和 aten::add_
转换器,而无需专用的转换器。
解包 LogSoftmax¶
将 aten::logsoftmax
解包为 aten::softmax
和 aten::log
。这使我们可以重用 aten::softmax
和 aten::log
转换器,而无需专用的转换器。
展开循环¶
展开兼容循环的操作(例如,足够短),以便您只需循环一次。
用 Repeat 替换 Tile¶
由于我们正在执行推理,因此删除 dropout 运算符。