tensorflow2 自定义层

介绍

有时需要自定义网络层,实现自己想要的功能

官方网址:
https://www.tensorflow.org/tutorials/customization/custom_layers?hl=zh_cn

keras网址:
https://keras-cn.readthedocs.io/en/latest/layers/writting_layer/

更多注意事项参见:
https://zhuanlan.zhihu.com/p/86886620

方法

__init__

初始化

build()

build(input_shape)

这是定义权重的方法,可训练的权应该在这里被加入列表self.trainable_weights中。其他的属性还包括self.non_trainabe_weights(列表)和self.updates(需要更新的形如(tensor, new_tensor)的tuple的列表)。你可以参考BatchNormalization层的实现来学习如何使用上面两个属性。这个方法必须设置self.built = True,可通过调用super([layer],self).build()实现

call()

call(x)

这是定义层功能的方法,除非你希望你写的层支持masking,否则你只需要关心call的第一个参数:输入张量

例子

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

  def build(self, input_shape):
    self.kernel = self.add_weight("kernel",
                                  shape=[int(input_shape[-1]),
                                         self.num_outputs])

  def call(self, input):
    return tf.matmul(input, self.kernel)

layer = MyDenseLayer(10)