BigTransfer (BiT) ファインチューニング

- 公式ブログ
  https://blog.tensorflow.org/2020/05/bigtransfer-bit-state-of-art-transfer-learning-computer-vision.html

- 公式サンプルコード (TF2, PyTorch, JAX)
  https://github.com/google-research/big_transfer/tree/master/colabs

ここから、BiTファインチューニング時に必要な情報（自分が遊んだときに使った／調べたとこ）を抜き出し、日本語でまとめています。



## 本文

- **事前学習済モデル** を参照・取得できるURL
  Cloud bucketには h5, npz形式が保存されている。TF以外でも使える。
  - TensorFlow Hub  
    https://tfhub.dev/google/collections/bit/1
  - Cloud bucket  
    https://storage.googleapis.com/bit_models/

- BiTモデルについてる **S, M, L** の意味
  ⇒ 学習させたデータセットの違い。**L** は非公開。

|  モデル名  |  データセット                 |
|  :--:     |  :---:                      |
|  BiT-S    |  ILSVRC-2012 (1.3M images)  |
|  BiT-M    |  ImageNet-21k (14M images)  |
|  BiT-L    |  JFT (300M images)          |

- BiTモデルについてる **R-??x?** の意味  
  BiTモデルはResNetを利用しているので、そこの情報。
  - 読み方  
    R50x3 → 50層のResNetで、各層の幅が通常の3倍。
  - パラメータ数  
    TFHubで公開されているモデル（出力層を除く）が持つパラメータ数の概数を示します。  

|  ResNet  |  パラメータ数（概数）  |
|  :----:  |  :----:  |
|  R50x1   |   23M    |
|  R101x1  |   42M    |
|  R50x3   |  211M    |
|  R101x3  |  381M    |
|  R152x4  |  928M    |

次のコードで、パラメータ数を確認しています。

```python
import tensorflow as tf
import tensorflow_hub as tfhub
model = tfhub.KerasLayer('https://tfhub.dev/google/bit/s-r50x1/1')
print(sum(tf.math.reduce_prod(w.shape).numpy() for w in model.weights))
```

### BiT-HyperRule

BiTをファインチューニングするために用意されているヒューリスティックな方法。「これ使えば、一発でイイ感じになる」

もちろんHyper-parameter探索をすれば、より良いモデルを得られる可能性はあります。コスト見合いです。

#### データセットの画像サイズに依存　→ リサイズ、クロップサイズ

データセットの画像サイズで、指定サイズにリサイズ・ランダムクロップする。
対応表は、公式ブログTable1から引用。
![BigTransfer (BiT): State-of-the-art transfer learning for computer vision; Table1](https://4.bp.blogspot.com/-54vXxLE1bqU/XsMCth3HvzI/AAAAAAAADEs/UVrCWT0o6wYPtCpat81ApNGPus16-CHzgCLcBGAsYHQ/s1600/table1.jpg)
併せて、ランダムに左右反転も入れる。バリデーション用データにはリサイズだけ行えば良い。

TF2だとこんな雰囲気になるかと。

```python
def augmentation(image, label):
    image = tf.cast(image, tf.float32) / 255.0
    image = tf.image.resize(image, [512, 512], method=tf.image.ResizeMethod.BILINEAR)
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_crop(image, [480, 480, 3])
    return image, label

ds_train: tf.data.Dataset
ds_train = (ds_train
            .shuffle(1024)
            .repeat()
            .map(augmentation, tf.data.experimental.AUTOTUNE)
            .batch(64)
            .prefetch(tf.data.experimental.AUTOTUNE))
```

**Note**

- 画像サイズがまばらな場合があると思うが・・
  - 公開ソースコード読んだが、一律でリサイズしてたので、気にしなくて良さそう。
  - 一部だけとか、気になるなら、フィルターして除去しちゃう。(おまけのサンプルでは除去してやってみた)

- 正解ラベルと乖離が出るので、タスクによっては行わないデータ拡張。
  - 物体の数え上げ ⇒ ランダムクロップはNG
  - 物体の位置特定 ⇒ ランダムフリップはNG


#### データセットのサンプル数に依存　→ 学習ステップ数, Mix-Up

公式ブログTable2から引用。
![BigTransfer (BiT): State-of-the-art transfer learning for computer vision; Table2](https://3.bp.blogspot.com/-2JVDJV2A5Uo/XsMT-JfSJtI/AAAAAAAADFQ/UqcoPn11wFAudJJxkdzYxb3tAxBpgoGMQCLcBGAsYHQ/s1600/table%2B2.jpg)

`boundaries`は、以降の学習率のスケジューリングで使用。

```python
if dataset_size < 20 * 10 ** 3:
    schedule_len, boundaries = 500, [200, 300, 400]
elif 20 * 10 ** 3 <= dataset_size < 500 * 10 ** 3:
    schedule_len, boundaries = 10000, [3000, 6000, 9000]
else:
    schedule_len, boundaries = 20000, [6000, 12000, 18000]
```

#### MixUp のやり方

参考: https://github.com/google-research/big_transfer/blob/master/input_pipeline_tf2_or_jax.py#L118

```python
import tensorflow_probability as tfp

def mixup(image, label):
    beta_dist = tfp.distributions.Beta(0.1, 0.1)  # alpha=0.1
    beta = tf.cast(beta_dist.sample([]), tf.float32)
    image = (beta * image + (1 - beta) * tf.reverse(image, axis=[0]))
    label = (beta * label + (1 - beta) * tf.reverse(label, axis=[0]))
    return image, label
```

引数は適宜変更。この関数の場合、データセットに適応するタイミングは **ミニバッチ化`batch()`の後**。

MixUpが絡むので、ラベルはOne-Hotベクトルにする。

ちなみに、上のコードは`tf.reverse`により、バッチ内で中心から対称の位置にあるデータの組でMixUpが行われる。
以下、バッチ数16の例。左上[0]と右下[15]でMixUpされてる。
![MixUp Cat-Dog](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/450728/65d4943a-2b93-0c55-a3ac-d3d0fc223156.png)


#### バッチサイズ = 512

メモリ的に無理なら、下げてもいいみたい。

TF2のサンプルコードでは、バッチサイズで学習率・ステップ数を調整している。
・・のだが、他のサンプルではしていない。学習率変更タイミングのステップ数も変えてないけど、理由はよくわからない。

```python
batch_size = 64
schedule_len = schedule_len * 512 / batch_size
lr = 0.003 * batch_size / 512
```

#### 最適化アルゴリズム = SGD

- Learning rate: 0.003
- Momentum: 0.9

学習率は初期値。学習中の学習率の変更を、以下のスケジューリングを行う。

##### 学習率のスケジューリング

学習の進捗が、全体の30%, 60%, 90%になるタイミングで、学習率を$\frac{1}{10}$ずつ減衰させる。

公式サンプルコードでは、厳密に30%, 60%, 90%で区切っていない。
例えば、サンプル数：20k未満の時、200ステップまで0.003、201~300ステップ間は0.0003みたいになる。

```python
lr_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
    boundaries=boundaries, values=[lr, lr * 1e-1, lr * 1e-2, lr * 1e-3])

optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule, momentum=0.9)
```

※TF2のサンプルコードでは、`values=[lr, lr*0.1, lr*0.001, lr*0.0001]`と$\frac{1}{10}$ごとではない。PyTorch版だと順々に減衰させてる。あえてなのか、Typoなのかわからん。

---

### おまけ

サンプルコードを参考に Cats vs Dogs にファインチューニングさせて遊んだときの[Jupyter Notebook](https://github.com/miya-sh/study-xyz/tree/master/BiT)。