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 倍。