tf.kerasでFReLUを実装

概要

FReLUという新たな活性化関数のkerasでの実装を紹介します。

FReLUについての詳細+pytorchでの実装は以下の記事を読んでみてください。
非常に丁寧に解説されています。
参考:新たな活性化関数「FReLU」誕生&解説！

本記事はtf.kerasでのFReLUの実装をしようという内容です。
DepthwiseConv2Dとtfのmaximumを使えばサクッと実装できてしまいます。

変更・追記（2020/9/30）

・Lambda関数が必要ありませんでしたので削除したコードに変更しました
・Layer版の実装を載せました
参考:カスタムレイヤー

定義

y= \max(x,\mathbb{T}(x))

$\mathbb{T}(\cdot)$はDepthwiseConv2Dです。

前提

tensorflow 2.3.0

実装:１つ１つレイヤーとして定義

FReLU.py

from tensorflow.keras.layers import DepthwiseConv2D,BatchNormalization
import tensorflow as tf
from tensorflow.keras.layers import Lambda
from tensorflow.keras import backend as K

def FReLU(inputs, kernel_size = 3):
    #T(x)の部分
    x = DepthwiseConv2D(kernel_size, strides=(1, 1), padding='same')(inputs)
    x = BatchNormalization()(x)
    #max(x, T(x))の部分
    x = tf.maximum(inputs, x)
    return x

実装:１つのLayerとして定義

FReLU.py

class FReLU(tf.keras.Model):
    def __init__(self, kernel_size=3):
        super(FReLU, self).__init__(name='')

        self.depconv = DepthwiseConv2D(kernel_size, strides=(1, 1), padding='same')
        self.bn = BatchNormalization()

    def call(self, input_tensor, training=False):
        x = self.depconv(input_tensor)
        x = self.bn(x, training=training)
        out = tf.maximum(input_tensor, x)
        return out

終わりに

実装自体は簡単ですね。
不明点、おかしい点ありましたらコメントよろしくお願いします。