在 PyTorch 中实现 ConvNext

一个击败变形金刚的新卷积网络——你好!今天我们将在 2020 年代的 A ConvNet 中在 PyTorch 中实现著名的 ConvNext。代码在这里,本文的交互式版本可以从这里下载。让我们开始吧!该论文提出了一种新的基于卷积的架构,它不仅超越了基于…

在 PyTorch 中实现 ConvNext

一个击败变压器的新卷积网络

你好!!今天我们将在 2020 年代的 A ConvNet 中在 PyTorch 中实现著名的 ConvNext。[0]

代码在这里,本文的交互式版本可以从这里下载。[0][1]

让我们开始吧!

论文提出了一种新的基于卷积的架构,不仅超越了基于 Transformer 的模型(如 Swin),而且随着数据量的增加而扩展!下图显示了针对不同数据集/模型大小的 ConvNext 准确度。

因此,作者首先采用众所周知的 ResNet 架构,并根据过去十年中的新最佳实践和发现对其进行迭代改进。作者专注于 Swin-Transformer,并密切关注其设计选择。这篇论文是一流的,我强烈推荐阅读它:)

下图显示了所有各种改进以及每一项改进之后的各自性能。

他们将路线图分为两部分:宏观设计和微观设计。宏观设计是我们从高层次的角度所做的所有改变,例如阶段的数量,而微设计更多的是关于较小的事情,例如使用哪个激活。

我们现在将从一个经典的 BottleNeck 块开始,并逐个应用每个更改。

起点:ResNet

如你所知(如果你没有我有一篇关于在 PyTorch 中实现 ResNet 的文章)ResNet 使用一个残差的 BottleNeck 块,这将是我们的起点。[0]

让我们检查它是否有效

torch.Size([1, 64, 7, 7])

让我们也定义一个 Stage,一个块的集合。每个阶段通常将输入下采样 2 倍,这是在第一个块中完成的。

torch.Size([1, 64, 4, 4])

酷,注意输入是如何从 7×7 减少到 4×4 的。

ResNet 也有所谓的 stem,这是模型中对输入图像进行大量下采样的第一层。

很酷,现在我们可以定义 ConvNextEncoder 来保存阶段列表,并将图像作为输入生成最终嵌入。

torch.Size([1, 2048, 7, 7])

这是你的普通 resnet50 编码器,如果你附加一个分类头,你会得到好的旧 resnet50 准备好在图像分类任务上进行训练。

宏观设计

改变阶段计算比率

在 ResNet 中我们有 4 个阶段,Swin Transformer 使用 1:1:3:1 的比例(所以第一阶段一个块,第二个一个块,第三个第三个……)。将 ResNet50 调整到这个比率 ((3, 4, 6, 3) -> (3, 3, 9, 3)) 导致性能从 78.8% 提高到 79.4%。

将词干更改为“Patchify”

ResNet stem 使用非常激进的 7×7 conv 和 maxpool 对输入图像进行大量下采样。然而,Transformers 使用“patchify”词干,这意味着它们将输入图像嵌入到补丁中。 Vision Transfomers 使用了非常激进的补丁(16×16),作者使用了 4×4 补丁,实现了 conv 层。准确率从 79.4% 变为 79.5%,表明修补有效。

ResNeXt-ify

ResNetXt 对 BottleNeck 中的 3×3 卷积层采用分组卷积来减少 FLOPS。在 ConvNext 中,他们使用深度卷积(如 MobileNet 和后来的 EfficientNet)。深度卷积是分组卷积,其中组数等于输入通道数。[0]

作者注意到这与 self-attention 中的加权求和操作非常相似,后者仅在空间维度上混合信息。使用 depth-wise convs 会降低精度(因为我们没有像 ResNetXt 那样增加宽度),这是意料之中的。

所以我们将 BottleNeck 块内的 3×3 conv 更改为

倒置瓶颈

我们的 BottleNeck 首先通过 1×1 conv 减少特征,然后应用重的 3×3 conv,最后将特征扩展为原始大小。倒置瓶颈块则相反。我有一整篇文章,对它们进行了很好的可视化。[0]

所以我们从宽 -> 窄 -> 宽到窄 -> 宽 -> 窄。

这与 Transformers 类似,由于 MLP 层遵循窄 -> 宽 -> 窄设计,MLP 中的第二个密集层将输入的特征扩展了四倍。

大内核大小

Modern Vision Transfomer 和 Swin 一样,使用更大的内核大小 (7×7)。增加内核大小将使计算成本更高,因此我们向上移动大的深度卷积,这样做我们将拥有更少的通道。作者指出,这类似于 Transformers 模型,其中多头自我注意 (MSA) 在 MLP 层之前完成。

这将准确度从 79.9% 提高到 80.6%

微设计

用 GELU 替换 ReLU

既然最先进的变压器都使用了 GELU,为什么不在我们的模型中使用它呢?作者报告准确性保持不变。在 nn.GELU 中的 PyTorch GELU 中。

更少的激活函数

我们的块有三个激活函数。而在 Transformer 模块中,只有一个激活函数,即 MLP 模块内部的激活函数。作者删除了所有激活,除了中间卷积层之后的激活。这将匹配 Swin-T 的准确度提高到 81.3%!

更少的归一化层

与激活类似,Transformers 块具有较少的归一化层。作者决定删除所有 BatchNorm,只保留中间转换之前的那个。

用 LN 代替 BN

好吧,他们用 LinearyNorm 代替了 BatchNorm 层。他们注意到在原始 ResNet 中这样做会损害性能,但经过我们所有的更改后,性能提高到 81.5%

所以,让我们应用它们

分离下采样层。

在 ResNet 中,下采样是通过 stride=2 conv 完成的。 Transformers(以及其他卷积网络)也有一个单独的下采样模块。作者删除了 stride=2 并使用 2×2 stride=2 conv 在三个 conv 之前添加了一个下采样块。在下采样操作之前需要进行归一化,以保持训练期间的稳定性。我们可以将此模块添加到我们的 ConvNexStage。最后,我们达到了超过 Swin 的 82.0%!

class ConvNexStage(nn.Sequential):
    def __init__(
        self, in_features: int, out_features: int, depth: int, **kwargs
    ):
        super().__init__(
            # add the downsampler
            nn.Sequential(
                nn.GroupNorm(num_groups=1, num_channels=in_features),
                nn.Conv2d(in_features, out_features, kernel_size=2, stride=2)
            ),
            *[
                BottleNeckBlock(out_features, out_features, **kwargs)
                for _ in range(depth)
            ],
        )

现在我们可以清理我们的 BottleNeckBlock

class BottleNeckBlock(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        expansion: int = 4,
    ):
        super().__init__()
        expanded_features = out_features * expansion
        self.block = nn.Sequential(
            # narrow -> wide (with depth-wise and bigger kernel)
            nn.Conv2d(
                in_features, in_features, kernel_size=7, padding=3, bias=False, groups=in_features
            ),
            # GroupNorm with num_groups=1 is the same as LayerNorm but works for 2D data
            nn.GroupNorm(num_groups=1, num_channels=in_features),
            # wide -> wide 
            nn.Conv2d(in_features, expanded_features, kernel_size=1),
            nn.GELU(),
            # wide -> narrow
            nn.Conv2d(expanded_features, out_features, kernel_size=1),
        )

    def forward(self, x: Tensor) -> Tensor:
        res = x
        x = self.block(x)
        x += res
        return x

我们终于到达了我们的最后一个街区!让我们测试一下

stage = ConvNexStage(32, 62, depth=1)
stage(torch.randn(1, 32, 14, 14)).shape
torch.Size([1, 62, 7, 7])

最后的润色

他们还添加了随机深度,也称为放置路径,(我有一篇关于它的文章)和图层比例。[0]

from torchvision.ops import StochasticDepth

class LayerScaler(nn.Module):
    def __init__(self, init_value: float, dimensions: int):
        super().__init__()
        self.gamma = nn.Parameter(init_value * torch.ones((dimensions)), 
                                    requires_grad=True)
        
    def forward(self, x):
        return self.gamma[None,...,None,None] * x

class BottleNeckBlock(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        expansion: int = 4,
        drop_p: float = .0,
        layer_scaler_init_value: float = 1e-6,
    ):
        super().__init__()
        expanded_features = out_features * expansion
        self.block = nn.Sequential(
            # narrow -> wide (with depth-wise and bigger kernel)
            nn.Conv2d(
                in_features, in_features, kernel_size=7, padding=3, bias=False, groups=in_features
            ),
            # GroupNorm with num_groups=1 is the same as LayerNorm but works for 2D data
            nn.GroupNorm(num_groups=1, num_channels=in_features),
            # wide -> wide 
            nn.Conv2d(in_features, expanded_features, kernel_size=1),
            nn.GELU(),
            # wide -> narrow
            nn.Conv2d(expanded_features, out_features, kernel_size=1),
        )
        self.layer_scaler = LayerScaler(layer_scaler_init_value, out_features)
        self.drop_path = StochasticDepth(drop_p, mode="batch")

        
    def forward(self, x: Tensor) -> Tensor:
        res = x
        x = self.block(x)
        x = self.layer_scaler(x)
        x = self.drop_path(x)
        x += res
        return x

etvoilà🎉我们有它!让我们看看它是否有效!

stage = ConvNexStage(32, 62, depth=1)
stage(torch.randn(1, 32, 14, 14)).shape
torch.Size([1, 62, 7, 7])

极好的!我们需要在编码器中创建丢弃路径概率

class ConvNextEncoder(nn.Module):
    def __init__(
        self,
        in_channels: int,
        stem_features: int,
        depths: List[int],
        widths: List[int],
        drop_p: float = .0,
    ):
        super().__init__()
        self.stem = ConvNextStem(in_channels, stem_features)

        in_out_widths = list(zip(widths, widths[1:]))
        # create drop paths probabilities (one for each stage)
        drop_probs = [x.item() for x in torch.linspace(0, drop_p, sum(depths))] 
        
        self.stages = nn.ModuleList(
            [
                ConvNexStage(stem_features, widths[0], depths[0], drop_p=drop_probs[0]),
                *[
                    ConvNexStage(in_features, out_features, depth, drop_p=drop_p)
                    for (in_features, out_features), depth, drop_p in zip(
                        in_out_widths, depths[1:], drop_probs[1:]
                    )
                ],
            ]
        )
        

    def forward(self, x):
        x = self.stem(x)
        for stage in self.stages:
            x = stage(x)
        return x
image = torch.rand(1, 3, 224, 224)
encoder = ConvNextEncoder(in_channels=3, stem_features=64, depths=[3,4,6,4], widths=[256, 512, 1024, 2048])
encoder(image).shape
torch.Size([1, 2048, 3, 3])

为了获得用于图像分类的最终 ConvNext,我们需要在编码器顶部应用分类头。我们还在最后一个线性层之前添加了一个 LayerNorm。

class ClassificationHead(nn.Sequential):
    def __init__(self, num_channels: int, num_classes: int = 1000):
        super().__init__(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(1),
            nn.LayerNorm(num_channels),
            nn.Linear(num_channels, num_classes)
        )
    
    
class ConvNextForImageClassification(nn.Sequential):
    def __init__(self,  
                 in_channels: int,
                 stem_features: int,
                 depths: List[int],
                 widths: List[int],
                 drop_p: float = .0,
                 num_classes: int = 1000):
        super().__init__()
        self.encoder = ConvNextEncoder(in_channels, stem_features, depths, widths, drop_p)
        self.head = ClassificationHead(widths[-1], num_classes)
image = torch.rand(1, 3, 224, 224)
classifier = ConvNextForImageClassification(in_channels=3, stem_features=64, depths=[3,4,6,4], widths=[256, 512, 1024, 2048])
classifier(image).shape
torch.Size([1, 1000])

在这里你有它!

Conclusions

在本文中,我们一步一步地看到了作者为从 ResNet 创建 ConvNext 所做的所有更改。我希望这很有用:)

感谢您阅读!

Francesco

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
社会演员多的头像社会演员多普通用户
上一篇 2022年4月25日
下一篇 2022年4月25日

相关推荐