[論文の解釈] FastViT: 構造再パラメータ化を使用した高速ハイブリッドビジョントランスフォーマー

1 はじめに

最先端のレイテンシと精度のトレードオフ

レイテンシは、広く使用されている 2 つのプラットフォーム (モバイル GPU とデスクトップ GPU) で最速です。

1.1 これまでの問題点

スキップ接続では、メモリアクセスコストが増加するため、レイテンシーの点で大きなオーバーヘッドが生じます。

1.1.1 スキップ接続によりメモリアクセスコストが増加するのはなぜですか?

これは主に次の 2 つの要因に関係しています。

メモリ要件: スキップ結合では追加の機能マップまたはテンソルが導入されるため、モデルのメモリ使用量が増加します。各スキップ接続では特定の数の機能マップを保存する必要があるため、メモリ使用量が増加する可能性があります。通常、この増加はモデルの全体的なメモリ使用量に大きな影響を与えませんが、場合によっては追加のメモリ負荷が発生する可能性があります。
メモリアクセス: ディープニューラルネットワークでは、特徴マップの計算と保存に大量のメモリアクセスが必要です。スキップ接続により、読み取りおよび書き込みが必要な機能マップの数が増加し、メモリアクセスコストの増加につながる可能性があります。

1.2 FastViT の改善

スキップ接続を削除するために、完全に再パラメータ化可能な RepMixer を導入しました。
すべての密な k×k 畳み込みを、その分解されたバージョン、つまりチャネルごとの畳み込みとポイントごとの畳み込みに置き換えます。線形トレイン時間オーバーパラメータ化が使用されます。これらの追加のブランチはトレーニング中にのみ導入され、推論時に再パラメータ化されます。
大規模な畳み込みカーネルを使用して、初期段階でセルフアテンション層を置き換えます

1.3 前提知識

トレーニングと推論の分離
パラメータ化
大規模なカーネルコンボリューションを使用する利点 (セルフアテンションレイヤーの代わりに大規模なカーネルコンボリューションを使用できる理由)

2.ファストビタミン

(a) トレーニング時間のアーキテクチャと推論時間のアーキテクチャを分離する FastViT アーキテクチャの概要。ステージ 1、2、および 3 は同じアーキテクチャを持ち、トークン混合に RepMixer を使用します。ステージ 4 では、セルフアテンションレイヤーがトークンの混合に使用されます。(b) 畳み込みシステムの構造。(c) convolution-ffn のアーキテクチャ (d) RepMixer ブロックの概要。推論時にスキップ接続を再パラメータ化します。

基本的な内容はすでに上の図にあります

[論文解釈の参考]モデル圧縮の解釈 24: FastViT: 高速畳み込み変換器のハイブリッドビジョンアーキテクチャ - Zhihu (zhihu.com)

3.コード

3.1 異なるモデル

FastViT には合計 6 つのモデルがあります。

3.2 RepCPE

fastvit_sa12、fastvit_sa24、fastvit_sa36、fastvit_ma36 およびその他のモデル

pos_embs = [None, None, None, partial(RepCPE, spatial_shape=(7, 7))]

推論中に、元の残差ブランチを削除し、前方推論関数のみを調べます。

#reparam_conv 和 self.pe其实是参数基本相同的卷积
 def forward(self, x: torch.Tensor) -> torch.Tensor:
        if hasattr(self, "reparam_conv"):
            x = self.reparam_conv(x)
            return x
        else:
            x = self.pe(x) + x
            return x

3.3 基本ブロック

token_mixer には RepMixerBlock と AttendanceBlock の 2 つのオプションがありますが、一般に AttendanceBlock は最後の層構造でのみ使用されます。

3.3.1 RepMixerブロック

他のミキサー構造と同様

3.3.1.1 レップミキサー

(1) 研修体制

残差構造は、分岐構造のある MobileOneBlock から分岐構造のない MobileOneBlock を減算することによって実装されます。

x = x + self.mixer(x) - self.norm(x)

(2) 重いパラメータ

重いパラメータの原理は RepVGG と同じですが、詳細は説明しません。

3.3.1.2 ConvFFN

実際、これは比較的単純であり、ConvFFN が複雑なパラメータ化を実行しないように見えることは言及する価値があります。

3.4 パッチ埋め込み

違いは一般的なもので、ReparamLargeKernelConv と MobileOneBlock を介したマッピングです。

ReparamLargeKernelConv からのカーネルの 31x31 へのスケールアップ: CNN における大規模カーネル設計の再考https://openaccess.thecvf.com/content/CVPR2022/papers/Ding_Scaling_Up_Your_Kernels_to_31x31_Revisiting_Large_Kernel_Design_CVPR_2022_paper.pdf

class PatchEmbed(nn.Module):
    """Convolutional patch embedding layer."""

    def __init__(
        self,
        patch_size: int,
        stride: int,
        in_channels: int,
        embed_dim: int,
        inference_mode: bool = False,
    ) -> None:

        super().__init__()
        block = list()
        block.append(
            ReparamLargeKernelConv(
                in_channels=in_channels,
                out_channels=embed_dim,
                kernel_size=patch_size,
                stride=stride,
                groups=in_channels,
                small_kernel=3,
                inference_mode=inference_mode,
            )
        )
        block.append(
            MobileOneBlock(
                in_channels=embed_dim,
                out_channels=embed_dim,
                kernel_size=1,
                stride=1,
                padding=0,
                groups=1,
                inference_mode=inference_mode,
                use_se=False,
                num_conv_branches=1,
            )
        )
        self.proj = nn.Sequential(*block)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj(x)
        return x