Vision TransformerでCIFAR10正解率99%を達成する方法

# はじめに
画像認識の新方式として期待されているVision Transformer(ViT)を使って、CIFAR10正解率99%に挑戦する。
公式のページでもCIFAR10の転移学習ができるColabのノートブックが提供されていて、さほど難しいことでもないが、そのまま実施しても面白くないので、ここではTensorFlow/Kerasの自作コードに学習済みの重みをロードして実行する。

## 環境
TensorFlow 2.4.0
tf.keras 2.3.0
Google Colab TPU

#ViTのモデルコード

主に下記のコードを参考にした。
[Google公式](https://github.com/google-research/vision_transformer/)
[Eunkwang Jeon氏によるPytorch移植](https://github.com/jeonsworld/ViT-pytorch)

Googleの公式コードはフレームワークとして[JAX](https://github.com/google/jax)を使っている。
Pytorch移植版はその忠実な再現で、コードも簡潔でこちらのほうが処理がわかりやすいのでおすすめ。AttentionMapの視覚化もできて素晴らしい。(記事作成後に同じ人がすでに[Keras版](https://github.com/jeonsworld/vit-keras)を作っていると気づいたが、こちらは参照してない。)

今回作成したtf.keras版のコードはこちら。(全文は最後のGoogle Colabのノートブックのリンク参照のこと)


```python
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.layers import  Dense, Dropout, Add, Concatenate, Reshape, Permute, Activation, Flatten, Conv2D, LayerNormalization
class PatchEmbedding(tf.keras.layers.Layer):
    def __init__(self, num_patches, d_model, **kwargs):
        self.num_patches = num_patches
        self.d_model = d_model
        super().__init__(**kwargs)
    def build(self, input_shape):
        self.pos_emb = self.add_weight("pos_emb", 
                                       shape=(1, self.num_patches + 1, self.d_model),
                                       initializer=tf.initializers.RandomNormal(stddev=0.02))
        self.class_emb = self.add_weight("class_emb", 
                                         shape=(1, 1, self.d_model),
                                         initializer=tf.initializers.Zeros())
        super().build(input_shape)
    def call(self,x):
        batch_size = tf.shape(x)[0]
        class_emb = tf.broadcast_to(
            self.class_emb, [batch_size, 1, self.d_model]
        )
        x = tf.concat([class_emb, x], axis=1)
        x = x + self.pos_emb
        return x
    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1]+1, input_shape[2], input_shape[3])


def vit(
    image_size = 224,
    patch_size = 16,
    channels = 3,
    d_model = 768,
    num_classes =21843,
    mlp_dim = 3072,
    num_layers = 12,
    num_heads = 12,
    dropout = 0.1,
    attention_dropout = 0.0):

    dense_kwargs = {
        'kernel_initializer':'glorot_uniform',
        'bias_initializer': tf.keras.initializers.RandomNormal(stddev=1e-6)
        }
    lnorm_kwargs = {'epsilon':1e-6}

    def MultiHead_SelfAttention(inputs, d_model, num_heads, dropout,prefix):
        projection_dim = d_model // num_heads

        query = Dense(d_model, name = prefix+"_Dense_query", **dense_kwargs)(inputs)
        key   = Dense(d_model, name = prefix+"_Dense_key"  , **dense_kwargs)(inputs)
        value = Dense(d_model, name = prefix+"_Dense_value", **dense_kwargs)(inputs)
        query = Reshape((-1, num_heads, projection_dim),name = prefix+"_Reshape_query")(query)
        key   = Reshape((-1, num_heads, projection_dim),name = prefix+"_Reshape_key"  )(key)
        value = Reshape((-1, num_heads, projection_dim),name = prefix+"_Reshape_value")(value)

        query = Permute((2, 1, 3), name = prefix+"_Permute_query")(query)
        key   = Permute((2, 1, 3), name = prefix+"_Permute_key")(key)
        value = Permute((2, 1, 3), name = prefix+"_Permute_value")(value)

        score = tf.matmul(query, key, transpose_b=True)
        score = score/K.sqrt(K.cast(projection_dim, 'float32'))
        attention_probs = Activation("softmax", name = prefix+"_Softmax_projection")(score) 
        if dropout != 0.0:
            attention_probs = Dropout(dropout, name=prefix+"_Dropout_projection")(attention_probs)

        attention = tf.matmul(attention_probs, value)
        attention = Permute((2, 1, 3), name = prefix+"_Permute_attention")(attention)
        attention = Reshape( (-1, d_model), name = prefix+"_Reshape_attention")(attention)
        output = Dense(d_model, name = prefix+"_Dense_out", **dense_kwargs)(attention)
        if dropout != 0.0:
            output = Dropout(dropout, name=prefix+"_Dropout_out")(output)

        return output

    x = inputs = tf.keras.layers.Input(shape=(image_size, image_size, channels), name="Input")
    
    num_patches = (image_size // patch_size) ** 2
    patch_dim = channels * patch_size ** 2
    x = Conv2D(d_model,patch_size,patch_size, padding='valid', name='Embedding_Conv2D')(x)
    x = Reshape((num_patches,d_model), name='Embedding_Reshape')(x)
    x = PatchEmbedding(num_patches,d_model, name="PatchEmbedding")(x)

    if dropout != 0.0:
       x = Dropout(dropout, name='Embedding_Dropout')(x)

    # Encoder
    for i in range(num_layers):
        prefix = f'Encoder_{i}'
        # Attention block.
        shortcut = x
        x = LayerNormalization(name = prefix+"_Attention_LayerNorm")(x)
        x = MultiHead_SelfAttention(x, d_model, num_heads, attention_dropout, prefix=prefix+"_MHSA" )
        x = Add(name=prefix+"_Attention_Add")([x,shortcut])
        # MLP block.
        shortcut = x
        x = LayerNormalization(name=prefix+"_MLP_LayerNorm",**lnorm_kwargs)(x)
        x = Dense(mlp_dim, name=prefix+"_MLP_Dense_0",**dense_kwargs)(x)
        x = Activation(tf.nn.gelu, name=prefix+"_MLP_GeLU")(x)
        if dropout != 0.0:
            x = Dropout(dropout, name=prefix+"_MLP_Dropout")(x)
        x = Dense(d_model, name=prefix+"_MLP_Dense_1", **dense_kwargs)(x)
        if dropout != 0.0:
            x = Dropout(dropout, name=prefix+"_MLP_Dropout_hidden")(x)
        x = Add(name=prefix+"_MLP_Add")([x,shortcut])

    x = LayerNormalization(name="Encoder_Norm", **lnorm_kwargs)(x)
    x = Dense(num_classes, kernel_initializer=tf.keras.initializers.Zeros() , name="Head_Dense")(x[:, 0])
    return tf.keras.Model(inputs,x) 

```

公式が提供する重みのデータをロード可能で、実際に精度が出ることから、コードに大きな瑕疵は無いと思われる。
以下コードを作成した際の個人的な感想など。

* 最初にConv2Dが来る
 * 開発者側としては無理に畳み込みを排除する姿勢はない
* クラスや位置情報のEmbeddingは学習可能な重みが必要なので、Layerを自作せざるを得ない模様
 * そのうちtf.kerasでも実装済みのLayerとして提供されるかもしれない
 * MultiHeadAttentionはtf.kerasで提供されるようなので、この辺ももっと簡単にかけるようになるはず　　
* Transformerの部分は本当に既存の言語処理用の処理とほぼ同じ
 * 参照([言語理解のためのTransformerモデル](https://www.tensorflow.org/tutorials/text/transformer?hl=ja))
* 最後はDense(pytorchでのLinear)がActivation無しで付け加えられているだけ

#ファインチューニング

公式の重みをロードしてファインチューニングを行う。今回はベースモデルの重みを使用。
画像は224x224にリサイズして入力とする。最後のDenseのみクラス数が違うので重みはロードせず0で初期化して、その後は普通の学習と同じようにfitさせれば良い。フリーズ等は特にしなくても良いようだ。
パラメータ数は比較的多いが、計算量自体はさほど大きくないようで、割と早くEpochが終わる。メモリ消費もそれほどでもなく、BatchSizeは512でも可能。
今回の各種設定は下記の通り。

| パラメータ | 値 |
|:--|:--|
| Model Name| ViT-B_16 |
| Batch Size | 200  |
| Warmup | 3 epochs(Linear)  |
| Cooldown | 12 epochs(Cosine)  |
| Horizontal/Vertical Shift | 0.2  |
| Horizontal Flip | True  |
| Cutout | 0.5 x 1  |
| Optimizer | SGD(momentum=0.9)|
| Learning Rate | 4e-2 |
| Label Smoothing | 0.1  |

15エポック（ColabのTPUで約20分）で大体99%までいける。もう少しエポック数を増やせば確実性が上がる。
以下、実際の学習の結果。


```
Epoch 1/15
250/250 [==============================] - 149s 321ms/step - loss: 1.6647 - acc: 0.8384 - val_loss: 0.5657 - val_acc: 0.9792
Epoch 2/15
250/250 [==============================] - 72s 284ms/step - loss: 0.6631 - acc: 0.9336 - val_loss: 0.5473 - val_acc: 0.9824
Epoch 3/15
250/250 [==============================] - 71s 283ms/step - loss: 0.6373 - acc: 0.9403 - val_loss: 0.5521 - val_acc: 0.9789
Epoch 4/15
250/250 [==============================] - 71s 282ms/step - loss: 0.6282 - acc: 0.9449 - val_loss: 0.5484 - val_acc: 0.9797
Epoch 5/15
250/250 [==============================] - 71s 283ms/step - loss: 0.6055 - acc: 0.9547 - val_loss: 0.5419 - val_acc: 0.9837
Epoch 6/15
250/250 [==============================] - 71s 282ms/step - loss: 0.5913 - acc: 0.9616 - val_loss: 0.5344 - val_acc: 0.9868
Epoch 7/15
250/250 [==============================] - 71s 282ms/step - loss: 0.5795 - acc: 0.9660 - val_loss: 0.5325 - val_acc: 0.9868
Epoch 8/15
250/250 [==============================] - 71s 282ms/step - loss: 0.5699 - acc: 0.9690 - val_loss: 0.5329 - val_acc: 0.9873
Epoch 9/15
250/250 [==============================] - 71s 282ms/step - loss: 0.5639 - acc: 0.9731 - val_loss: 0.5321 - val_acc: 0.9867
Epoch 10/15
250/250 [==============================] - 71s 282ms/step - loss: 0.5584 - acc: 0.9743 - val_loss: 0.5318 - val_acc: 0.9867
Epoch 11/15
250/250 [==============================] - 71s 282ms/step - loss: 0.5476 - acc: 0.9802 - val_loss: 0.5303 - val_acc: 0.9884
Epoch 12/15
250/250 [==============================] - 71s 282ms/step - loss: 0.5452 - acc: 0.9808 - val_loss: 0.5292 - val_acc: 0.9893
Epoch 13/15
250/250 [==============================] - 72s 286ms/step - loss: 0.5385 - acc: 0.9835 - val_loss: 0.5279 - val_acc: 0.9898
Epoch 14/15
250/250 [==============================] - 71s 283ms/step - loss: 0.5364 - acc: 0.9846 - val_loss: 0.5281 - val_acc: 0.9895
Epoch 15/15
250/250 [==============================] - 71s 283ms/step - loss: 0.5385 - acc: 0.9836 - val_loss: 0.5282 - val_acc: 0.9900
```

[Google Colab Notebookへのリンク](https://gist.github.com/T-STAR-LTD/5c78f4a460541191a68708c29b3b2e75)

# まとめ

以前に[別の記事](https://qiita.com/T-STAR/items/de933c67b52a47f9efef)でEfficentNetで99%達成する方法を書いたが、あちらでは２時間以上かかった。工夫すればもう少し短くできるが結局１時間は切れなかった。それに比べれば大幅に短い時間で達成できたので驚いた。
現状のViTの欠点としては、転移学習を使わないと低性能に終わる、ということが当初から指摘されている。Qiitaでもいくつか[記事](https://qiita.com/hitonimonarukurohyo/items/b2812188d3ad24a1b629)があるが、今回作成したモデルでも最初からCIFAR10を学習させると80%にすら到達できなかった[^1]。この辺りは、[T2T-ViT](https://arxiv.org/abs/2101.11986)のような改良版が出てきているようなので、今後に期待。


# 参考

[画像認識の大革命。AI界で話題爆発中の「Vision Transformer」を解説！](https://qiita.com/omiita/items/0049ade809c4817670d7)
[Vision Transformerのコードの写経](https://qiita.com/nishiha/items/9b02ef14f4fb1cd86a26)

[^1]: その後、Cutmixでデータ拡張子し800エポック程度学習させて88%まで到達できた。