使用TensorFlow的时候，经常会遇到要自己构建层的问题，下面是一个标准的自定义层的demo，方便后续复制。

• 构造函数__init__：会接受各种超参数，还会接受**kwargs参数，需要把kwargs传给父类。
  kwargs包括input_shape, trainable还有name等等。
• build()：创建权重，第一次使用自定义层的时候会调用它，来保存权重。需要提供改成的形状。最后需要调用父类的build()，来告知父类这一层被build了，其中只是设置了built=True
• call()：执行计算操作
• compute_output_shape()：返回这层的形状，一般是最后一个维度被换成当前层神经元的数量。
• get_config()：调用activations.serialize()来保存激活函数的完整配置

首先导入包，以及加载mnist数据集

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers
from keras.datasets import mnist

train_X, train_y = mnist.load_data()[0]
train_y = keras.utils.to_categorical(train_y, num_classes=10)

接下来构建一个自定义的层，这个层的功能和Dense完全相同

class MyDense(keras.layers.Layer):
    def __init__(self, units, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.activation = keras.activations.get(activation)
    
    def build(self, batch_input_shape):
        self.kernel = self.add_weight(
            name='kernel', shape=[batch_input_shape[-1], self.units],
            initializer='glorot_normal')
        self.bias = self.add_weight(
            name='bias', shape=[self.units], initializer='zeros')
        super().build(batch_input_shape)
    
    def call(self, X):
        return self.activation(X @ self.kernel + self.bias)
    
    def compute_output_shape(self, batch_input_shape):
        return tf.TensorShape(batch_input_shape.as_list()[:-1] + [self.units])
    
    def get_config(self):
        base_config = super().get_config()
        return {**base_config, 'units' : self.units,
                   'activation' : keras.activations.serialize(self.activation)}
        

最后可以尝试用自定义层来做预测

model = keras.Sequential()
model.add(layers.Flatten(input_shape=[28, 28]))
model.add(MyDense(28*28, activation='relu'))
model.add(MyDense(10, activation='softmax'))

model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['acc'])

model.fit(train_X, train_y, epochs=10)