EvoNormの実装と評価

#はじめに

DeepLearningのネットワークにBatch NormalizationとReLUを組み合わせて使用することはもはや常識的になっているが、これを置き換えるものとしてEvoNormという手法が提案されている。本記事ではこれを実装して評価する。


#EvoNormとは何か

以下が発表論文。
[Evolving Normalization-Activation Layers](https://arxiv.org/abs/2004.02967)

Batch NormalizationとReLUの組み合わせに代わる機能を持つグラフをAutoMLで探索し、最終的に出来上がったものを一つのレイヤーとしたものがEvoNorm-B0。

下記は論文に載っている構造を示す図で、いかにもAutoMLによる構成という感じになっている。

![スクリーンショット 2021-06-14 17.37.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/a298f89c-e89a-3f43-716b-a6a98ea97700.png)

各種ネットワークでEvoNorm-B0に置き換えた効果が以下の通り論文に掲載されている。

![スクリーンショット 2021-06-14 17.50.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/bb65db62-94d0-a8df-70b0-df10a60a6dcb.png)


同様にGroup Normalizationに代わるものを探索して出来上がったものがEvoNorm-S0のようで、論文内で提示されている効果は以下の通り。

![スクリーンショット 2021-06-14 17.56.43.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/3c20afbf-088f-4409-2302-385700480827.png)


#実装

公式実装は[こちら](https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_model.py)。
ResNetV1/V2/RSのモデル実装の一部としてコードがある。プラットフォームはTensorflow。そもそも筆者がEvoNormを知ったのはResNet-RSの実装確認のためソースを見ていた際にたまたまこの実装を見たため。
非公式だが[pytorchの実装](https://github.com/digantamisra98/EvoNorm)もある。

この記事ではtf.kerasのレイヤーとして実装したものを掲載＆実験結果報告する。tensorflowのバージョンは2.5.0。
実装を簡単にするため、２次元データしか対応していない。


##EvoNorm-B0

EvoNorm-B0の実装を以下に示す。

```python
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.python.ops import variables as tf_variables

class EvoNormB0(tf.keras.layers.Layer):
    def __init__(self, momentum=0.9, epsilon=1e-5, **kwargs):
        super().__init__(**kwargs)

        self.momentum = momentum
        self.epsilon = epsilon

    def build(self, input_shape):
        var_shape, axes = (1, 1, 1, input_shape[3]), [0, 1, 2]
        self.axes = axes
        self.gamma = self.add_weight(name="gamma", shape=var_shape, initializer=tf.initializers.Ones())
        self.beta = self.add_weight(name="beta", shape=var_shape, initializer=tf.initializers.Zeros())
        self.v_1 = self.add_weight(name="v1", shape=var_shape, initializer=tf.initializers.Ones())

        self.running_variance = self.add_weight(name="running_variance",trainable=False, shape=var_shape, 
                                                synchronization=tf_variables.VariableSynchronization.ON_READ,
                                                aggregation=tf_variables.VariableAggregation.MEAN,
                                                initializer=tf.initializers.Ones())

    def call(self, inputs, training=True):
        def instance_std(x, epsilon):
            _, var = tf.nn.moments(x, axes=[1, 2], keepdims=True)
            return tf.sqrt(var + epsilon)
        
        def batch_std(x, epsilon):
            if training:
                _, var = tf.nn.moments(x, self.axes, keepdims=True)
                self.running_variance.assign(self.momentum * self.running_variance + (1 - self.momentum) * var)
            else:
                var = self.running_variance
            return tf.sqrt(var + epsilon)

        right = batch_std(inputs, self.epsilon)        
        left = instance_std(inputs, self.epsilon) + self.v_1 * inputs
        denominator = tf.maximum(left, right)
        return inputs * self.gamma / denominator + self.beta

    def get_config(self):
        config = {'momentum': self.momentum, 'epsilon':self.epsilon}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

```
##Batch Normalization

比較のため、標準的なBatchNormalizationの実装を以下に示す。
ただし、BN+ReLUのレイヤーとして使えるように、activationとしてreluも使用できるようになっているが、これは標準的ではない。

```python
class BatchNorm2D(tf.keras.layers.Layer):
    def __init__(self, momentum=0.9, epsilon=1e-5, relu=False, **kwargs):
        super().__init__(**kwargs)

        self.momentum = momentum
        self.epsilon = epsilon
        self.relu = relu

    def build(self, input_shape):
        var_shape, axes = (1, 1, 1, input_shape[3]), [0, 1, 2]
        self.axes = axes
        self.gamma = self.add_weight(name="gamma", shape=var_shape, initializer=tf.initializers.Ones())
        self.beta = self.add_weight(name="beta", shape=var_shape, initializer=tf.initializers.Zeros())

        self.running_variance = self.add_weight(name='running_variance', trainable=False, shape=var_shape, 
                                            synchronization=tf_variables.VariableSynchronization.ON_READ,
                                            aggregation=tf_variables.VariableAggregation.MEAN,
                                            initializer=tf.initializers.Ones())
        self.running_mean = self.add_weight(name='running_mean',trainable=False, shape=var_shape, 
                                            synchronization=tf_variables.VariableSynchronization.ON_READ,
                                            aggregation=tf_variables.VariableAggregation.MEAN,
                                            initializer=tf.initializers.Zeros())

    def call(self, inputs, training=True):
        if training:
            mean, variance = tf.nn.moments(inputs, self.axes, keepdims=True)
            self.running_variance.assign(self.momentum * self.running_variance + (1 - self.momentum) * variance)
            self.running_mean.assign(self.momentum * self.running_mean + (1 - self.momentum) * mean)
        else:
            variance = self.running_variance
            mean = self.running_mean
        
        denominator = tf.sqrt(variance + self.epsilon)
        outputs = (inputs-mean) * self.gamma / denominator + self.beta
        if self.relu:
            outputs = tf.maximum( outputs, 0 )
        return outputs 

    def get_config(self):
        config = {'momentum': self.momentum, 'epsilon':self.epsilon, 'relu':self.relu}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))
```


##EvoNorm-S0

EvoNorm-S0の実装を以下に示す。

```python
import tensorflow as tf
from tensorflow.python.ops import variables as tf_variables
class EvoNormS0(tf.keras.layers.Layer):
    def __init__(self, groups=32, epsilon=1e-5, **kwargs):
        super().__init__(**kwargs)
        self.groups = groups
        self.epsilon = epsilon

    def build(self, input_shape):
        var_shape, axes = (1, 1, 1, input_shape[3]), [0, 1, 2]
        self.axes = axes
        self.gamma = self.add_weight(name="gamma", shape=var_shape, initializer=tf.initializers.Ones())
        self.beta = self.add_weight(name="beta", shape=var_shape, initializer=tf.initializers.Zeros())
        self.v_1 = self.add_weight(name="v1", shape=var_shape, initializer=tf.initializers.Ones())

    def call(self, inputs, training=True):
        def group_std(inputs, groups, epsilon):
            input_shape = tf.shape(inputs)
            groups = tf.minimum(input_shape[3], groups)
            group_shape = [input_shape[0], input_shape[1], input_shape[2], groups, input_shape[3]//groups]
            grouped_inputs = tf.reshape(inputs, group_shape)
            _, var = tf.nn.moments(grouped_inputs, [1, 2, 4], keepdims=True)

            std = tf.sqrt(var + epsilon)
            std = tf.broadcast_to(std, tf.shape(grouped_inputs))
            return tf.reshape(std, input_shape)

        denominator = group_std(inputs, groups=self.groups, epsilon=self.epsilon)
        return (inputs * tf.sigmoid(self.v_1 * inputs)) / denominator * self.gamma + self.beta

    def get_config(self):
        config = {'groups': self.groups, 'epsilon':self.epsilon}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))
```

##Group Normalization

比較のため、普通のGroup Normalizationのコードを示す。

```python
class GroupNorm(tf.keras.layers.Layer):
    def __init__(self, groups=32, epsilon=1e-5, relu=False,**kwargs):
        super().__init__(**kwargs)
        self.groups = groups
        self.epsilon = epsilon
        self.relu = relu

    def build(self, input_shape):
        var_shape, axes = (1, 1, 1, input_shape[3]), [0, 1, 2]
        self.axes = axes
        self.gamma = self.add_weight(name="gamma", shape=var_shape, initializer=tf.initializers.Ones())
        self.beta = self.add_weight(name="beta", shape=var_shape, initializer=tf.initializers.Zeros())

    def call(self, inputs, training=True):
        def group_std(inputs, groups, epsilon):
            input_shape = tf.shape(inputs)
            groups = tf.minimum(input_shape[3], groups)
            group_shape = [input_shape[0], input_shape[1], input_shape[2], groups, input_shape[3]//groups]
            grouped_inputs = tf.reshape(inputs, group_shape)
            mean, var = tf.nn.moments(grouped_inputs, [1, 2, 4], keepdims=True)

            std = tf.sqrt(var + epsilon)
            std = tf.broadcast_to(std, tf.shape(grouped_inputs))
            mean = tf.broadcast_to(mean, tf.shape(grouped_inputs))
            return tf.reshape(mean, input_shape), tf.reshape(std, input_shape)

        mean, std = group_std(inputs, groups=self.groups, epsilon=self.epsilon)
        outputs =(inputs-mean) / std * self.gamma + self.beta

        if self.relu:
            outputs = tf.maximum( outputs, 0 )
        return outputs

    def get_config(self):
        config = {'groups': self.groups, 'epsilon':self.epsilon, 'relu':self.relu}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

```

#実験

WideResNet22-8でCIFAR10を学習させた結果で比較する。
150エポックでの最も高いValidationの認識率を採用し、各５回実施での最高/最低/中央値を掲載する。
バッチサイズは256で、TPU環境で実施。Cutout等のデータ拡張が入っている。

## EvoNorm-B0 と Batch Normalizationの比較

momentum=0.9/epsilonは1e-5で実施。
論文ではBN+ReLUの他にBN+SiLUも比較対象となっているので、そちらも実験した。
BNはtf.keras標準の実装を使用した。

| Layer | Min | Median | Max |
|:-:|:-:|:-:|:-:|
| BN+ReLU  | 96.36 |96.50 |96.65  |
| BN+SiLU  | **96.72** | **96.78** | **96.89**  |   
| EvoNorm-B0  | 96.34 | 96.60 |96.69 |  

![EvoNormB0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/9f17296d-c34d-3331-fdc1-2c10b8d46f2e.png)


以下、所見。

* BN+ReLUとEvoNorm-B0では最終結果には大きな違いが見られなかったが、若干EvoNorm-B0の方が良い
* ReLUをSiLUに置き換えた場合の方が明確に性能向上が見られた
* グラフを比較すると、EvoNorm-B0の方が若干学習が遅く、揺れは少ないように見える

## EvoNorm-S0 と Group Normalizationの比較

Group数は32、epsilonは1e-5で実施。
GNはTensorflow Addonsの実装を使用した。
GN+SiLUの結果も合わせて掲載。

| Layer | Min | Median | Max |
|:-:|:-:|:-:|:-:|
|GN+ReLU | 95.46 | 95.58 | 95.85 |
|GN+SiLU | 95.80 | 95.91 | 96.07 |
|EvoNorm-S0|  **96.31** | **96.43** | **96.61** |

![EvoNormS0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/824193/370d5154-3bb5-8897-2177-7664f8aa33fb.png)


以下、所見。

* EvoNorm-S0の性能が最も良かった
* ReLUをSiLUに差し替えた場合でも明確な性能向上が見られるが、EvoNorm-S0には及ばない
* グラフを比較すると、EvoNorm-S0の方が学習が速く、揺れも少ないようだ

#まとめ

EvoNorm-B0とEvoNorm-S0の実装と評価を行なった。
EvoNorm-B0の方はBatch Normalizaton+ReLUに比べてあまり効果が出なかったが、EvoNorm-S0の方はGroup Normalization+ReLUに比べて大きな効果が確認できた。
今回の実験ではBatch Normalization+SiLUの組み合わせが最も良かった。SiLUはSwishとも呼ばれてEfficientNet等で採用されているが、他の既存のモデルでも差し替えるだけでそれなりに性能向上する可能性がある。
ただし、上記の結論は今回の実験条件での結果によるもので、条件が変われば別の結論になる可能性はある。