【人工智能笔记】第七节：Tensorflow 2.0 自定义Layer

 自定义Layer：

1. 继承 tf.keras.layers.Layer。
2. 在__init__构造函数里调用父类构造函数，并把参数传到父类。
3. build方法里初始化权重变量。
4. call方法里编写计算图。

构造函数参数：

trainable：布尔值，层的变量是否可训练。
name：图层的字符串名称。
dtype：图层的计算和权重的dtype（在TensorFlow 2中None使用均值的默认值 tf.keras.backend.floatx，或TensorFlow 1中第一个输入的类型）。
dynamic：将其设置为True仅在您的图层只急切运行且不应用于生成静态计算图的情况下。例如，对于Tree-RNN或递归网络，或者对于使用Python控制流操纵张量的任何层，通常就是这种情况。如果为False，则假定可以安全地使用该层来生成静态计算图。

例子： 

import tensorflow as tf
import numpy as np

class MyDenseLayer(tf.keras.layers.Layer):
    def __init__(self, num_outputs, **args):
        super(MyDenseLayer, self).__init__(**args)
        self.num_outputs = num_outputs

    def build(self, input_shape):
        num_inputs = input_shape[-1]
        self.kernel = self.add_weight('kernel',
                                      shape=[num_inputs,
                                             self.num_outputs],
                                      initializer=tf.keras.initializers.he_normal()
                                      )
        self.bias = self.add_weight("bias",
                                    shape=[self.num_outputs],
                                    initializer=tf.keras.initializers.Zeros
                                    )

    @tf.function
    def call(self, input):
        num_batch = tf.shape(input)[0]
        num_inputs = tf.shape(input)[-1]
        output = tf.zeros([num_batch, self.num_outputs], dtype=self.dtype)
        for i in tf.range(num_inputs):
            reverse_index = i % self.num_outputs
            kernel_one = tf.matmul(tf.reshape(input[:,i],[num_batch, 1]), tf.reshape(self.kernel[1,:],[1, self.num_outputs]))
            kernel_one = tf.concat([kernel_one[:, reverse_index:],
                                    kernel_one[:, :reverse_index]], axis=1)
            output = tf.add(output, kernel_one)
        return output + self.bias


def main():
    mnist = tf.keras.datasets.mnist

    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0

    model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        MyDenseLayer(128),
        tf.keras.layers.Activation(activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(x_train, y_train, epochs=5)
    model.evaluate(x_test, y_test)


if __name__ == '__main__':
    main()