TensorFlow2自定义

自定义层

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

class MyLayer(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        # 初始化代码

    def build(self, input_shape): 
        #首次调用代码
        super.build(input_shape)

    def call(self, inputs):
        # 模型调用的代码）
        return output

关于build()方法

build()方法将会在第一次调用实例时执行，传入参数input_shape内储存有首次输入数据的形状信息（层的一个特性就是首次调用时自动确定权重类型），通过input_shape我们得到输入数据的形状后，应通过self.kernel = self.add_weight()的方法为层添加权重矩阵或偏置矩阵，最后调用super.build(input_shape)结束层的构建。

如果已经可以完全确定变量的形状，也可以在__init__部分创建变量，不选择对build的重写。

add_weight()方法

使用 self.add_weight()方法创建权重矩阵，可自动完成一些层内的配置，比如可被get_weights()获取到，该方法说明如下：

add_weight(
    name=None, shape=None, dtype=None, initializer=None, regularizer=None,
    trainable=None, constraint=None, partitioner=None, use_resource=None,
    synchronization=tf.VariableSynchronization.AUTO,
    aggregation=tf.compat.v1.VariableAggregation.NONE, **kwargs
)

常用参数：

• name: 张量名称
• shape: 张量形状，默认是张量
• dtype: 张量类型默认是 self.dtype 或 float32.
• initializer: 初始化器
• regularizer: 正则化器
• trainable: Boolean, 能否被优化器优化

返回：

返回一个Variable类型的张量

自定义实现全连接层

class LinearLayer(tf.keras.layers.Layer):
    def __init__(self, units):
        super().__init__()
        self.units = units

    def build(self, input_shape):     
        #w参数张量,实现矩阵乘法只需要获取输入张量最后一个维度的长度,即input_shape[-1]
        self.w = self.add_weight(name='w',
            shape=[input_shape[-1], self.units], initializer=tf.zeros_initializer())
        #b从参数张量,形状为units,执行张量加法时可通过broadcast自动扩展
        self.b = self.add_weight(name='b',
            shape=[self.units], initializer=tf.zeros_initializer())

    def call(self, inputs):
        #y=wx+b
        y_pred = tf.matmul(inputs, self.w) + self.b
        return y_pred

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

class LinearModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.layer = LinearLayer(units=1)

    def call(self, inputs):
        output = self.layer(inputs)
        return output

自定义损失函数

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

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

自定义评价函数

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

class SparseCategoricalAccuracy(tf.keras.metrics.Metric):
    def __init__(self):
        super().__init__()
        self.total = self.add_weight(name='total', dtype=tf.int32, initializer=tf.zeros_initializer())
        self.count = self.add_weight(name='count', dtype=tf.int32, initializer=tf.zeros_initializer())

    def update_state(self, y_true, y_pred, sample_weight=None):
        values = tf.cast(tf.equal(y_true, tf.argmax(y_pred, axis=-1, output_type=tf.int32)), tf.int32)
        self.total.assign_add(tf.shape(y_true)[0]) #assign_add 对应元素相加
        self.count.assign_add(tf.reduce_sum(values))

    def result(self):
        return self.count / self.total

本文摘录、翻译自

Keras中文文档
https://keras.io/zh/

TensorFlow官方
https://tensorflow.google.cn/

最全Tensorflow2.0 入门教程持续更新
https://zhuanlan.zhihu.com/p/59507137

简单粗暴TensorFlow2
https://tf.wiki/zh/preface.html