PyTorch を使用した効率的な画像セグメンテーション: パート 4

1. 説明

        この 4 部構成のシリーズでは、PyTorch の深層学習技術を使用して画像セグメンテーションを基礎から説明します。このセクションでは、ビジュアル トランスフォーマー ベースの画像セグメンテーション モデルを実装する方法に焦点を当てます。

 

図 1: Vision Transformer モデル アーキテクチャを使用して画像セグメンテーションを実行した結果。

        上から下に、入力イメージ、グラウンド トゥルース セグメンテーション マスク、予測セグメンテーション マスクです。出典: 著者

2. 記事の概要

この記事では、ディープ ラーニングの世界を席巻したTransformer アーキテクチャ        について説明します。Transformer は、音声、視覚、音声などのさまざまなモダリティをモデル化できるマルチモーダル アーキテクチャです。

        この記事では、

1. 変圧器のアーキテクチャと関連する重要な概念について学びます
2. Vision Transformer アーキテクチャを理解する
3. すべての構成要素と可動部分を理解できるように、ゼロから作成された Visual Converter モデルの導入
4. このモデルに供給された入力テンソルを追跡し、形状がどのように変化するかを確認します
5. このモデルを使用して、Oxford IIIT Pets データセットで画像セグメンテーションを実行します。
6. このセグメンテーション タスクの結果を観察します。
7. セマンティック セグメンテーションのための最先端のビジュアル トランスフォーマーである SegFormer の簡単な紹介

この記事では、このノートブックのコードと結果をモデルのトレーニングに        参照します。結果を再現したい場合は、最初のノートブックが適切な時間内に実行されるようにするための GPU が必要になります。

3. このシリーズの記事

        このシリーズは、あらゆる深層学習経験レベルの読者を対象としています。ディープ ラーニングとビジュアル AI について、確かな理論と実践的な経験を積んで実際に学びたいなら、ここが最適な場所です。これは次の記事からなる 4 部構成のシリーズになります。

1. コンセプトとアイデア
2. CNN ベースのモデル
3. 深さ方向の分離可能な畳み込み
4. ビジョントランスフォーマーベースのモデル (この論文)

        Transformer アーキテクチャの概要と直感的な理解から、視覚的な Transformer ツアーを始めましょう。

4. トランスのアーキテクチャ

トランスフォーマーのアーキテクチャは、インターリーブされた通信層と計算層        の組み合わせと考えることができます。図 2 は、このアイデアを視覚的に示しています。トランスフォーマーには N 個の処理ユニット (図 3 の N は 2) があり、各ユニットは入力の 1/N 部分の処理を担当します。これらの処理ユニットが意味のある結果を生成するには、各処理ユニットが入力の全体的なビューを持っている必要があります。したがって、システムは各処理ユニット内のデータに関する情報を他のすべての処理ユニットに繰り返し通信します。これは、各処理ユニットから他のすべての処理ユニットへの赤、緑、青の矢印を使用して示されています。次に、この情報に基づいていくつかの計算を示します。このプロセスを十分に繰り返した後、モデルは期待どおりの結果を生み出すことができました。

図 2: トランスにおけるインターリーブ通信と計算。この図では、通信と計算の 2 層のみを示しています。

「必要なのは注意だけ」というタイトルの論文で説明されている        ように、ほとんどのオンライン リソースでは通常、トランスのエンコーダとデコーダについて説明していることに注意してください。ただし、この記事ではトランスのエンコーダ部分のみについて説明します。

        Transformer における通信と計算の構成要素を詳しく見てみましょう。

4.1 変圧器内の通信: 注意

        Transformer では、通信はアテンション レイヤーと呼ばれるレイヤーによって実装されます。PyTorch では、これは MultiHeadtentionと呼ばれます。名前の理由については後ほど説明します。

        ドキュメントには次のように書かれています。

「論文で述べられているように、モデルが異なる表現部分空間からの情報に共同で焦点を当てることができるようになります。必要なのは注意だけです。

        アテンション メカニズムは、形状 (バッチ、長さ、特徴) の入力テンソル xを受け取り、同様の形状のテンソル yを生成します。これにより、各入力の特徴は、同じインスタンス内でテンソルが注目する他の入力に基づいて更新されます。したがって、サイズ「長さ」のインスタンスでは、長さ「特徴」の各テンソルの特徴が、他のすべてのテンソルに関して更新されます。ここで、注意メカニズムの二次コストが登場します。

図 3: 文中の他の単語と比較して表示される単語「it」に対する注意。「それ」が同じ文内の「animal」、「too」、「tire(d)」という単語に注目していることがわかります。 

        ビジョン トランスフォーマーのコンテキストでは、トランスフォーマーへの入力は画像です。これが 128 x 128 (幅、高さ) の画像だとします。それを複数の小さなサイズのチャンク (16 x 16) に分割します。128 x 128 の画像の場合、64 個のパッチ (長さ)、1 行あたり 8 個のパッチ、1 行あたり 8 個のパッチが得られます。

        サイズ 16 x 16 ピクセルのこれらの 64 ブロックのそれぞれは、トランスフォーマー モデルへの個別の入力とみなされます。詳細には立ち入らず、このプロセスは 64 個の異なる処理ユニットによって駆動され、それぞれが単一の 16x16 画像パッチを処理すると考えるだけで十分です。

        各ラウンドで、各処理ユニットのアテンション メカニズムは、担当する画像パッチを調べ、残りの 63 の各処理ユニットに関連性​​があり、独自の処理を効率的に行うのに役立つ情報を問い合わせます。画像パッチ。

        注意によるコミュニケーションのステップの後に計算が続きます。これを次に調べます。

4.2 トランスフォーマーの計算: 多層パーセプトロン

        Transformer での計算は、多層パーセプトロン (MLP) ユニットにすぎません。このユニットは、間に GeLU 非線形性を備えた 2 つの線形層で構成されています。他の非線形性も考慮することができます。このユニットはまず入力を 4 倍のサイズに投影し、次に入力と同じサイズの 1 倍に再投影します。

        ノートブックにあるコードでは、このクラスは多層パーセプトロンと呼ばれます。コードを以下に示します。

class MultiLayerPerceptron(nn.Sequential):
    def __init__(self, embed_size, dropout):
        super().__init__(
            nn.Linear(embed_size, embed_size * 4),
            nn.GELU(),
            nn.Linear(embed_size * 4, embed_size),
            nn.Dropout(p=dropout),
        )
    # end def
# end class

        Transformer アーキテクチャが高レベルでどのように機能するかを理解したので、画像のセグメンテーションを実行する Vision Transformer に焦点を当てましょう。

5. ビジョンコンバーター

        ビジョン トランスフォーマーは、もともと「 16x16 ワードの価値のある画像: 大規模画像認識のためのトランスフォーマー」というタイトルの論文で紹介されました。この記事では、著者らが標準の Transformer アーキテクチャを画像分類の問題にどのように適用したかについて説明します。これは、イメージをサイズ 16x16 のパッチに分割し、各パッチをモデルへの入力トークンとして扱うことによって行われます。Transformer-Encoder モデルにはこれらの入力トークンが供給され、入力イメージのクラスを予測するように求められます。

図 4: 出典:大規模画像認識用のトランスフォーマー。

        私たちの場合、画像のセグメンテーションに興味があります。各ピクセルのターゲット クラスを予測することを目的としているため、これをピクセル レベルの分類タスクと考えることができます。

        バニラ ビジョン コンバーターに小さいながらも重要な変更を加え、MLP ヘッダーを置き換えて、ピクセルごとの分類が MLP ヘッダーによって行われるようにしました。出力には、各パッチによって共有される線形レイヤーがあり、そのセグメンテーション マスクはビジョン トランスフォーマーによって予測されます。この共有線形層は、モデルへの入力として送信された各パッチのセグメンテーション マスクを予測します。

        ビジュアル トランスフォーマーの場合、サイズ 16x16 のパッチは、特定のタイム ステップにおける単一の入力トークンと同等であると見なされます。

図 5: 画像セグメンテーションのための Vision Transformer のエンドツーエンドの動作。このノートブックを使用して生成された画像。

5.1 ビジュアルコンバーターでテンソル次元を構築するための直感

        ディープ CNN を使用する場合、主に使用するテンソル次元は (N、CH、W) です。ここで、文字は以下を表します。

• N: バッチサイズ
• C: チャンネル数
• H: 高さ
• W：幅

        このフォーマットは、画像の特性に非常に特有の匂いがするため、2D 画像処理に特化していることがわかります。

        一方、トランスフォーマーを使用すると、物事はより汎用的になり、ドメインに依存しなくなります。以下に示す内容は、入力データがシーケンスとして表現できるビジョン、テキスト、NLP、オーディオ、またはその他の問題に当てはまります。驚くべきことに、テンソルがビジョン コンバーターを通過する際のテンソルの表現には、ビジョン固有のバイアスがほとんどありません。

        トランスフォーマーを使用する場合、および一般に、テンソルが次の形状を持つことが期待されます: (B、T、C)。ここで、文字は次を表します。

• B：バッチサイズ（CNNと同じ）
• T: 時間次元またはシーケンスの長さ。この寸法は L とも呼ばれます。ビジュアル トランスフォーマーの場合、各画像パッチはこの次元に対応します。画像パッチが 16 個ある場合、T 寸法の値は 16 になります。
• C: チャネルまたは埋め込みサイズの寸法。この寸法は E とも呼ばれます。画像を処理するとき、サイズ 3x16x16 (チャネル、幅、高さ) の各パッチは、パッチ埋め込みレイヤーを介してサイズ C の埋め込みにマッピングされます。これを行う方法については後で説明します。

        セグメンテーション マスクを予測するプロセスで入力イメージ テンソルがどのように変更され、処理されるかを詳しく見てみましょう。

5.2 ビジョンコンバーターにおけるテンソルの旅

        ディープ CNN では、テンソルの過程は次のようになります (UNet、SegNet、またはその他の CNN ベースのアーキテクチャの場合)。

        入力テンソルの形状は通常 (1、3、128、128) です。このテンソルは一連の畳み込みと最大プーリング操作を経て、空間次元が減少し、チャネル次元が通常それぞれ 2 倍ずつ増加します。これを特徴エンコーダと呼びます。この後、逆の操作を実行して、空間次元を増加させ、チャネル次元を減少させます。これを特徴デコーダと呼びます。デコード プロセスの後、形状 (1、64、128、128) のテンソルを取得します。次に、これは、(128, C, 1, 1) として 1x128 の不偏ポイントワイズ畳み込みを使用して、必要な数の出力チャネル C に投影されます。

図 6: 画像セグメンテーションのためのディープ CNN によるテンソル形状の典型的な進行。 

        ビジョントランスフォーマーを使用する場合、プロセスはさらに複雑になります。以下の画像の 1 つを見て、テンソルが各ステップで形状をどのように変換するかを理解してみましょう。

図 7: 画像セグメンテーションのためのビジュアル トランスフォーマーによるテンソル形状の一般的な進行。 

        各ステップをさらに詳しく見て、ビジョン トランスフォーマーを流れるテンソルの形状がどのように更新されるかを見てみましょう。これをよりよく理解するために、テンソル次元の具体的な値を見てみましょう。

1. バッチ正規化:入力テンソルと出力テンソルの形状は (1、3、128、128) です。形状は同じままですが、値はゼロ平均と単位分散に正規化されます。
2. イメージからパッチへ:形状 (1、3、128、128) の入力テンソルは、16x16 イメージのスタック ブロックに変換されます。出力テンソルの形状は (1, 64, 768) です。
3. パッチエンベディング:パッチエンベディングレイヤーは、768 の入力チャンネルを 512 のエンベディングチャンネルにマッピングします (この例では)。出力テンソルの形状は (1, 64, 512) です。パッチ埋め込み層は基本的に単なる NN です。PyTorch の線形レイヤー。
4. 位置エンベディング:位置エンベディング層は入力テンソルを受け取りませんが、パッチエンベディングと同じ形状の学習可能なパラメーター (PyTorch のトレーニング可能なテンソル) を効果的に提供します。これは (1、64、512) の形状です。
5. 追加:パッチと位置のエンベディングが区分的に加算されて、Vision Transformer エンコーダーへの入力が生成されます。このテンソルの形状は (1, 64, 512) です。ビジョン トランスフォーマーの主力であるエンコーダーが、基本的にこのテンソル形状を不変に保つことに気づくでしょう。
6. トランスフォーマー エンコーダー:形状 (1、64、512) の入力テンソルは、それぞれが複数のアテンション ヘッド (通信) を備えた複数のトランスフォーマー エンコーダー ブロックを通過し、その後に MLP 層 (計算) が続きます。テンソルの形状は (1, 64, 512) のように変わりません。
7. 線形出力投影:各イメージが 10 クラスに分割されると仮定すると、サイズ 16x16 のパッチごとに 10 チャネルが必要になります。出力投影の nn.linear 層は、512 個の埋め込みチャネルを 16x16x10 = 2560 個の出力チャネルに変換します。このテンソルは (1, 64, 2560) のようになります。上のグラフでは C' = 10 です。「MLP は汎用関数近似器」であるため、理想的にはこれは多層パーセプトロンになりますが、これは教育的な演習であるため、単一の線形層を使用します。
8. Patch to Image:このレイヤーは、(64, 1, 64) テンソルとしてエンコードされた 2560 個のパッチを、セグメンテーション マスクのようなものに変換して戻します。これは 10 個の単一チャネル イメージ、またはこの場合は単一の 10 チャネル イメージであり、各チャネルは 10 個のクラスのいずれかのセグメンテーション マスクです。出力テンソルの形状は (1、10、128、128) です。

         以上です。Vision Transformer を使用して入力画像をセグメント化することに成功しました。次に、実験と結果を見てみましょう。

5.3 ビジョントランスフォーマーの実用化

        このノートブックには、このセクションのすべてのコードが含まれています。

        コードとクラス構造の点では、上記のブロック図をよく似ています。上で説明した概念のほとんどは、このノートブックのクラス名に 1 対 1 で対応しています。

        モデルの重要なハイパーパラメータであるアテンション レイヤーに関連する概念がいくつかあります。強気派の焦点の詳細については、この記事の範囲を超えているため、先ほどは触れませんでした。Transformers のアテンション メカニズムの基本を理解していない場合は、先に進む前に上記の参考文献を読むことを強くお勧めします。

        セグメンテーションには、Vision Transformer の次のモデル パラメーターを使用します。

1. パッチ埋め込み層の 768 埋め込み次元
2. 12 トランスエンコーダブロック
3. 各トランスエンコーダーブロックに 8 つのアテンションヘッド
4. MLP では多頭注意と 20% のドロップアウト

この構成は、VisionTransformerArgs Python データ クラスで確認できます。

@dataclass
class VisionTransformerArgs:
    """Arguments to the VisionTransformerForSegmentation."""
    image_size: int = 128
    patch_size: int = 16
    in_channels: int = 3
    out_channels: int = 3
    embed_size: int = 768
    num_blocks: int = 12
    num_heads: int = 8
    dropout: float = 0.2
# end class

        モデルのトレーニングと検証中に、以前と同様の構成が使用されました。構成は次のように指定します。

1. ランダムな水平反転とカラー ジッタリングデータ拡張がトレーニング セットに適用され、過剰適合を防止します
2. サイズ変更操作を維持しながら、非アスペクト比で画像を 128x128 ピクセルにサイズ変更します
3. 画像に入力正規化を適用せず、代わりにバッチ正規化レイヤーをモデルの最初のレイヤーとして使用します。
4. モデルは、LR 50.0 の Adam オプティマイザーと、0004 エポックごとに学習率を 0.8 の係数で減衰する StepLR スケジューラーを使用して、12 エポックにわたってトレーニングされます。
5. クロスエントロピー損失関数は、ピクセルをペット、背景、またはペットの境界に属するものとして分類するために使用されます。

        このモデルには 8,628 万個のパラメータがあり、85 トレーニング エポック後の検証精度は 89.50% です。これは、88 トレーニング エポック後にディープ CNN モデルによって達成される 28.20% の精度よりも低くなります。これは、実験的に検証する必要があるいくつかの要因によるものである可能性があります。

1. 最後の出力投影層は単一の NN です。多層パーセプトロンではなく線形パーセプトロン
2. 16x16 のパッチ サイズは大きすぎて、より詳細な詳細をキャプチャすることはできません
3. 研修期間が足りない
4. トレーニング データが不十分 - Transformer モデルを効果的にトレーニングするには、ディープ CNN モデルよりも多くのデータが必要であることが知られています。
5. 学習率が低すぎる

モデルが検証セット内の 21 枚の画像のセグメンテーション マスクを予測する方法を学習する方法を示す gif を描画します。

図 8: 画像セグメンテーション モデルのビジュアル トランスフォーマーによって予測されるセグメンテーション マスクの進行状況を示す gif。 

        私たちは初期のトレーニング期間中に興味深いことに気づきました。予測されたセグメンテーション マスクには、奇妙なブロッキング アーティファクトがいくつかあります。考えられる唯一の理由は、画像をサイズ 16x16 のパッチに分解し、非常に少ないトレーニング エポックの後、モデルは通常ペッティングまたは背景ピクセル オーバーレイであるこの 16x16 パッチについて有益な情報を何も学習しないことです。

図 9: 画像セグメンテーションにビジュアル トランスフォーマーを使用すると、予測セグメンテーションで見られるオクルージョン アーティファクトがマスクされます。 

        基本的なビジュアル トランスフォーマーを見てきたので、セグメンテーション タスク用の最先端のビジュアル トランスフォーマーに注目してみましょう。

5.4 SegFormer: トランスフォーマーを使用したセマンティック セグメンテーション

        この論文では、2021 年の SegFormer アーキテクチャを提案します。上で見たコンバータは、SegFormer アーキテクチャの簡略化されたバージョンです。

図 10: SegFormer アーキテクチャ。ソース： 

        最も注目に値するのは、SegFormer です。

1. サイズ 16x16 のパッチを含む 1 つのパッチ イメージの代わりに、サイズ 4x4、8x8、16x16、および 32x32 のパッチを含む 4 セットのイメージを生成します。
2. トランスフォーマー エンコーダー ブロックを 1 つだけではなく 4 つ使用します。モデルアンサンブルって感じです
3. 自己注意の前段階と後段階で畳み込みを使用する
4. 位置埋め込みを使用しないでください
5. 各変換モジュールは空間解像度 H/4 x W/4、H/8 x W/8、H/16 x W/16 および H/32、W/32 で画像を処理します。
6. 同様に、空間次元が減少するとチャネルも増加します。これはディープ CNN に似た感じです
7. 複数の空間次元で予測をアップサンプリングし、デコーダー内でそれらをマージする
8. MLP はこれらすべての予測を組み合わせて最終的な予測を提供します。
9. 最終的な予測は、H,W ではなく、空間次元 H/4,W/4 で行われます。

6. 結論

このシリーズのパート 4 では、特にトランスフォーマー アーキテクチャとビジョン トランスフォーマーについて説明しました。ビジョン トランスフォーマーの動作方法と、ビジョン トランスフォーマーの通信フェーズと計算フェーズに関与する基本的な構成要素について直感的に理解できます。ビジュアル トランスフォーマーがセグメンテーション マスクを予測し、その予測を結合するために採用した独自のパッチベースのアプローチを確認しました。

ビジュアル トランスフォーマーの動作を示し、深い CNN アプローチとの結果の比較を可能にする実験をレビューします。当社のビジョン コンバーターは最先端のものではありませんが、かなり良好な結果を達成できます。SegFormer などの最先端のメソッドを垣間見ることができます。

Transformer にはディープ CNN ベースのメソッドよりも多くの可動部分があり、より複雑であることはもう明らかです。生の FLOP の観点から見ると、変圧器は効率の向上を約束します。Transformer では、計算量の多い実際の層は nn だけです。線形。これは、ほとんどのアーキテクチャで最適化された行列乗算を使用して実装されます。このアーキテクチャの単純さにより、Transformer はディープ CNN ベースの手法と比較して最適化と高速化が容易になることが期待されています。

ここまで到達できておめでとうございます! PyTorch での効率的な画像セグメンテーションに関するシリーズをお読みいただき、嬉しく思います。ご質問やコメントがございましたら、お気軽にコメント欄に残してください。

7. 多読

注意メカニズムの詳細については、この文書の範囲を超えています。さらに、注意のメカニズムについて詳しく学ぶために、多くの高品質のリソースを参照することができます。強くお勧めするものをいくつかご紹介します。

1. 回路図トランスフォーマー
2. PyTorch を使用して NanoGPT をゼロから作成する

ビジュアル コンバーターについて詳しく説明した記事へのリンクを以下に示します。

1. PyTorch での Vision Transformer (ViT) の実装 : このペーパーでは、 PyTorch での画像分類のための Vision Transformerの実装について詳しく説明します。これらの実装では einops が使用されていることに注意してください。これは教育に重点を置いた演習であるため避けています (コードの可読性を向上させるために、 einopsを学習して使用することをお勧めします  )。代わりに、ネイティブの PyTorch 演算子を使用して、テンソル次元を並べ替え、再配置します。また、作者はいくつかの場所で線形レイヤーの代わりに Conv2d を使用しています。私たちは、畳み込み層をまったく使用しないビジョン トランスフォーマーの実装を構築したいと考えています。
2. ビジョンコンバーター: AIの夏
3. PyTorch での SegFormer の実装

ドリュー・マッタニ

·