CVPR 2022|从原理和代码详解FAIR的惊艳之作:全新的纯卷积模型ConvNeXt
本文首發于極市平臺,作者科技猛獸,轉載請獲得授權并標明出處。
本文目錄
7 匹敵 Transformer 的2020年代的卷積網絡
(來自 FAIR,UCB)
7.1 ConvNeXt 原理分析
7.2 ConvNeXt 代碼解讀
7 匹敵 Transformer 的2020年代的卷積網絡
論文名稱:A ConvNet for the 2020s
論文地址:
https://arxiv.org/pdf/2201.03545.pdf
7.1 ConvNeXt 原理分析
7.1.1 Motivation
回顧2010年代,這十年來以深度學習取得了巨大的進步和影響力。主要的驅動力是神經網絡的復興,尤其是卷積神經網絡 (ConvNets)。在過去的十年里,視覺識別領域成功地從設計特征提取器轉移到設計神經網絡架構。盡管反向傳播的發明可以追溯到20世紀80年代,但直到2012年末,我們才看到它真正的潛力。 AlexNet 的發明誕生了 “ImageNet moment”,開啟了計算機視覺的新時代。 此后,該領域發展迅速。像 VGGNet,Inceptions,ResNe(X)t,DenseNet,MobileNet,EfficientNet 和 RegNet 這樣的代表性 ConvNets 側重于精度、效率和可擴展性的不同方面,并推廣了許多有用的設計原則 (design principles)。ConvNets 本身是計算高效的,因為它以滑動窗口的方式 (sliding-window manner) 進行計算,使得計算是共享的。ConvNets 的默認應用場景包括數字識別,人臉識別,行人識別等等。
與此同時, 自然語言處理 (NLP) 的神經網絡設計走了一條非常不同的道路,因為 Transformer 模型取代了 RNN,成為主導的主干架構。
來到2020年代,盡管語言和視覺領域之間的任務存在差異,但這兩個領域的主干架構卻出人意料地融合在一起,改變了網絡設計的面貌。在這些視覺識別模型中,Vision Transformer 是研究熱點,在識別任務上首次超過了卷積模型。除了一開始的圖片分塊操作,原始的 ViT 結構沒有引入任何歸納偏置。雖然圖像識別任務的這些結果十分令人鼓舞,但是計算機視覺不限于圖像分類。原始的 ViT 結構在檢測,分割等等通用性視覺任務上收到輸入圖片分辨率導致的計算復雜度限制,最大的挑戰是 ViT 的全局注意力設計 (global attention design),它相對于輸入圖片的大小 NNN 具有二次復雜度 O(N2)O(N^2)O(N2) 。這對于 ImageNet 分類任務來講可能是可以接受的,但是對于更高分辨率的輸入來說會很快變得難以處理。
金字塔結構 (Hierarchical Transformers) 這種卷積模型先驗 (如 Swin-T, PVT 等) 的引入解決了這一問題,使得 ViT 可以被用做其他視覺任務的骨干網絡。Swin Transformer 使得 “滑動窗口” 策略被重新引入 Transformer 模型,使它們的行為更類似于卷積模型。Swin 的成功也揭示了卷積的本質并沒有變得無關緊要。相反,它仍然備受期待,從未褪色。然而,這種混合方法的有效性仍然很大程度上歸功于 ViT 模型的內在優勢,而不是卷積固有的歸納偏置 (inductive bias)。
ConvNets 和 Hierarchical Transformers 都具備相似的歸納偏置,但在訓練過程和宏/微觀層次的架構設計 (macro/micro-level architecture design) 上有顯著的差異。在這項工作中,作者重新檢查和審視了卷積模型的設計空間,這項研究旨在彌合 pre-ViT 時代和 post-ViT 時代模型性能上的差距,并想要探索出一個純的卷積網絡所能夠達到的性能極限。
為了做到這一點,作者從一個標準的 ResNet (ResNet-50) 開始,用改進的方法進行訓練,逐漸將架構 “現代化 (modernize)”。這個研究想回答的問題是: ViT 模型中的設計決策會如何影響 ConvNets 的性能?
作者現了導致性能差異的幾個關鍵因素。因此,提出了一個名為 ConvNeXt 的系列純卷積模型。作者在各種視覺任務上評估 ConvNeXts,例如 ImageNet 圖像識別,COCO 上的物體檢測/分割,ADE20K 上的語義分割。作者希望這些新的觀察和討論可以鼓勵人們重新思考卷積在計算機視覺中的重要性。
7.1.2 2020年代的卷積網絡
在本文中,作者提出了一個從 ResNet 到與 Transformer 類似的 ConvNet 的 trajectory。作者根據計算了提出了兩種不同的大小范式:FLOPs 約為 4.5×e9 的 ResNet-50 / Swin-T 和 FLOPs 約為 15.0×e9 的 ResNet-200 / Swin-B。從 ResNet-50 開始,以下所有的模型都在 ImageNet-1k 上進行訓練和驗證。
7.1.3 訓練策略
除了網絡架構的設計,訓練過程也會影響最終的性能。Vision Transformer 不僅僅帶來了一套新的模塊和架構設計策略,也引入了不同的訓練技術 (例如 AdamW 優化器)。 這主要與優化策略和相關的超參數設置有關。因此,作者探索的第一步是用 Vision Transformer 的訓練策略訓練一個 ResNet50/200 基線模型。 在研究中,作者使用了一個接近 DeiT 和 Swin Transformer 的訓練方法,訓練輪數從90 epochs 提升到了 300 epochs。使用了AdamW 優化器,數據增強技術包括 Mixup, Cutmix, RandAugment, Random Erasing。正則化方案包括 Stochastic Depth 和 Label Smoothing,如下圖2所示。
這種增強的訓練方案將 ResNet-50 模型的性能從76.1% 提高到了78.8% (+2.7%),這意味著傳統 ConvNets 和Vision Transformer 之間的性能差異的很大一部分可能是由于訓練策略的不同所造成的。
7.1.4 宏觀設計
Swin Transformer 遵循 ConvNet 使用多階段設計,其中每個階段都有不同的特征分辨率。作者借鑒了Swin-T的兩個設計:
小的 Swin Transformer 不同 stage 的層數之比是1:1:3:1,大的 Swin Transformer 不同 stage 的層數之比是1:1:9:1。根據這個設計,作者將每個階段的塊數從 ResNet-50 中的 (3,4,6,3) 調整為 (3,3,9,3) ,這也將 FLOPs 與Swin-T 對齊。這將模型精度從78.8%提高到79.4%。
第2點,作者不再使用 k=7,s=2k=7, s=2k=7,s=2 的卷積加上 MaxPooling 進行下采樣了,而是像 Swin-T 模型那樣把圖片分成4×4 的 Patches (具體是通過 k=4,s=4k=4, s=4k=4,s=4 的一層卷積對輸入圖片進行下采樣),這一步稱為 “Changing stem to Patchify”。這樣每次卷積操作的感受野不重疊,準確率由79.4%提升至79.5%。
7.1.5 模仿 ResNeXt 的設計
在這一部分中,作者試圖采用 ResNeXt 的思想,它比普通的 ResNet 具有更好的 FLOPs/Accuracy Trade-off。核心部分是分組卷積 (grouped convolution),其中卷積濾波器被分成不同的組。具體在這里使用的是 group 數與 channel 數相等的 Depth-wise Convolution。作者注意到 Depth-wise Convolution 類似于 Self-attention 中的加權和運算,即,僅混合空間維度中的信息。深度卷積的使用有效地降低了網絡 FLOPs,但是也降低了精度。按照 ResNeXt 中提出的策略,將網絡寬度增加到與 Swin-T 相同的 channel 數 (從64增加到96)。隨著 FLOPs (5.3G) 的增加,網絡性能提高到80.5%。
7.1.6 Inverted Bottleneck
每個 Transformer Block 中的一個重要設計是,MLP 層的 hidden dimension 比 Input dimension 寬四倍。有趣的是,這種 hidden dimension 比 Input dimension 寬四倍的設計與 MobileNet 的 inverted bottleneck 設計很類似。所以作者也采用了這種辦法,將 block 由下圖3(a)變為3(b)。 這一變化將整個網絡的浮點運算量降至4.6G,這是由于兩個1×1 卷積層的 FLOPs 顯著減少。有趣的是,這導致性能略有提高 (80.5%到80.6%)。在ResNet-200/Swin-B體系中,這一步帶來了更大的增益 (81.9%至82.6%),同時還降低了 FLOPs。
7.1.7 增大卷積kernel
作者認為更大的感受野是 ViT 性能更好的可能原因之一,作者嘗試增大卷積的kernel,使模型獲得更大的感受野。 Swin Transformers 的 Window 大小至少是7×7的,明顯大于 3×3 的 ResNe(X)t 卷積核大小。
首先是把計算復雜度比較高的 depthwise conv layer 往前移動,將 block 由下圖3(b)變為3(c)。使得復雜的模塊(MSA,大內核 conv) 將有更少的channel,而高效、密集的1×1層將有更多的channel。這一中間步驟將 FLOPs降至4.1G,導致性能暫時下降至79.9%。
接下來就是采用更大的卷積核大小。作者試驗了幾種內核大小,包括3,5,7,9,11。網絡的性能從79.9% (3×3)提高到80.6% (7×7),而網絡的 FLOPs 基本保持不變。結果顯示使用7×7的卷積核是最優的。
7.1.8 微觀設計
- 將 ReLU 替換為 GELU: GELU,可以被認為是 ReLU 的平滑變體,在最先進的變形金剛中使用,包括谷歌的BERT 和 OpenAI 的 GPT-2 ,以及最近的 ViT。在的 ConvNet 中,ReLU 也可以用 GELU 代替,但是精度保持不變 (80.6%)。
- 更少的激活函數和歸一化層 (僅在1×1卷積之間使用激活函數,僅在7×7卷積和1×1卷積之間使用歸一化層): Transformer 和 ResNet Block 之間的一個區別是 Transformer 的激活函數較少。MLP 塊只有一個激活函數。相比之下,卷積網絡的激活函數較多,每個卷積層都有一個激活函數,包括1×1卷積層。 作者從殘差塊中消除了所有的 GELU 層,復制了 Transformer 的樣式。這個過程將結果提高到81.3%,實際上與 Swin-T 的性能相當,如下圖4所示。Transformer Block 通常也具有較少的標準化層。這里作者去掉了兩個 BN 層,在 1×1 卷積層之前只剩下一個 BN 層。這個過程將結果提高到81.4%。
- BN 替換為 LN: 作者觀察到,ConvNet 模型在訓練時使用 LN 沒有任何困難;事實上將結果提高到81.5%。
- 將下采樣層單獨分離出來: 最后仿照 Swin-T,作者將下采樣層單獨分離出來,單獨使用2×2卷積層進行下采樣。為保證收斂,在下采樣后加上 Layer Norm 歸一化。最終加強版 ResNet-50 準確率82.0% (FLOPs 4.5G)。
ResNet-50,ConvNeXt-T 和 Swin-T 的結構差別如下圖4,5:
ConvNeXt 變體,ConvNeXt-T/S/B/L 與 Swin-T/S/B/L 的復雜度相似,不同大小的模型配置如下:
實驗結果
Training on ImageNet-1K
300 epochs,優化器:AdamW,初始學習率:4e-3,linear warmup:20 epochs,學習率變化策略:cosine decaying schedule,batch size:4096,weight decay:0.05,EMA。
數據增強:Mixup , Cutmix , RandAugment , 和 Random Erasing。
Pre-training on ImageNet-22K
90 epochs,優化器:AdamW,初始學習率:4e-3,linear warmup:5 epochs,學習率變化策略:cosine decaying schedule,weight decay:0.05。
Fine-tuning on ImageNet-1K
30 epochs,優化器:AdamW,初始學習率:5e-5,學習率變化策略:layer-wise learning rate decay,batch size:512,weight decay:1e-8,EMA。
從實驗結果上看 ConvNeXt-T 性能收益較為突出,在 Accuracy-Computation Trade-off 以及 Inference throughputs 方面,ConvNeXt 與兩個強大的 ConvNet 基線模型 (RegNet 和 EfficientNet) 相當。整體表現也優于具有類似復雜性的 Swin Transformer。
Object detection and segmentation on COCO
作者在 COCO 數據集上微調了帶有 ConvNeXt backbone 的 Mask R-CNN 和 Cascade Mask R-CNN 的檢測模型,訓練策略是 multi-scale training, AdamW optimizer, 3x schedule,結果如下圖7所示。對于不同復雜度的模型,ConvNeXt 都實現了比 Swin Transformer 更好的性能。
Semantic segmentation on ADE20K
作者還用 UperNet 評估了 ADE20K 語義分割任務中的 ConvNeXt 主干模型。 所有模型都經過 160K iterations 的訓練,Batch size 為16,結果如下圖8所示。 ConvNeXt 模型可以在不同的模型容量上實現有競爭力的性能。
7.2 ConvNeXt 代碼解讀
代碼來自:
1 ConvNeXt Block
有2種實現方案:
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
二者的區別是第1種的 LN 是 channels_first 的,第2種的 LN 是 channels_last 的,實測第2種實現方案速度更快些。
class Block(nn.Module):r""" ConvNeXt Block. There are two equivalent implementations:(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute backWe use (2) as we find it slightly faster in PyTorchArgs:dim (int): Number of input channels.drop_path (float): Stochastic depth rate. Default: 0.0layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6."""def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):super().__init__()self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise convself.norm = LayerNorm(dim, eps=1e-6)self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layersself.act = nn.GELU()self.pwconv2 = nn.Linear(4 * dim, dim)self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True) if layer_scale_init_value > 0 else Noneself.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()def forward(self, x):input = xx = self.dwconv(x)x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)x = self.norm(x)x = self.pwconv1(x)x = self.act(x)x = self.pwconv2(x)if self.gamma is not None:x = self.gamma * xx = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)x = input + self.drop_path(x)return x這里需要注意的是 PyTorch 的 nn.LayerNorm 沒辦法直接對維度是 (N, H, W, C) 的張量使用,PyTorch 的 nn.LayerNorm 有2種情況:
# NLP Example batch, sentence_length, embedding_dim = 20, 5, 10 embedding = torch.randn(batch, sentence_length, embedding_dim) layer_norm = nn.LayerNorm(embedding_dim) # Activate module layer_norm(embedding) # Image Example N, C, H, W = 20, 5, 10, 10 input = torch.randn(N, C, H, W) # Normalize over the last three dimensions (i.e. the channel and spatial dimensions) # as shown in the image below layer_norm = nn.LayerNorm([C, H, W]) output = layer_norm(input)可以參考 PyTorch 的官方網站:
也就是說輸入張量的維度要么是 N, C, H, W,那么使用時需要 nn.LayerNorm([C, H, W]),輸入張量的維度要么是 N, L, D,那么使用時需要 nn.LayerNorm(D)。
所以對于輸入維度是 的張量,就需要使用:
F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
如果是 channels_last 數據結構,就直接 F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)。
如果是 channels_first 數據結構,就按照下面的代碼形式手動計算 LN。
class LayerNorm(nn.Module):r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width)."""def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):super().__init__()self.weight = nn.Parameter(torch.ones(normalized_shape))self.bias = nn.Parameter(torch.zeros(normalized_shape))self.eps = epsself.data_format = data_formatif self.data_format not in ["channels_last", "channels_first"]:raise NotImplementedError self.normalized_shape = (normalized_shape, )def forward(self, x):if self.data_format == "channels_last":return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)elif self.data_format == "channels_first":u = x.mean(1, keepdim=True)s = (x - u).pow(2).mean(1, keepdim=True)x = (x - u) / torch.sqrt(s + self.eps)x = self.weight[:, None, None] * x + self.bias[:, None, None]return x2 ConvNeXt 整體結構
class ConvNeXt(nn.Module):r""" ConvNeXtA PyTorch impl of : `A ConvNet for the 2020s` -https://arxiv.org/pdf/2201.03545.pdfArgs:in_chans (int): Number of input image channels. Default: 3num_classes (int): Number of classes for classification head. Default: 1000depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]drop_path_rate (float): Stochastic depth rate. Default: 0.layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1."""def __init__(self, in_chans=3, num_classes=1000, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,):super().__init__()self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layersstem = nn.Sequential(nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),LayerNorm(dims[0], eps=1e-6, data_format="channels_first"))self.downsample_layers.append(stem)for i in range(3):downsample_layer = nn.Sequential(LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),)self.downsample_layers.append(downsample_layer)self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocksdp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] cur = 0for i in range(4):stage = nn.Sequential(*[Block(dim=dims[i], drop_path=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])])self.stages.append(stage)cur += depths[i]self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layerself.head = nn.Linear(dims[-1], num_classes)self.apply(self._init_weights)self.head.weight.data.mul_(head_init_scale)self.head.bias.data.mul_(head_init_scale)def _init_weights(self, m):if isinstance(m, (nn.Conv2d, nn.Linear)):trunc_normal_(m.weight, std=.02)nn.init.constant_(m.bias, 0)def forward_features(self, x):for i in range(4):x = self.downsample_layers[i](x)x = self.stages[i](x)return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C)def forward(self, x):x = self.forward_features(x)x = self.head(x)return x3 不同架構的 ConvNeXt
@register_model def convnext_tiny(pretrained=False, **kwargs):model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)if pretrained:url = model_urls['convnext_tiny_1k']checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)model.load_state_dict(checkpoint["model"])return model@register_model def convnext_small(pretrained=False, **kwargs):model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)if pretrained:url = model_urls['convnext_small_1k']checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")model.load_state_dict(checkpoint["model"])return model@register_model def convnext_base(pretrained=False, in_22k=False, **kwargs):model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)if pretrained:url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")model.load_state_dict(checkpoint["model"])return model@register_model def convnext_large(pretrained=False, in_22k=False, **kwargs):model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)if pretrained:url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")model.load_state_dict(checkpoint["model"])return model@register_model def convnext_xlarge(pretrained=False, in_22k=False, **kwargs):model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)if pretrained:assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"url = model_urls['convnext_xlarge_22k']checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")model.load_state_dict(checkpoint["model"])return model總結
ConvNeXt 可以看做是把 Swin Transformer 包括 ViT 的所有特殊的設計 (包括結構,訓練策略等等) 集于一身之后的卷積網絡進化版,升級了 ResNet 架構,看看借助了2020年代 CV 設計范式之后的卷積網絡的性能極限在哪里。其中 Swin Transformer 的 Self-Attention 層可以和 ConvNeXt 的 DW Conv 等價,所以作者將自注意力層替換為 DW Conv 模塊,其他部分和 Swin Transformer 盡量保持一致 (金字塔結構的 details,1×1,GeLU 激活函數,Layer Normalization 等等)。作者從 ResNet-50 開始,逐步 Swin Transformer 化,最終得到了ConvNeXt-T 模型,性能超過了 Swin-T。ConvNeXt 在不同的 FLOPs 均可以超過 Swin,如果采用ImageNet21K 預訓練后,模型性能有進一步的提升。ConvNeXt 是一個很好的工作,且它在吞吐量上的優勢和魯棒性使得 ConvNeXt 在工業部署上更有價值。我們想,不同神經網絡架構的性能不同,并不是某一兩種結構所帶來的,而是神經網絡的整體架構 (激活函數,歸一化,金字塔不同階段的配置等等) 和訓練方式 (優化器,學習率,數據增強,數據集分辨率,數據集大小) 共同造成的。從這一點上而言,ConvNeXt 給我們帶來了很好的設計范例和參考。
總結
以上是生活随笔為你收集整理的CVPR 2022|从原理和代码详解FAIR的惊艳之作:全新的纯卷积模型ConvNeXt的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: CVPR 2022 57 篇论文分方向整
- 下一篇: 小目标检测、图像分类、图像识别等开源数据