TResNet: High Performance GPU-Dedicated Architecture 来自阿里的达摩院,发布于2021 WACV,该论文引入了一系列架构修改,旨在提高神经网络的准确性,同时保持其 GPU 训练和推理效率。

论文首先讨论了面向 FLOP 的优化引起的瓶颈。然后建议更好地利用 GPU 结构的设计。最后引入了一个新的 GPU 专用模型,称其为 TResNet。

动机

上表将 ResNet50 与流行的较新架构进行了比较,具有相似的 ImageNet top-1 精度——ResNet50-D [11]、ResNeXt50 [43]、SEResNeXt50 (SENet+ResNeXt) [13]、EfficientNet-B1 [36] 和 MixNet-L (MixConv)[37]。与 ResNet50 相比,新提出的网络中 FLOP 减少和新技巧的使用并未转化为 GPU 吞吐量的提高。

最新的一些网络,如 EfficientNet、ResNeXt 和 MixNet (MixConv) 广泛使用深度和 1×1 卷积,它们提供的 FLOP 明显少于 3×3 卷积。但是GPU 通常受限于内存访问成本而不是计算数量,尤其是对于低 FLOP 层。ResNeXt 和 MixNet (MixConv) 等网络广泛使用了多路径。对于训练这会创建大量需要存储以进行反向传播的激活图,占用大量的显存肯定会减小批量大小,从而降低 GPU 吞吐量。

而TResNet 的提出旨在实现高精度的同时保持高 GPU 利用率。

TResNet:ResNet 的改进和变化

包含三个变体,TResNet-M、TResNet-L 和 TResNet-XL,它们仅在深度和通道数上有所不同。

SpaceToDepth Stem

ResNet50 stem 由一个 stride-2 conv7×7 和一个最大池化层组成。ResNet-D 将 conv7×7 替换为三个 conv3×3 层。这种设计确实提高了准确性,但代价是降低了训练吞吐量。论文使用了专用的 SpaceToDepth 转换层 [33],将空间数据块重新排列为深度。SpaceToDepth 层之后是简单的卷积,以匹配所需通道的数量。

Anti-Alias Downsampling (AA)

stride-2 卷积被 stride-1 卷积替换,然后是一个 3×3 的步长为 2的blur filter。

In-Place Activated BatchNorm (Inplace-ABN)

所有 BatchNorm+ReLU 层都被 Inplace-ABN [32] 层取代,该层将 BatchNorm 激活为单个inplace操作,从而显着减少训练深度网络所需的内存,而计算成本仅略有增加。并且使用 Leaky-ReLU 代替了 ResNet50 的普通 ReLU。

Novel Block-Type Selection

Bottleneck 层比 BasicBlock 层具有更高的 GPU 使用率,并且提供更好的准确性。但是BasicBlock 层具有更大的感受野,因此它们可能更适合放置在网络的早期阶段。由于 BasicBlock 层具有较大的感受野,因此它们被放置在网络的前两个阶段,而Bottleneck层则位于最后两个阶段。与 [10] 和 [36] 类似,也修改了初始通道数和第 3 阶段中的残差块数。架构详情如上表。

Optimized SE Layers

TResNet BasicBlock 和Bottleneck设计(stride 1)。IBN = Inplace-BatchNorm,r = 缩减因子

SE 层仅放置在网络的前三个阶段,以获得最大的速度-准确度优势。对于Bottleneck单元,在conv3×3操作之后添加SE模块,缩减因子为8(r = 8)。对于 BasicBlock 单元,在残差和之前添加 SE 模块,缩减因子为 4 (r=4)。

除了架构改进之外,进行了下面的一些代码优化。

JIT 编译能够在执行时将高级代码动态编译成高效、优化的机器代码。这与通过解释器动态运行代码的默认 Pythonic 选项形成对比。对于 AA 和 SpaceToDepth 模块,发现 JIT 编译几乎可以将 GPU 成本降低两倍。

Inplace操作直接更改给定张量的内容,无需内存之间的复制,这样可以防止创建不需要的激活映射不需要进行反向传播。所以尽量使用Inplace操作。TResNet-M 的最大批量大小几乎是 ResNet50-512 的两倍,

Fast Global Average Pooling 是 GAP 的一个简单专用实现,针对 (1,1) 空间输出的特定情况优化了代码,比 GPU 上的样板实现快 5 倍。

实验结果

ImageNet

TResNet-M 的 GPU 吞吐量与 ResNet50 相似,在 ImageNet 上的验证准确率显着提高(+1.8%)。在 8×V100 GPU 机器上训练 TResNet-M 和 ResNet50 模型耗时不到 24 小时,可见该训练方案也是高效经济的。TResNet 模型的另一个优势是能够处理比其他模型大得多的批量大小。

消融研究

  1. 虽然预期 GPU 吞吐量会提高,但在用“更便宜”的 SpaceToDepth 单元替换 ResNet stem时,准确度也会(略微)提高这一事实有点令人惊讶。
  2. 块类型选择为所有方法提供了显着改进。
  3. Inplace-ABN 显着增加了批量大小。Inplace-ABN 的影响好坏参半:虽然推理速度有所提高,但训练速度有所降低。
  4. 优化的 SE + Anti-Aliasing 层显着提高了 ImageNet top-1 精度,但代价是降低了模型 GPU 吞吐量。

ResNet-M 模型中的代码优化对推理速度、训练速度和最大批量大小的影响

在优化中,inplace操作提供了最大的提升——它不仅提高了 GPU 吞吐量,而且还显着增加了批量大小,因为它避免了为反向传播创建不需要的激活图。

High-Resolution Fine-Tuning

使用 224 的 ImageNet 预训练的 TResNet 模型作为起点,并通过 10 次 epoch 微调到 448 的输入分辨率。TResNet 模型可以很好地扩展到高分辨率。即使是 TResNet-M 这种相对较小紧凑的模型,在高分辨率输入的 ImageNet 上也可以达到 83.2% 的 top-1 准确率。

与 EfficientNet 模型的比较

沿着 top-1 准确度曲线,TResNet 模型比 EfficientNet 模型提供更好的推理速度准确度和训练速度准确度权衡。

迁移学习

TResNet 在迁移学习数据集上与最先进模型的比较(仅基于 ImageNet 的迁移学习结果)使用 ImageNet 预训练,并对模型进行了 80 个 epoch 的微调。TResNet 在 4 个数据集中的 3 个数据集上超过或匹配了最先进的精度,GPU 推理速度提高了 8-15 倍。

基于 TResNet 的解决方案明显优于之前 MSCOCO 多标签数据集的顶级解决方案,将已知 SOTA 从 83.7 mAP 大幅提高到 86.4 mAP。

使用 FCOS 作为目标检测器,TResNet-M 在这个目标检测任务上的表现优于 ResNet50,将 COCO mAP 得分从 42.8 提高到 44.0。

论文

[2021 WACV] TResNet: High Performance GPU-Dedicated Architecture

https://avoid.overfit.cn/post/d032d7b8ca65463b95c70fc4950a6e1a

作者:Sik-Ho Tsang


deephub
111 声望89 粉丝