自定义层、损失函数和评估指标 *

可能你还会问，如果现有的这些层无法满足我的要求，我需要定义自己的层怎么办？事实上，我们不仅可以继承 tf.keras.Model 编写自己的模型类，也可以继承 tf.keras.layers.Layer 编写自己的层。

自定义层

自定义层需要继承 tf.keras.layers.Layer 类，并重写 init 、 build 和 call 三个方法，如下所示：

  1. class MyLayer(tf.keras.layers.Layer):
  2.     def __init__(self):
  3.         super().__init__()
  4.         # 初始化代码
  5.  
  6.     def build(self, input_shape):     # input_shape 是一个 TensorShape 类型对象，提供输入的形状
  7.         # 在第一次使用该层的时候调用该部分代码，在这里创建变量可以使得变量的形状自适应输入的形状
  8.         # 而不需要使用者额外指定变量形状。
  9.         # 如果已经可以完全确定变量的形状，也可以在__init__部分创建变量
  10.         self.variable_0 = self.add_weight(...)
  11.         self.variable_1 = self.add_weight(...)
  12.  
  13.     def call(self, inputs):
  14.         # 模型调用的代码（处理输入并返回输出）
  15.         return output

例如，如果我们要自己实现一个 本章第一节 中的全连接层（ tf.keras.layers.Dense ），可以按如下方式编写。此代码在 build 方法中创建两个变量，并在 call 方法中使用创建的变量进行运算：

  1. class LinearLayer(tf.keras.layers.Layer):
  2.     def __init__(self, units):
  3.         super().__init__()
  4.         self.units = units
  5.  
  6.     def build(self, input_shape):     # 这里 input_shape 是第一次运行call()时参数inputs的形状
  7.         self.w = self.add_variable(name='w',
  8.             shape=[input_shape[-1], self.units], initializer=tf.zeros_initializer())
  9.         self.b = self.add_variable(name='b',
  10.             shape=[self.units], initializer=tf.zeros_initializer())
  11.  
  12.     def call(self, inputs):
  13.         y_pred = tf.matmul(inputs, self.w) + self.b
  14.         return y_pred

在定义模型的时候，我们便可以如同 Keras 中的其他层一样，调用我们自定义的层 LinearLayer：

  1. class LinearModel(tf.keras.Model):
  2.     def __init__(self):
  3.         super().__init__()
  4.         self.layer = LinearLayer(units=1)
  5.  
  6.     def call(self, inputs):
  7.         output = self.layer(inputs)
  8.         return output

自定义损失函数和评估指标

自定义损失函数需要继承 tf.keras.losses.Loss 类，重写 call 方法即可，输入真实值 y_true 和模型预测值 y_pred ，输出模型预测值和真实值之间通过自定义的损失函数计算出的损失值。下面的示例为均方差损失函数：

  1. class MeanSquaredError(tf.keras.losses.Loss):
  2.     def call(self, y_true, y_pred):
  3.         return tf.reduce_mean(tf.square(y_pred - y_true))

自定义评估指标需要继承 tf.keras.metrics.Metric 类，并重写 init 、 update_state 和 result 三个方法。下面的示例对前面用到的 SparseCategoricalAccuracy 评估指标类做了一个简单的重实现：

  1. class SparseCategoricalAccuracy(tf.keras.metrics.Metric):
  2.     def __init__(self):
  3.         super().__init__()
  4.         self.total = self.add_weight(name='total', dtype=tf.int32, initializer=tf.zeros_initializer())
  5.         self.count = self.add_weight(name='count', dtype=tf.int32, initializer=tf.zeros_initializer())
  6.  
  7.     def update_state(self, y_true, y_pred, sample_weight=None):
  8.         values = tf.cast(tf.equal(y_true, tf.argmax(y_pred, axis=-1, output_type=tf.int32)), tf.int32)
  9.         self.total.assign_add(tf.shape(y_true)[0])
  10.         self.count.assign_add(tf.reduce_sum(values))
  11.  
  12.     def result(self):
  13.         return self.count / self.total
  • LeCun1998
    • LeCun, L. Bottou, Y. Bengio, and P. Haffner. “Gradient-based learning applied to document recognition.” Proceedings of the IEEE, 86(11):2278-2324, November 1998. http://yann.lecun.com/exdb/mnist/

  • Graves2013

  • Graves, Alex. “Generating Sequences With Recurrent Neural Networks.” ArXiv:1308.0850 [Cs], August 4, 2013. http://arxiv.org/abs/1308.0850.

  • Mnih2013

  • Mnih, Volodymyr, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. “Playing Atari with Deep Reinforcement Learning.” ArXiv:1312.5602 [Cs], December 19, 2013. http://arxiv.org/abs/1312.5602.