keras 自定义层input_NLP实战篇之tf2自定义Layer与Model

本文是基于tensorflow2.2.0版本，介绍了自定义Layer和自定义Model。Layer是tf.keras中的核心抽象之一，是状态(权重)与前向传递(输入到输出的计算转换过程)的封装。Model与Layer结构有些类似，但也有明显区别，tf.keras中构建模型主要有三种形式：sequential、subclass、functional，前面文章已经介绍过前两种，本文则介绍后面两种方式。

实战系列篇章中主要会分享，解决实际问题时的过程、遇到的问题或者使用的工具等等。如问题分解、bug排查、模型部署等等。相关代码实现开源在：https://github.com/wellinxu/nlp_store ，更多内容关注知乎专栏(或微信公众号)：NLP杂货铺。

• Layer
  • 权重与计算
  • 延迟创建权重
  • add_loss与add_metric
  • 序列化
  • 递归组合的Layer
• Model
  • subclass
  • functional

Layer

权重与计算

Layer是tf.keras中的核心抽象之一，是状态(权重)与前向传递(输入到输出的计算转换过程)的封装。以简单的全连接层为例，它具有两个权重：w、b，如下面代码所示，可以根据输入维度，在init函数中创建它们，可以使用tf.Variable的方式，也可以使用self.add_weight的方式，后面一种更加简洁。层的计算过程则写在call函数中，call函数需要一个默认的inputs参数，全连接层的计算过程也如下面代码所示。

class Linear(tf.keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__init__()
        # add_weight与下面注释部分是等价的，只是更方便
        self.w = self.add_weight(shape=(input_dim, units),
                                 initializer=tf.random_normal_initializer,
                                 trainable=True)
        self.b = self.add_weight(shape=(units, ),
                                 initializer=tf.zeros_initializer,
                                 trainable=True)
        # w_init = tf.random_normal_initializer()
        # self.w = tf.Variable(initial_value=w_init(shape=(input_dim, units),
        #                                           dtype=tf.float32),
        #                      trainable=True)
        # b_init = tf.zeros_initializer()
        # self.b = tf.Variable(initial_value=b_init(shape=(units,),
        #                                           dtype=tf.float32),
        #                      trainable=True)

    def call(self, inputs, training=None, mask=None):
        #if training:
        #    inputs = tf.nn.dropout(inputs, rate=0.9)
        return tf.matmul(inputs, self.w) + self.b

如上面代码所示，构建权重时可以设置其他参数，比如trainable参数，可以控制该权重是否要被学习(如果为False，则训练过程中不会更新此权重的值)，add_weight的其他参数如下所示：

1. name: 变量名称
2. shape: 变量形状
3. dtype: 变量数据类型，默认为self.dtype 或 float32
4. initializer: 可调用的初始化实例
5. regularizer: 可调用的正则化实例
6. trainable: Boolean, 是否是可训练参数
7. constraint: 可调用的约束实例
8. partitioner:
9. use_resource:
10. synchronization:
11. aggregation:
12. **kwargs: 额外的关键字参数

类似的，call函数也可以有其他参数，比如training，可以区分训练过程还是预测过程，从而来判断是否要进行dropout等处理。当然也可以传递mask或其他关键字参数。

延迟创建权重

在上面实现中，Linear在init函数中就传入了input_dim参数，但在很多情况下，可能实现不知道输入的维度大小，需要在知道输入维度之后的时间点上进行延迟权重创建。tf2中可以使用build函数来处理这种情况，如下面代码所示：

# 动态创建权重，在知道输入形状之后再创建权重
class Linear2(tf.keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear2, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer=tf.random_normal_initializer,
                                 trainable=True)
        self.b = self.add_weight(shape=(self.units, ),
                                 initializer=tf.zeros_initializer,
                                 trainable=True)

    def call(self, inputs, training=None, mask=None):
        self.add_loss(0.5 * tf.reduce_sum(tf.square(self.w)))    # L2正则
        # self.add_metric()
        return tf.matmul(inputs, self.w) + self.b

    def get_config(self):
        return {"units": self.units}

add_loss与add_metric

如上面代码所示，在call函数中可以使用add_loss函数，来添加自定义的loss，不过普通参数的正则化操作可以通过add_weight的regularizer来实现。同样的也可以使用add_metric来添加一些计算过程中需要的度量结果。add_loss与add_metric都可以被模型的fit()函数跟踪，如果是自定义训练过程，则需要手动添加，自定义训练相关内容参考后续文章。

序列化

通过实现get_config方法，可以对层进行序列化，然后使用from_config方法来反序列化，如果反序列化比较复杂，也可以重写from_config方法。Layer的序列化与反序列化简单示例如下：

layer = Linear2(64)
config = layer.get_config()
print(config)
new_layer = Linear2.from_config(config)

递归组合的Layer

Layer对应于文献中的“块”(卷积层、全连接层等)或者“块”(编码块、解码块等)。如下面代码所示，可以将多个Layer组合成一个新的Layer：

# 层之间的递归组合
class MLPBlock(tf.keras.layers.Layer):
    def __init__(self):
        super(MLPBlock, self).__init__()
        self.linear_1 =  Linear2(2)
        self.linear_2 =  Linear2(2)
        self.linear_3 =  Linear2(1)

    def call(self, inputs, **kwargs):
        x = self.linear_1(inputs)
        x = tf.nn.relu(x)
        x = self.linear_2(x)
        x = tf.nn.relu(x)
        return self.linear_3(x)

Model

tf.keras构建模型主要有三种形式：sequential、subclass、functional。前面两种在【NLP实战篇之tensorflow2.0快速入门】上已经做过简单介绍。

subclass

通常我们会用Layer来定义内部计算块，用Model来定义外部模型，Model与Layer具有非常类似的api，但也有着明显的区别：

1. Model内置了训练(fit())、评估(evaluate())和预测(predict())循环。
2. Model通过layers属性获取其内部层的列表。
3. Model具有保存和序列化api(save()、save_weights()...)。

跟自定义Layer类似，我们可以自定义Model，如下面代码所示，init函数中包含了模型的各个子结构(层、块)，call函数中则体现了整个模型从输入到输出的计算流程。

# 使用模型子类化构建模型
class MyModel(tf.keras.models.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = tf.keras.layers.Flatten()
        self.d1 = tf.keras.layers.Dense(128, activation="relu")
        self.dropout = tf.keras.layers.Dropout(0.2)
        self.d2 = tf.keras.layers.Dense(10, activation="softmax")

    def call(self, inputs, training=None, mask=None):
        # inputs: [batch_size, 28, 28]
        x = self.flatten(inputs)    # [batch_size, 28 * 28]
        x = self.d1(x)    # [batch_size, 128]
        x = self.dropout(x)    # [batch_size, 128]
        x = self.d2(x)    # [batch_size, 10]]
        return x

functional

上面我们都是使用的面相对象的方式进行开发的，其实也可以用函数式api来进行构建模型，更棒的是，选择哪种方式其实并不影响，可以将两种方式混合使用，这样可以更加灵活地构建模型。下面代码展示了如何用函数式api进行构建模型，更具一般性的，inputs与outputs可以是多个值(多个输入、输出)。

# 使用函数式api构建模型
inputs = tf.keras.Input(shape=(256,))
emb = tf.keras.layers.Embedding(vocab_size, 16)(inputs)
avg_pool = tf.keras.layers.GlobalAveragePooling1D()(emb)
d1 = tf.keras.layers.Dense(16, activation="relu")(avg_pool)
outputs = tf.keras.layers.Dense(1, activation="sigmoid")(d1)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

参考

https://www.tensorflow.org/guide/keras/custom_layers_and_models