tf.kerasでEfficientNetV2の実装

#はじめに

EfficientNetの改良版というEffcientNetV2が発表されたので、実装して確認してみる。

#EfficientNetV2とは

元論文は[こちら](https://arxiv.org/abs/2104.00298)。

詳細は既にいくつか記事があるので、そちらを読んだ方が早いだろう。
[NFNetを超える速度と精度でEfficientNetが帰ってきた！！EfficientNetV2論文まとめ](https://qiita.com/nktankta/items/14797bcbc392a88c8f2e)
[2021年最強になるか！？最新の画像認識モデルEfficientNetV2を解説](https://qiita.com/omiita/items/1d96eae2b15e49235110)

改良により学習/推論時間が短くなり、精度も上がったと主張されている。

[公式実装](https://github.com/google/automl/tree/master/efficientnetv2)が既に公開されている。これはTensorFlowで実装されているが、この記事ではこれを元にKerasのフレームワークに移植した。元がTensorFlowなのであまり違いがないが、元のコードでは冗長と思われる点を削除等したので、こちらの方が見通しが良いコードにはなったと思う。

純粋なモデル変更の他に、学習データセットの強化や、徐々に画像サイズやデータ拡張等の強度を増やす、等の改良も含まれるが、この記事ではモデルの違いのみに注目して、オリジナルのEfficeintNetと性能差の確認も試みる。

以下、論文や公式githubから各モデルのImageNetの正解率/パラメータ数/FLOPsを表にまとめた。
論文とgithub上の表に数値上の相違がある場合があるが、その場合はgithub上に載っている数値を優先した。
比較のため、EfficentNetV2とEfficeintNet両方(以降V2/V1と略す)示す。

| Model | Top1-Acc | Params(M) | FLOPs(B) |
|:-:|:-:|:-:|:-:|
| V2-B0 | 78.7  | 7.1  | 0.72  |
| V2-B1 | 79.8  | 8.1  | 1.2  |
| V2-B2 | 80.5  | 10.1  | 1.7  |
| V2-B3 | 82.1  | 14.4  | 3.0  |
| V2-S  | 83.9  | 21.5  | 8.4  |
| V2-M  | 85.2  | 54.1  | 24.7  |
| V2-L  | 85.7  | 119.5  | 56.3  |

| Model | Top1-Acc | Params(M) | FLOPs(B) |
|:-:|:-:|:-:|:-:|
| V1-B0 | 77.1  | 5.3  | 0.39  |
| V1-B1 | 79.1  | 7.8  | 0.70  |
| V1-B2 | 80.1  | 9.2  | 1.0  |
| V1-B3 | 81.6  | 12  | 1.8  |
| V1-B4 | 82.9  | 19  | 4.2  |
| V1-B5 | 83.6  | 30  | 9.9  |
| V1-B6 | 84.0  | 42 | 19  |
| V1-B7 | 84.3  | 66 | 37  |

V2の論文にはB0/B1/B2/B3は出てこないようだが、github上には存在して重みもダウンロードできる。B4以降をS/M/Lという名前にしているようだ。V2-S以降が正式なV2のモデルで、B0からB3は、比較参考用という扱いのようだ。
V2ではImageNet21kで学習させてからImageNetの分類をさせてもいて、そちらの方が当然成績が良いのだが、比較目的なので上の表には含まない。ちなみに21kで学習させた場合は同じモデルでも約1%正解率が向上するようだ。




#実装

モデル設定の生成と、それを元にしたモデル生成の２段階に分ける。
V1とV2両方対応している。

##モデル設定

公式実装からかなりコピーしている。V1とV2の差異はこちらを参照した方がわかりやすいと思う。

```python
"""EfficientNet V1 and V2 model configs."""

#################### EfficientNet V1 configs ####################
v1_b0_block_str = [
        'r1_k3_s1_e1_i32_o16_se0.25',
        'r2_k3_s2_e6_i16_o24_se0.25',
        'r2_k5_s2_e6_i24_o40_se0.25',
        'r3_k3_s2_e6_i40_o80_se0.25',
        'r3_k5_s1_e6_i80_o112_se0.25',
        'r4_k5_s2_e6_i112_o192_se0.25',
        'r1_k3_s1_e6_i192_o320_se0.25',
]

#################### EfficientNet V2 configs ####################
v2_base_block = [    # The baseline config for v2 models.
        'r1_k3_s1_e1_i32_o16_c1',
        'r2_k3_s2_e4_i16_o32_c1',
        'r2_k3_s2_e4_i32_o48_c1',
        'r3_k3_s2_e4_i48_o96_se0.25',
        'r5_k3_s1_e6_i96_o112_se0.25',
        'r8_k3_s2_e6_i112_o192_se0.25',
]

v2_s_block = [    # about base * (width1.4, depth1.8)
        'r2_k3_s1_e1_i24_o24_c1',
        'r4_k3_s2_e4_i24_o48_c1',
        'r4_k3_s2_e4_i48_o64_c1',
        'r6_k3_s2_e4_i64_o128_se0.25',
        'r9_k3_s1_e6_i128_o160_se0.25',
        'r15_k3_s2_e6_i160_o256_se0.25',
]


v2_m_block = [    # about base * (width1.6, depth2.2)
        'r3_k3_s1_e1_i24_o24_c1',
        'r5_k3_s2_e4_i24_o48_c1',
        'r5_k3_s2_e4_i48_o80_c1',
        'r7_k3_s2_e4_i80_o160_se0.25',
        'r14_k3_s1_e6_i160_o176_se0.25',
        'r18_k3_s2_e6_i176_o304_se0.25',
        'r5_k3_s1_e6_i304_o512_se0.25',
]


v2_l_block = [    # about base * (width2.0, depth3.1)
        'r4_k3_s1_e1_i32_o32_c1',
        'r7_k3_s2_e4_i32_o64_c1',
        'r7_k3_s2_e4_i64_o96_c1',
        'r10_k3_s2_e4_i96_o192_se0.25',
        'r19_k3_s1_e6_i192_o224_se0.25',
        'r25_k3_s2_e6_i224_o384_se0.25',
        'r7_k3_s1_e6_i384_o640_se0.25',
]

v2_xl_block = [    # only for 21k pretraining.
        'r4_k3_s1_e1_i32_o32_c1',
        'r8_k3_s2_e4_i32_o64_c1',
        'r8_k3_s2_e4_i64_o96_c1',
        'r16_k3_s2_e4_i96_o192_se0.25',
        'r24_k3_s1_e6_i192_o256_se0.25',
        'r32_k3_s2_e6_i256_o512_se0.25',
        'r8_k3_s1_e6_i512_o640_se0.25',
]
efficientnet_params = {
        # (block, width, depth, train_size, eval_size, dropout, randaug, mixup, aug)
        'efficientnetv2-s':    # 83.9% @ 22M
                (v2_s_block, 1.0, 1.0, 300, 384, 0.2, 10, 0, 'randaug'),
        'efficientnetv2-m':    # 85.2% @ 54M
                (v2_m_block, 1.0, 1.0, 384, 480, 0.3, 15, 0.2, 'randaug'),
        'efficientnetv2-l':    # 85.7% @ 120M
                (v2_l_block, 1.0, 1.0, 384, 480, 0.4, 20, 0.5, 'randaug'),

        'efficientnetv2-xl':
                (v2_xl_block, 1.0, 1.0, 384, 512, 0.4, 20, 0.5, 'randaug'),

        # For fair comparison to EfficientNetV1, using the same scaling and autoaug.
        'efficientnetv2-b0':    # 78.7% @ 7M params
                (v2_base_block, 1.0, 1.0, 192, 224, 0.2, 0, 0, 'effnetv1_autoaug'),
        'efficientnetv2-b1':    # 79.8% @ 8M params
                (v2_base_block, 1.0, 1.1, 192, 240, 0.2, 0, 0, 'effnetv1_autoaug'),
        'efficientnetv2-b2':    # 80.5% @ 10M params
                (v2_base_block, 1.1, 1.2, 208, 260, 0.3, 0, 0, 'effnetv1_autoaug'),
        'efficientnetv2-b3':    # 82.1% @ 14M params
                (v2_base_block, 1.2, 1.4, 240, 300, 0.3, 0, 0, 'effnetv1_autoaug'),

        'efficientnet-b0': (v1_b0_block_str, 1.0, 1.0, 224, 224, 0.2),
        'efficientnet-b1': (v1_b0_block_str, 1.0, 1.1, 240, 240, 0.2),
        'efficientnet-b2': (v1_b0_block_str, 1.1, 1.2, 260, 260, 0.3),
        'efficientnet-b3': (v1_b0_block_str, 1.2, 1.4, 300, 300, 0.3),
        'efficientnet-b4': (v1_b0_block_str, 1.4, 1.8, 380, 380, 0.4),
        'efficientnet-b5': (v1_b0_block_str, 1.6, 2.2, 456, 456, 0.4),
        'efficientnet-b6': (v1_b0_block_str, 1.8, 2.6, 528, 528, 0.5),
        'efficientnet-b7': (v1_b0_block_str, 2.0, 3.1, 600, 600, 0.5),
        'efficientnet-b8': (v1_b0_block_str, 2.2, 3.6, 672, 672, 0.5),
        'efficientnet-l2': (v1_b0_block_str, 4.3, 5.3, 800, 800, 0.5),
}



class Struct:
        def __init__(self, **entries):
                self.__dict__.update(entries)

def make_config(model_name):
    def round_filters(filters, multiplier, divisor, min_depth=None, skip=False):
        """Round number of filters based on depth multiplier."""
        if skip or not multiplier:
            return filters

        filters = filters*multiplier
        min_depth = min_depth or divisor
        new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor)
        return int(new_filters)


    def round_repeats(repeats, multiplier, skip=False):
        """Round number of filters based on depth multiplier."""
        if skip or not multiplier:
            return repeats
        return int(math.ceil(multiplier * repeats))    
        
    cfg = efficientnet_params[model_name]
    width_coefficient = cfg[1]
    depth_coefficient = cfg[2] 

    block_args_list = []
    for block_string in cfg[0]:
        ops = block_string.split('_')
        options = {}
        for op in ops:
            splits = re.split(r'(\d.*)', op)
            if len(splits) >= 2:
                key, value = splits[:2]
                options[key] = value

        input_filters = round_filters(int(options['i']), width_coefficient, 8)
        output_filters = round_filters(int(options['o']), width_coefficient, 8)

        block0_args =    Struct(
                kernel_size=int(options['k']),
                input_filters=input_filters,
                output_filters=output_filters,
                expand_ratio=int(options['e']),
                se_ratio=float(options['se']) if 'se' in options else 0.0,
                strides=int(options['s']),
                conv_type=int(options['c']) if 'c' in options else 0,
        )

        block_args_list.append(block0_args)
        repeats = round_repeats(int(options['r']), depth_coefficient)
        for _ in range(repeats - 1):
            block_args = copy.deepcopy(block0_args)
            block_args.input_filters=int(output_filters)
            block_args.strides=int(1)
            block_args_list.append(block_args)
    
    mconfig = Struct(
        name=model_name,
        block_args_list = block_args_list,
        train_size = cfg[3],
        eval_size = cfg[4],
        dropout_rate = cfg[5],
        num_classes = 1000,
        feature_size = round_filters( 1280, width_coefficient, 8),
        depth_divisor = 8,
        bn_momentum = 0.9,
        bn_epsilon = 1e-3,
        min_depth = 8,
        survival_prob = 0.8,
        weight_decay = 0.0,
        stem_stride=2,

    )
    return mconfig

```

##モデル生成

先に示したmake_configを使って作成したパラメータを渡してKerasのモデルを作成する。
こちらはオリジナルのコードが多いが、基本的な構成は元コードに従っている。
純粋にモデルとして見た場合の大きな違いはFusedMBconvの導入だが、普通モデルがバージョンアップするとモデルが複雑になる場合が多いと思うのだが、FusedMBConvはMBConvより構造が単純になっている。

```python
def conv_kernel_initializer(shape, dtype=None, partition_info=None):
    del partition_info
    kernel_height, kernel_width, _, out_filters = shape
    fan_out = int(kernel_height * kernel_width * out_filters)
    return tf.random.normal(
            shape, mean=0.0, stddev=np.sqrt(2.0 / fan_out), dtype=dtype)

def dense_kernel_initializer(shape, dtype=None, partition_info=None):
    del partition_info
    init_range = 1.0 / np.sqrt(shape[1])
    return tf.random.uniform(shape, -init_range, init_range, dtype=dtype)

def create_efficientnet(config, input_shape=(None,None,3) , name=None, verbose=0):
    mconfig = config
    BN_KWARGS = {'momentum':mconfig.bn_momentum, 'epsilon':mconfig.bn_epsilon}
    CONV_KWARGS = {'use_bias':False, 'padding':'same', 
                    'kernel_initializer': conv_kernel_initializer, 
                    'kernel_regularizer': tf.keras.regularizers.l2(mconfig.weight_decay)}

    def activation(fn='silu', name=None):
        if fn=='sigmoid':
            return layers.Lambda( lambda y: tf.sigmoid(y), name=name)
        return layers.Lambda( lambda y: tf.nn.silu(y), name=name)

    def conv_bn_act(x, filters, kernel, stride, prefix):
        x = layers.Conv2D( filters, kernel, stride, name=prefix+'_conv2d',    **CONV_KWARGS)(x)
        x = layers.BatchNormalization(name=prefix+'_bnorm', **BN_KWARGS)(x)
        x = activation(name=prefix+'_act')(x)
        return x

    def se(inputs, se_filters, filters, prefix):
        SE_CONV_KWARGS = {'use_bias':True, 'padding':'same', 
                        'kernel_initializer': conv_kernel_initializer,
                        'kernel_regularizer': tf.keras.regularizers.l2(mconfig.weight_decay)}
        x = layers.Lambda(lambda y:tf.reduce_mean(y, [1, 2], keepdims=True), name=prefix+'_se_mean')(inputs)
        x = layers.Conv2D(se_filters, 1, 1, name=prefix+'_se_conv2d_s', **SE_CONV_KWARGS)(x)
        x = activation(name=prefix+'_se_act_s')(x)
        x = layers.Conv2D(filters, 1, 1, name=prefix+'_se_conv2d_e', **SE_CONV_KWARGS)(x)
        x = activation('sigmoid', name=prefix+'_se_act_e')(x)
        x = layers.Multiply(name=prefix+'_se_mult')( [x ,inputs ])
        return x

    def mb_conv(inputs, block_args, prefix):
        filters = block_args.input_filters * block_args.expand_ratio
        se_filters = int(block_args.input_filters*block_args.se_ratio)
        x = inputs
        if block_args.expand_ratio != 1:
            x = layers.Conv2D( filters, 1, 1, name=prefix+'_conv2d_expand', **CONV_KWARGS)(x)
            x = layers.BatchNormalization(name=prefix+'_bnorm_expand', **BN_KWARGS)(x)
            x = activation(name=prefix+'_act_expand')(x)

        x = layers.DepthwiseConv2D( block_args.kernel_size, block_args.strides, name=prefix+'_dwconv2d', **CONV_KWARGS)(x)
        x = layers.BatchNormalization(name=prefix+'_bnorm', **BN_KWARGS)(x)
        x = activation(name=prefix+'_act')(x)
        if se_filters!=0:
            x = se(x, se_filters, filters, prefix=prefix)
        x = layers.Conv2D( block_args.output_filters, 1, 1, name=prefix+'_conv2d_proj', **CONV_KWARGS)(x)
        x = layers.BatchNormalization(name=prefix+'_bnorm_proj', **BN_KWARGS)(x)

        return x

    def fused_mb_conv(inputs, block_args, prefix):
        filters = block_args.input_filters * block_args.expand_ratio
        se_filters = int(block_args.input_filters*block_args.se_ratio)
        x = inputs
        if block_args.expand_ratio != 1:
            x = layers.Conv2D( filters, block_args.kernel_size, block_args.strides, 
                name=prefix+'_conv2d_expand', **CONV_KWARGS)(x)
            x = layers.BatchNormalization(name=prefix+'_bnorm_expand', **BN_KWARGS)(x)
            x = activation(name=prefix+'_act_expand')(x)
            if se_filters!=0:
                x = se(x, se_filters, filters, prefix=prefix)
            x = layers.Conv2D( block_args.output_filters, 1, 1, 
                name = prefix+'_conv2d', **CONV_KWARGS )(x)
            x = layers.BatchNormalization(name=prefix+'_bnorm', **BN_KWARGS)(x)
        else:
            if se_filters!=0:
                x = se(x, se_filters, filters, prefix=prefix)
            x = layers.Conv2D( block_args.output_filters, block_args.kernel_size, block_args.strides,
                name = prefix+'_conv2d', **CONV_KWARGS )(x)
            x = layers.BatchNormalization(name=prefix+'_bnorm', **BN_KWARGS)(x)
            x = activation(name=prefix+'_act')(x) # add act if no expansion.
        
        return x

    # Input
    x = input = tf.keras.Input(shape=input_shape, name='input')


    # Stem
    stem_filters = mconfig.block_args_list[0].input_filters
    x = conv_bn_act(x, stem_filters , 3, mconfig.stem_stride, prefix='stem')
    if verbose!=0:
        print('stem', 3, stem_filters, mconfig.stem_stride )

    # Blocks
    for i, block_args in enumerate(mconfig.block_args_list):
        shortcut = x
        if block_args.conv_type==0:
            prefix = f'block_{i}'
            x = mb_conv(x, block_args, prefix)
        else:
            prefix = f'block_{i}_fused'
            x = fused_mb_conv(x, block_args, prefix)
        
        survival_prob = mconfig.survival_prob
        if block_args.strides == 1 and block_args.input_filters==block_args.output_filters:
            if survival_prob:
                # StochasticDepth
                drop_rate = 1.0 - survival_prob
                drop_rate = drop_rate * float(i+1) / len(mconfig.block_args_list)
                survival_prob = 1.0 - drop_rate
                x = layers.Dropout(rate=drop_rate, noise_shape=(None,1,1,1), name=prefix+'_stochdepth')(x)
            x = layers.Add(name=prefix+'_add')([shortcut,x])
        else:
            survival_prob = None

        if verbose!=0:
            print(i, block_args.kernel_size,block_args.input_filters, block_args.output_filters, 
                block_args.strides, block_args.expand_ratio, survival_prob    )

    # Head
    x = conv_bn_act( x,mconfig.feature_size , 1, 1, prefix='head')
    x = layers.GlobalAveragePooling2D(name='head_gap')(x)
    if mconfig.dropout_rate > 0:
        x = layers.Dropout(mconfig.dropout_rate,name='head_dropout')(x)
    output = layers.Dense(mconfig.num_classes,
            kernel_initializer = dense_kernel_initializer,
            kernel_regularizer = tf.keras.regularizers.l2(mconfig.weight_decay),
            use_bias = True, name='head_fc')(x)
    if verbose!=0:
        print('head', mconfig.feature_size, mconfig.num_classes, mconfig.dropout_rate    )

    model = tf.keras.Model(inputs=[input], outputs=[output], name=name)
    if verbose!=0:
        print('params', model.count_params())
    return model

```

#実験

公式実装にはGoogleColabで動作するチュートリアルが含まれ、そこではCIFAR10での転移学習も実行できるのだが、本記事ではスクラッチでCIFAR100を学習させてみた。
V2-B0とV1-B0で同じ学習をさせて比較する。

以下、実験に関する細かい解説など。

* TPU使用、バッチサイズ128、エポック数150で３回実施
* 複雑なデータ拡張は行わず、簡単なフリップやカラーシフトと、Cutoutで済ませた。Cutoutの確率は学習中で線形に増加させるようにしてある。
* Optimizerには論文通りRMSpropでmomentum=0.9を使った。筆者の認識ではRMSpropでmomentumを使ったらほぼAdamと同じになるはずだが、SGDよりこちらの方が良いという判断なのだろう。論文中には書いていないが、公式実装ではepsilon=0.001としている。
* V2では学習時と推論時で入力画像のサイズを変えるような設計なのだが、ここはV1との比較のために推論と同じサイズ(224x224)で学習させた。
* 論文中に"We use exponential moving average with 0.9999 decay rate"とあって、最初何のことかよくわからなかったのだが、コードをみる限りtf.train.ExponentialMovingAverageを使うようだ。tfa.optimizers.MovingAverageを使えば同等になるような実装もできそうな気がするが、よくわからないので対応しなかった。
* その他、DropoutやStochasticDepthはデフォルト通りに適用。

##実験結果

以下、学習結果を掲載。Timeは３回中の中央値。

| Model | Params(M) | FLOPs(B) | Time(sec.) | Accuracy(%) |
|:-:|:-:|:-:|:-:|:-:|
| V2-B0  | 7.1  | 0.72 | 2953  | 81.48/81.92/82.16  | 
| V1-B0  | 5.3  | 0.39  | 3442  | 80.84/80.90/80.95  |

上記はミスでWeight Decayなしで計測したもの。論文通りWeight Decay 1e-5 の場合の結果は以下の通り。

| Model | Params(M) | FLOPs(B) | Time(sec.) | Accuracy(%) |
|:-:|:-:|:-:|:-:|:-:|
| V2-B0  | 7.1  | 0.72 | 2870  | 82.93/83.15/83.35  | 
| V1-B0  | 5.3  | 0.39  | 3488  | 82.25/82.48/82.51  |

以下、所見。

* V2-B0の方が正解率が高くなった。
  * パラメータ数/FLOPsが大きいので、当たり前の結果にも見える。
* V2-B0の方が学習にかかる時間は短かった。
  * パラメータ数/FLOPsも大きいにも関わらず短くなったのは、総レイヤー数が少ないことなど、構造が簡単になっているからと思われる。 
* 学習/推論時間が短いにも関わらず、正解率が高くなったので、論文の主張に沿う結果にはなったように見える。
  * ただし、同じ設定でCIFAR10を学習させたら認識率はV1の方が高くなった。少し調整が必要かもしれない。

[実験用ノートブック](https://colab.research.google.com/gist/T-STAR-LTD/10406551075fe0f73d9d3550ea8c5303/efficientnetv2.ipynb?authuser=1)

#まとめ

EfficeintNetV2をtf.kerasで実装し、CIFAR100でV1と比較した。
実際は転移学習で使うことが多いはずなので、CIFAR100を最初から学習させることに意義があるかは微妙なところだが、V1-B0とV2-B0の比較では少なくともCIFAR100では論文での主張通り性能向上が見られた。