Keras编写 Mish 激活函数

Mish激活函数
网上给出的Mish激活函数多是pytorch写的，今天折腾了一会儿写了keras框架下的Mish函数。如有不对，请大家指正：

from keras.engine.base_layer import Layer
class Mish(Layer):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        '''
        Forward pass of the function.
        '''
        return x * (K.tanh(K.softplus(x)))

def bn_mish(input):
     return Mish()(BatchNormalization(axis=CHANNEL_AXIS)(input))