在 PyTorch 中实现 ConvNext
一个击败变压器的新卷积网络
你好！！今天我们将在 2020 年代的 A ConvNet 中在 PyTorch 中实现著名的 ConvNext。[0]
让我们开始吧！
论文提出了一种新的基于卷积的架构，不仅超越了基于 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
文章出处登录后可见！
