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

news2024/11/20 10:35:17

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

论文题目：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内存消耗和计算成本。