Keras - Customized Layer

Keras - 自定义层

虽然 tf.keras.layers 提供了全面的内置层集合，但您可能会遇到需要具有特定、非标准功能的层的情况。这可能是为了实现新颖的研究思路，创建一个具有独特状态行为的层，或者仅仅是为了以可重用方式组合操作。

tf.keras 允许您通过子类化基础 tf.keras.layers.Layer 类来创建自己的自定义层。一旦定义，您的自定义层就可以像任何内置层一样在 Sequential 或 Functional 模型中无缝使用。

自定义层的结构

要创建自定义层，您通常需要在子类中实现三个关键方法：

__init__(self, **kwargs)： （构造函数）用于执行独立于输入形状的初始化。在这里定义任何非权重属性或特定于您的层的超参数。务必调用 super().__init__(**kwargs)。
build(self, input_shape)： 这个方法在层第一次使用时自动调用，此时已知输入张量的形状。这是使用 self.add_weight() 方法创建层的可训练权重的标准位置。input_shape 参数提供输入张量的形状。确保在最后调用 super().build(input_shape)。
call(self, inputs, **kwargs)： 这个方法定义了层的前向传播逻辑——层执行的核心计算。它接收输入张量作为参数并返回输出张量。在这里，您使用 TensorFlow 操作或调用其他层来实现层的变换。
get_config(self) (可选，但推荐)： 为了让您的层（以及包含它的模型）能够正确保存和加载，请实现 get_config。这个方法应返回一个字典，包含重新创建层实例所需的构造函数参数。同时包括父类的配置。

示例：一个简单的全连接层

让我们从头开始创建一个基本的自定义 Dense（全连接）层，以说明整个过程。这个层将执行操作：output = activation(dot(input, kernel) + bias)。

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import activations # 导入 activations 模块

class SimpleDense(layers.Layer):
    """一个基本的自定义 Dense 层。"""

    def __init__(self, units, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        # 从标识符（字符串或函数）获取激活函数
        self.activation = activations.get(activation)

    def build(self, input_shape):
        # input_shape 是一个 TensorShape 对象（例如，(batch_size, input_dim)）
        input_dim = input_shape[-1]

        # 添加可训练的权重变量 kernel
        self.kernel = self.add_weight(
            name='kernel',
            shape=(input_dim, self.units),
            initializer='glorot_uniform', # Dense 层的常用初始化器
            trainable=True
        )

        # 添加可训练的权重变量 bias
        self.bias = self.add_weight(
            name='bias',
            shape=(self.units,),
            initializer='zeros',
            trainable=True
        )

        # 确保调用父类的 build 方法
        super().build(input_shape)

    def call(self, inputs):
        # 定义前向传播逻辑
        # 使用 tf.matmul 进行矩阵乘法
        z = tf.matmul(inputs, self.kernel) + self.bias

        # 如果指定了激活函数，则应用
        if self.activation is not None:
            return self.activation(z)
        return z

    def compute_output_shape(self, input_shape):
        # 可选方法，用于符号计算输出形状
        # Keras 通常可以推断出来，但定义它可以提供帮助
        return (input_shape[0], self.units)

    def get_config(self):
        # 使层可序列化
        config = super().get_config()
        config.update({
            'units': self.units,
            'activation': activations.serialize(self.activation) # 序列化激活函数
        })
        return config

使用自定义层

现在您可以像使用内置 Keras 层一样使用 SimpleDense 了：

# 示例用法
model = tf.keras.Sequential([
    tf.keras.Input(shape=(16,)), # 输入层
    SimpleDense(units=32, activation='relu'), # 使用自定义层
    layers.Dropout(0.1), # 可以与标准层混合使用
    SimpleDense(units=8, activation='softmax') # 再次使用自定义层
])

# 像往常一样编译和训练
# model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.summary()

运行 model.summary() 将显示您的自定义层（simple_dense）与标准层一起列出。

# 示例摘要输出：
# Model: "sequential"
# _________________________________________________________________
#  Layer (type)                Output Shape              Param #
# =================================================================
#  simple_dense (SimpleDense)  (None, 32)                544
#
#  dropout (Dropout)           (None, 32)                0
#
#  simple_dense_1 (SimpleDen   (None, 8)                 264
#  se)
#
# =================================================================
# Total params: 808 (3.16 KB)
# Trainable params: 808 (3.16 KB)
# Non-trainable params: 0 (0.00 Byte)
# _________________________________________________________________

何时创建自定义层

实现研究论文中的新颖层类型。
创建带有需要与标准权重区别管理的内部状态的层。
封装一个复杂、可重用的操作块。
现有层无法提供确切所需的功能或灵活性时。

子类化 tf.keras.layers.Layer 提供了一种强大的机制，用于扩展 Keras 并根据特定的研究或应用需求进行定制。