TensorFlowでModel・Layer・Metricをカスタマイズする

## はじめに

TensorFlowにおいてモデルを自分で実装するときには既に用意されたクラスを継承するのが便利だが，理解するのに結構時間がかかったので忘れないようにメモ．

##### 確認した環境

- Python 3.6 & TensorFlow 1.15
- Python 3.7 & TensorFlow 2.2


### モデルのカスタマイズ

[公式のチュートリアル](https://www.tensorflow.org/tutorials/customization/custom_layers?hl=ja#モデル：レイヤーの組み合わせ)を参考にして作成．2つの畳込み層を通過した後に入力値を足し合わせ，出力のサイズが入力の半分になるようなResidual Blockを作る．

```
from tf.keras.layers import Conv2D, BatchNormalization, Add, Activation

class ResidualBlock(tf.keras.Model):
    
    def __init__(self, filters, kernel_size=2, block_name=''):

        # おまじない
        super(ResidualBlock, self).__init__()

        # 入力値を足し合わせるためにサイズを出力に合わせる
        self.conv0 = Conv2D(filters, 1, strides=2, padding='same', name=block_name+'_conv0')
        self.bn0 = BatchNormalization(name=block_name+'_bn0')

        # 畳み込み層（1層目でサイズを半分にする）
        self.conv1 = Conv2D(filters, kernel_size, strides=2, padding='same', activation='relu', name=block_name+'_conv1')
        self.bn1 = BatchNormalization(name=block_name+'_bn1')
        self.conv2 = Conv2D(filters, kernel_size, padding='same', activation='relu', name=block_name+'_conv2')
        self.bn2 = BatchNormalization(name=block_name+'_bn2')
        
        # 入力と出力を足し算
        self.add = Add(name=block_name+'_add')
        self.out = Activation('relu', name=block_name+'_out')
        
    def call(self, x):
        
        shortcut = self.conv0(x)
        shortcut = self.bn0(shortcut)
        
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.conv2(x)
        x = self.bn2(x)
       
        x = self.add([shortcut,x])
        x = self.out(x)
        
        return x
```

必要な層は```__init__```で作っておき，```call```で計算の流れを実装する．

### レイヤーのカスタマイズ

カスタマイズしたモデルや```tf.keras.applications```で用意されているモデルは，```tf.keras.models.Sequential```を使って組み合わせることで新しいモデルを作ることができる．ヒートマップを作るとき等にモデル内部のレイヤーにアクセスしたいときがあるが，それが自分で調べる限りだと出来なさそう．そこで，モデルの出力をそのまま返すような層を単体で作っておく．

```
class IdentityLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(IdentityLayer, self).__init__()
    def call(self, x):
        return x
```

モデルのカスタマイズとほとんど同じ．

### メトリックのカスタマイズ

[公式のAPI](https://www.tensorflow.org/api_docs/python/tf/keras/metrics/Metric?hl=en)を参考にして作成．Categoricalな入出力に対応するTruePositivesを作る．

```
class TP_metric(tf.keras.metrics.Metric):

    def __init__(self, name='TP', **kwargs):
        
        # おまじない
        super(TP_metric, self).__init__(name=name, **kwargs)
        # 内部変数の定義
        self.value = self.add_weight(name='tp', initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):

        # one_hotベクトルをラベルの配列に変換
        y_true = tf.argmax(y_true, axis=-1)
        y_true = tf.cast(y_true, K.floatx())
        y_pred = tf.argmax(y_pred, axis=-1)
        y_pred = tf.cast(y_pred, K.floatx())
        
        # 真値と予測値がともに1ならばTPは1増える
        tmp = tf.equal(tf.add(y_true, y_pred), 2)
        tmp = tf.cast(tmp, K.floatx())

        # 内部変数の更新
        self.value.assign_add(tf.reduce_sum(tmp))

    def result(self):
        return self.value
```

けっこう型にはうるさい（Pythonなのに）ので，こまめにcastをするように心がける．最初はSparse Categoricalに対応するように作ろうとしたが，何故か```y_true```のshapeが```y_pred```と同じ値になっていた．代わりにラベルのデータセットの方を調整することで対処．

```
def one_hot(x):
    return tf.one_hot(x, カテゴリー数)

label_ds = tf.data.Dataset.from_tensor_slices(tf.cast(label_list, tf.int32))
label_ds = label_ds.map(one_hot)
```

### おまけ

tf 1.xとtf 2.xを両方使うようになったことで，```eager_execution```が少し分かった気になった．