深度学习论文: TRT-ViT: TensorRT-oriented Vision Transformer及其PyTorch实现

深度学习论文: TRT-ViT: TensorRT-oriented Vision Transformer及其PyTorch实现
TRT-ViT: TensorRT-oriented Vision Transformer
PDF: https://arxiv.org/pdf/2205.09579.pdf
PyTorch代码: https://github.com/shanglianlm0525/CvPytorch
PyTorch代码: https://github.com/shanglianlm0525/PyTorch-Networks

1 概述

通过广泛的实验，总结了在TensorRT上设计高效网络的4个设计准则：

• Stage-Level：在后期Stage使用Transformer Block可以最大化效率和性能；
• Stage-Level：由浅入深的Stage模式可以提升性能；
• Block-Level：Transformer和BottleNeck混合Block比Transformer效率更高；
• Block-Level：先全局后局部的块结构可以显著提高性能

2 Practical Guidelines for Efficient Network Design on TensorRT

分别定义 TeraFLOPS（计算密度） 和 TeraParams（参数密度） 来描述一个操作或块在硬件上的效率和容量。
其中：


通过这两个硬件上效率的直接度量总结了在 TensorRT 上有效设计高效网络的4个实用准则。

准则1：Transformer Block在后期使用可以最大限度地提高效率和性能之间的平衡

Transformer旨在建立特征之间的全局连接。而卷积只捕获局部信息。因此Transformer块比卷积块更强大但效率更低。当输入分辨率降低时，Transformer和BottleNeck之间的效率差距会缩小。因此在设计高效网络时将Transformer模块置于后期阶段可以平衡性能和效率。

遵循准则1，在后期将BottleNeck块替换为Transformer块，即 MixNetV。可以看到MixNetV比ViT更快，比ResNet更准确。

准则2：从浅到深的Stage模式可以提高性能

ResNet中BottleNeck 和ViT中Transformer 块的 TeraParams 在后期变得更加突出。因此将网络结构的早期阶段设计更浅，晚期阶段设计更深，可以在提高模型容量并在不降低效率的情况下获得更好的性能。

在ResNet50、PvT-Small 和 MixNetV上做调整，均可以看到性能提升。

准则3：BottleNeck Block和Transformer Block混合使用比Transformer更有效

基于BottleNeck Block和Transformer Block的特点，混合使用二者可能带来更好的性能和效率。基于此设计如下混合结构：

用MixBlockA和MixBlockB替换MixNetV中的Transformer块，构建了MixNetA和MixNetB。如下图，其中C、T和M分别表示Convolution、Transformer和MixBlock。

从实验可以看出，混合块比Transformer块更有效。而且对应的TeraParams也更好，这表明混合块有实现更好性能的潜力。

准则4：Global-then-Local Block有助于弥补性能问题

感受野通常很大的Transformer从输入特征中提取全局信息，并实现信息交换。相比之下，卷积的小感受野只挖掘局部信息。因此先获取全局信息然后再局部进行提炼的结构（ MixBlock C）性能会更好，实验也正好验证了这一点。

3 TRT-ViT

按照ResNet的基本配置构建了TRT-ViT 架构，其中采用特征金字塔结构，特征图的分辨率随着网络深度而降低，而通道数增加。

接下来，应用提出的设计准则来构建TRT-ViT。整个架构分为5个阶段，后期Stage只使用MixBlockC，而卷积层用于早期Stage。此外，使用了由浅入深的Stage模式，与ResNet中的Stage模式相比，早期Stage更浅，后期Stage更深。