【Tensorflow 2.0 正式版教程】自定义层

本文介绍如何将自定义的运算或操作封装为一个层。

首先，任何自定义层都需要继承Layer类，并且需要实现4个类方法

class MyLayer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(MyLayer, self).__init__(**kwargs)
    
    def build(self, input_shape):
    
    def call(self, inputs):
        return output

    def compute_output_shape(self, input_shape):
        return output_shape

__init__()一般记录层初始化时所需的参数。
build()是网络编译时执行的操作，一般在这里利用类方法self.add_weight()完成权重初始化，并且输入参数input_shape会是call中第一个输入参数的形状，在定义的层有多个输入的时候需要多加注意。
call()是定义该层的计算。
compute_output_shape需要手动计算出输出tensor的形状，但是通常情况下tensorflow是能自行计算出call中输出tensor的形状的。

下面以卷积层为例，由于tensorflow官方实现的卷积层必须指定卷积核大小，默认padding形式为valid，但实际场景中卷积核通常为(3, 3)，padding为same，所以我们来重新实现一个更合理的卷积层。

class Conv2D(tf.keras.layers.Layer):
    def __init__(self, output_dim, kernel_size=(3, 3), strides=(1, 1, 1, 1), **kwargs):
        self.output_dim = output_dim
        self.kernel_size = kernel_size
        self.strides = strides
        super(Conv2D, self).__init__(**kwargs)

    def build(self, input_shape):
        kernel_shape = tf.TensorShape((self.kernel_size[0], self.kernel_size[1], input_shape[-1], self.output_dim))
        self.kernel = self.add_weight(name='kernel',
                                      shape=kernel_shape,
                                      initializer=tf.initializers.he_normal(),
                                      trainable=True)

    def call(self, inputs):
        output = tf.nn.conv2d(inputs, filters=self.kernel, strides=self.strides, padding='SAME')
        return output

    def compute_output_shape(self, input_shape):
        shape = tf.TensorShape(input_shape).as_list()
        shape[-1] = self.output_dim
        return tf.TensorShape(shape)

这里也可以看到self.add_weight()方法的具体用法，通常需要指定的参数有name，用于计算loss或加载权重时的筛选；shape指定参数的形状；initializer指定初始化方法，这里选用了Kaiming He的论文中所用的初始化方法，普通的如全零初始化可直接initializer=tf.initializers.zeros；trainable是设置参数是否参与训练。

tensorflow对层的定义是可嵌套的，即可以将多个基本层组合到一起定义为一个新的层，比如我们可以将一个残差块（residual block）定义为一个层。

class ResBlock(tf.keras.layers.Layer):
    def __init__(self, output_dim, strides=(1, 1, 1, 1), **kwargs):
        self.strides = strides
        if strides != (1, 1, 1, 1):
            self.shortcut = Conv2D(output_dim, kernel_size=(1, 1), strides=self.strides)
        self.conv_0 = Conv2D(output_dim, strides=self.strides)
        self.conv_1 = Conv2D(output_dim)
        self.bn_0 = BatchNormalization(momentum=0.9, epsilon=1e-5)
        self.bn_1 = BatchNormalization(momentum=0.9, epsilon=1e-5)
        super(ResBlock, self).__init__(**kwargs)

    def call(self, inputs, training):
        net = self.bn_0(inputs, training=training)
        net = tf.nn.relu(net)

        if self.strides != (1, 1, 1, 1):
            shortcut = self.shortcut(net)
        else:
            shortcut = inputs

        net = self.conv_0(net)
        net = self.bn_1(net, training=training)
        net = tf.nn.relu(net)
        net = self.conv_1(net)

        output = net + shortcut
        return output

这种情况下只需要在初始化时完成对所需的层的实例化，并在call中指定计算方式即可。这里call中有一个额外的参数training，这是因为bn层在训练和预测是有不同的行为。这也意味着在定义整体网络时，不需要给出training值，而在网络运算时不仅给出网络的输入还要给出training值。

更多的代码可以在我的github上找到
https://github.com/Apm5/tensorflow_2.0_tutorial/blob/master/CNN/Layers.py