TResNet: ResNet改进,实现高精度的同时保持高 GPU 利用率

终于开题,抓紧发文,然后放飞​​​​​​,来由就是想搞一篇论文,但是增加了某个东西之后吧,速度变慢了,所以导师提议加个这玩意看看能不能快点。

论文题目:TResNet: High Performance GPU-Dedicated Architecture

论文地址:https://arxiv.org/abs/2003.13630

代码:https://github.com/mrT23/TResNet

包含三个变体,TResNet-M、TResNet-L 和 TResNet-XL,它们仅在深度和通道数上有所不同。TResNet架构包含以下改进: SpaceToDepth stem,Anti-Alias downsampling,In-Place Activated BatchNorm,Blocks selection and SE layers。有些改进会增加模型的吞吐量,而有些则会降低模型的吞吐量。 就这5个,挨个说。

1.  SpaceToDepth stem

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

代码:

class SpaceToDepth(nn.Module):
    def __init__(self, block_size=4):
        super().__init__()
        assert block_size == 4
        self.bs = block_size

    def forward(self, x):
        N, C, H, W = x.size()
        x = x.view(N, C, H // self.bs, self.bs, W // self.bs, self.bs)  # (N, C, H//bs, bs, W//bs, bs)
        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # (N, bs, bs, C, H//bs, W//bs)
        x = x.view(N, C * (self.bs ** 2), H // self.bs, W // self.bs)  # (N, C*bs^2, H//bs, W//bs)
        return x

2.  Anti-Alias downsampling

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

class AADownsample(nn.Module):
    def __init__(self, filt_size=3, stride=2, channels=None):
        super(AADownsample, self).__init__()
        self.filt_size = filt_size
        self.stride = stride
        self.channels = channels


        assert self.filt_size == 3
        a = torch.tensor([1., 2., 1.])

        filt = (a[:, None] * a[None, :])
        filt = filt / torch.sum(filt)

        # self.filt = filt[None, None, :, :].repeat((self.channels, 1, 1, 1))
        self.register_buffer('filt', filt[None, None, :, :].repeat((self.channels, 1, 1, 1)))

    def forward(self, input):
        input_pad = F.pad(input, (1, 1, 1, 1), 'reflect')
        return F.conv2d(input_pad, self.filt, stride=self.stride, padding=0, groups=input.shape[1])

3. In-Place Activated BatchNorm (Inplace-ABN)

在整个架构中,作者将所有BatchNorm + ReLU层替换为Inplace-ABN 层,该层将BatchNorm和activation作为一个单独的就地操作来实现,从而大大减少了训练深度网络所需的内存,而计算量的增加可忽略不计成本。

在TResNet模型中使用Inplace-ABN具有以下优点:

BatchNorm层是GPU内存的主要消耗者。 用Inplace-ABN替换BatchNorm层实际上可使最大批处理大小增加一倍,从而提高了GPU吞吐量。
对于TResNet来说,Leaky-ReLU比普通的ReLU提供更好的准确率。虽然一些现代的激活函数,例如Swish和Mish,也可能与ReLU相比,它们提供了更好的精度,它们的GPU内存消耗更高,并且其计算成本更高。相反,Leaky-ReLU具有与普通ReLU完全相同的GPU内存消耗和计算成本。

4. Blocks Selection

下图左边为ResNet34 使用的BasicBlock,右边为ResNet50使用的Bottleneck, Bottleneck使用GPU更高,但是可以得到更高精度, BasicBlock有更大的感受野.

因此, TResNet在前两阶段使用BasicBlock,后两阶段使用Bottleneck

5. SE Layers

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

然后就是各种的对比实验结果加上消融实验。开题结束了,要写小论文了,我就是个大混子,毕业万岁!

猜你喜欢

转载自blog.csdn.net/Zosse/article/details/127783334