Keras深度学习框架第三讲：通过子类化来创建新的层和模型

一、绪论

本问将详细 讨论构建子类化层和模型的一些内容，主要包括：

• 层（Layer）类
• add_weight() 方法
• 可训练权重和不可训练权重
• build() 方法
• 确保您的层可以与任何后端一起使用
• add_loss() 方法
• call() 方法中的 training 参数
• call() 方法中的 mask 参数
• 确保您的层可以被序列化

二、操作实践

2.1 操作设置

import numpy as np
import keras
from keras import ops
from keras import layers

2.2 层（Layer）类：状态（权重）和计算的结合

在Keras中，层（Layer）类是一个核心抽象。一个层封装了状态（即层的“权重”）以及从输入到输出的转换（即“调用”，也就是层的前向传播）。

下面是一个密集连接层（全连接层）的例子。它有两个状态变量：变量w和b。

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super().__init__()
        self.w = self.add_weight(
            shape=(input_dim, units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(shape=(units,), initializer="zeros", trainable=True)

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b

就像调用Python函数一样可以通过张量输入上调用层来使用它。

x = ops.ones((2, 2))
linear_layer = Linear(4, 2)
y = linear_layer(x)
print(y)

[[ 0.085416   -0.06821361 -0.00741937 -0.03429271]
 [ 0.085416   -0.06821361 -0.00741937 -0.03429271]]

 注意，权重 w 和偏置 b 在被设置为层（layer）的属性时，会自动被该层追踪

assert linear_layer.weights == [linear_layer.w, linear_layer.b]

2.3 不可训练的权重

除了可训练的权重之外，你还可以向层中添加不可训练的权重。这些权重在训练层时不会参与反向传播过程。

以下是添加和使用不可训练权重的方法：

import tensorflow as tf  
from tensorflow.keras.layers import Layer  
  
class MyLayer(Layer):  
    def __init__(self, **kwargs):  
        super(MyLayer, self).__init__(**kwargs)  
  
    def build(self, input_shape):  
        # 添加一个可训练的权重  
        self.trainable_weight = self.add_weight(name='trainable_weight',  
                                                shape=(input_shape[1],),  
                                                initializer='uniform',  
                                                trainable=True)  
          
        # 添加一个不可训练的权重  
        self.non_trainable_weight = self.add_weight(name='non_trainable_weight',  
                                                    shape=(input_shape[1],),  
                                                    initializer='uniform',  
                                                    trainable=False)  
          
        # 一定要调用super().build(input_shape)  
        super(MyLayer, self).build(input_shape)  
  
    def call(self, inputs):  
        # 使用权重执行某些操作，但请注意在训练过程中只有trainable_weight会更新  
        # 例如，一个简单的线性变换  
        output = tf.matmul(inputs, self.trainable_weight) + self.non_trainable_weight  
        return output  
  
# 实例化并使用你的自定义层  
layer = MyLayer()  
# 假设你有一个形状为(batch_size, input_dim)的输入张量x  
# x = ... # 你需要提供一个适当的输入  
# output = layer(x)  
  
# 你可以通过以下方式检查权重是否可训练  
print(layer.trainable_weight.trainable)  # 应该输出True  
print(layer.non_trainable_weight.trainable)  # 应该输出False

2.4 实践经验之在形状确认后再创建权重

在深度学习和神经网络编程中，一种最佳实践是将权重的创建推迟到输入的形状已知之后。在上文的的示例中，Linear层（或称为全连接层）在 __init__方法中接收input_dim作为参数，并使用这个参数来计算权重w和偏置b的形状。然而，这种方法可能不是最优的，因为它假设了输入的形状在创建层时就是已知的。

在实际情况中，特别是在构建动态模型或处理变长输入时，输入的形状可能直到运行时才知道。因此，一种更好的做法是在层的 build方法中创建权重，此时输入的形状会被传递给该方法。build方法在层第一次被调用时（即调用call方法时）自动被调用，并接收输入的形状作为参数。

import tensorflow as tf  
from tensorflow.keras.layers import Layer  
  
class Linear(Layer):  
    def __init__(self, output_dim, **kwargs):  
        super(Linear, self).__init__(**kwargs)  
        self.output_dim = output_dim  
  
    def build(self, input_shape):  
        # Create a trainable weight variable for this layer.  
        self.kernel = self.add_weight(name='kernel',  
                                      shape=(input_shape[1], self.output_dim),  
                                      initializer='uniform',  
                                      trainable=True)  
        # Create a trainable bias variable for this layer.  
        self.bias = self.add_weight(name='bias',  
                                    shape=(self.output_dim,),  
                                    initializer='zeros',  
                                    trainable=True)  
        # Be sure to call this at the end  
        super(Linear, self).build(input_shape)  
  
    def call(self, inputs):  
        return tf.matmul(inputs, self.kernel) + self.bias  
  
# 使用示例  
# 假设 x 是一个形状为 (batch_size, input_dim) 的张量  
# layer = Linear(output_dim=10)  
# output = layer(x)

在Keras API中，当您不知道输入的确切大小，或者您希望推迟到确定输入大小后再创建权重时，推荐在层的build(self, input_shape)方法中创建权重。

build方法是一个特殊的层方法，它会在层第一次被调用（即call方法被调用）之前，并且输入的形状已知时被调用。这允许您根据输入的形状动态地设置权重的大小。

在build方法中，您可以使用self.add_weight方法来创建权重，并指定它们的形状、初始化器、是否可训练等属性。

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super().__init__()
        self.units = units

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

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b

当您创建了一个自定义层并定义了build方法后，该层的方法会在第一次接收到输入时自动调用__call__()方法。这种机制使得层的实现变得“懒惰”（lazy），也就是说，权重和其他内部状态直到需要时才会被创建。

# At instantiation, we don't know on what inputs this is going to get called
linear_layer = Linear(32)

# The layer's weights are created dynamically the first time the layer is called
y = linear_layer(x)

2.5 层递归可组合特性

层（Layer）是递归可组合的，这意味着您可以将一个层实例作为另一个层的属性来创建更复杂的模型结构。当您这样做时，外部层会开始跟踪内部层创建的权重。

在自定义层时，推荐在__init__()方法中创建子层（sublayers），并将它们作为属性分配给外部层。然后，当外部层第一次被调用（即__call__()方法被调用）时，如果内部层还没有构建权重，那么它们的权重就会被自动构建。

class MLPBlock(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.linear_1 = Linear(32)
        self.linear_2 = Linear(32)
        self.linear_3 = Linear(1)

    def call(self, inputs):
        x = self.linear_1(inputs)
        x = keras.activations.relu(x)
        x = self.linear_2(x)
        x = keras.activations.relu(x)
        return self.linear_3(x)


mlp = MLPBlock()
y = mlp(ops.ones(shape=(3, 64)))  # The first call to the `mlp` will create the weights
print("weights:", len(mlp.weights))
print("trainable weights:", len(mlp.trainable_weights))

eights: 6
trainable weights: 6

2.6 后端无关层与后端特定层

框架的设计目标是实现后端无关性（backend-agnostic），这意味着只要层仅使用来自 keras.op 命名空间（或其他 Keras 命名空间keras.activations、keras.random、keras.layers）的 API，那么这些层就可以与任何后端一起使用，无论是 TensorFlow、JAX 还是 PyTorch。

keras.ops命名空间提供了对多种操作的访问，这些操作是神经网络中常见的，包括 NumPy 风格的 API（如 ops.matmul、ops.sum、ops.reshape、ops.stack）以及神经网络特定的 API（如 ops.softmax、ops.conv、ops.binary_crossentropy、ops.relu等）。

然而，您也可以在层中使用特定后端的原生 API（如 TensorFlow 的 tf.nn函数）。但是，如果您这样做，那么您的层将仅能与该特定后端一起使用。例如，如果您使用 JAX 的 jax.numpy来编写一个特定的 JAX 层，那么该层就只能与 JAX 一起使用。

以下是一个使用 jax.numpy编写的 JAX 特定层的示例：

import jax

class Linear(keras.layers.Layer):
    ...

    def call(self, inputs):
        return jax.numpy.matmul(inputs, self.w) + self.b

等效的张量流特定层如下：

import tensorflow as tf

class Linear(keras.layers.Layer):
    ...

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

等效的PyTorch特定层

import torch

class Linear(keras.layers.Layer):
    ...

    def call(self, inputs):
        return torch.matmul(inputs, self.w) + self.b

2.7 add_loss()方法

当编写一个层的call()方法时，确实可以在其中创建损失张量（loss tensors），这些损失张量稍后在编写训练循环时会用到。通过调用self.add_loss(value)，你可以将这些损失张量添加到层的损失列表中。

在Keras框架中，add_loss()方法允许你在层的call()方法内部定义额外的损失项，这些损失项将自动添加到模型的总损失中。这在一些复杂的层（如正则化层、注意力机制层等）中特别有用，其中你可能需要计算额外的损失。

# A layer that creates an activity regularization loss
class ActivityRegularizationLayer(keras.layers.Layer):
    def __init__(self, rate=1e-2):
        super().__init__()
        self.rate = rate

    def call(self, inputs):
        self.add_loss(self.rate * ops.mean(inputs))
        return inputs

这些损失（包括任何内部层创建的损失）可以通过layer.losses这个属性来获取。这个属性在每个顶层层的__call__()方法开始时会被重置，以确保layer.losses总是包含上一次前向传播过程中创建的损失值。

在Keras中，当你定义一个自定义层并在其 call()方法中使用self.add_loss() 添加了额外的损失时，这些损失会被收集起来并存储在层的losses属性中。但是，这个 losses属性是针对特定层的，它只包含该层自身在call()方法中创建的损失。

对于包含多个子层的模型（如 Sequential模型或函数式API构建的模型），你可能想要获取整个模型在前向传播过程中产生的所有损失。在这种情况下，你可以通过调用模型的losses属性来获取这些损失。这个属性包含了模型中所有层在最近一次前向传播过程中通过 add_loss()添加的损失。

这里有一个重要的点需要注意：layer.losses和 model.losses是在每个前向传播过程中动态更新的。这意味着，在每次调用模型或层的__call__()方法时，这些列表都会被清空，然后重新填充在当前前向传播过程中创建的损失。因此，如果你需要在某个地方（比如自定义训练循环中）引用这些损失，你应该确保在调用 __call__()方法之后立即这样做。

class OuterLayer(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.activity_reg = ActivityRegularizationLayer(1e-2)

    def call(self, inputs):
        return self.activity_reg(inputs)


layer = OuterLayer()
assert len(layer.losses) == 0  # No losses yet since the layer has never been called

_ = layer(ops.zeros((1, 1)))
assert len(layer.losses) == 1  # We created one loss value

# `layer.losses` gets reset at the start of each __call__
_ = layer(ops.zeros((1, 1)))
assert len(layer.losses) == 1  # This is the loss created during the call above

loss属性（无论是 layer.losses还是 model.losses）不仅包含了在call()方法中通过self.add_loss()显式添加的损失，还包含了由 Keras 自动计算并添加的权重正则化损失。

class OuterLayerWithKernelRegularizer(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.dense = keras.layers.Dense(
            32, kernel_regularizer=keras.regularizers.l2(1e-3)
        )

    def call(self, inputs):
        return self.dense(inputs)


layer = OuterLayerWithKernelRegularizer()
_ = layer(ops.zeros((1, 1)))

# This is `1e-3 * sum(layer.dense.kernel ** 2)`,
# created by the `kernel_regularizer` above.
print(layer.losses)

[Array(0.00217911, dtype=float32)]

这些损失在编写自定义训练循环时应该被考虑进去。

它们也与fit()方法无缝协作（如果有主损失的话，这些损失会被自动求和并添加到主损失中）

inputs = keras.Input(shape=(3,))
outputs = ActivityRegularizationLayer()(inputs)
model = keras.Model(inputs, outputs)

# If there is a loss passed in `compile`, the regularization
# losses get added to it
model.compile(optimizer="adam", loss="mse")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))

# It's also possible not to pass any loss in `compile`,
# since the model already has a loss to minimize, via the `add_loss`
# call during the forward pass!
model.compile(optimizer="adam")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))

1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 60ms/step - loss: 0.2650
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 15ms/step - loss: 0.0050

<keras.src.callbacks.history.History at 0x146f71960>

2.8 选择性地为层启用序列化功能

在 TensorFlow/Keras 中，如果你希望你的自定义层能够作为函数式模型的一部分进行序列化（例如，当你保存和加载模型时），你可以选择性地实现一个get_config()方法。

get_config()方法是 Keras 层和模型的一个可选方法，它应该返回一个包含层的配置信息的字典。当你保存一个模型时，Keras 会调用这个方法来获取每个层的配置，并将这些信息保存在保存的模型文件中。当你加载模型时，Keras 会使用这些信息来重新创建层。

对于自定义层，你应该在get_config()方法中返回任何你需要重新创建层实例所需的参数。这些参数应该与你的 __init__()方法中的参数相对应。

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super().__init__()
        self.units = units

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

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b

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


# Now you can recreate the layer from its config:
layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)

{'units': 64}

基础Layer类的__init__()方法接受一些关键字参数，特别是name和dtype。将这些参数传递给父类并在__init__()方法中以及在层的配置中，定义这是参数将有效增加程序的可读性。

class Linear(keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super().__init__(**kwargs)
        self.units = units

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

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b

    def get_config(self):
        config = super().get_config()
        config.update({"units": self.units})
        return config


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

{'name': 'linear_7', 'trainable': True, 'dtype': 'float32', 'units': 64}

2.9 在call()方法中的特权训练参数（Privileged training argument）

在深度学习中，有些层（特别是BatchNormalization层和Dropout层）在训练和推理时有不同的行为。对于这些层，在call()方法中暴露training（布尔型）参数是标准做法。

通过在call()方法中暴露这个参数，你可以让内置的训练和评估循环（如fit()方法）在训练和推理时正确地使用这些层。

class CustomDropout(keras.layers.Layer):
    def __init__(self, rate, **kwargs):
        super().__init__(**kwargs)
        self.rate = rate
        self.seed_generator = keras.random.SeedGenerator(1337)

    def call(self, inputs, training=None):
        if training:
            return keras.random.dropout(
                inputs, rate=self.rate, seed=self.seed_generator
            )
        return inputs

2.9 模型类（Model class）

通常，你会使用Layer类来定义内部计算块，并使用Model类来定义外部模型——即你要训练的对象。

例如，在ResNet50模型中，你会有多个继承自Layer的ResNet块，以及一个包含整个ResNet50网络的单一Model。

Model类与Layer类有相同的API，但有以下区别：

1. 它提供了内置的训练、评估和预测循环（model.fit()、model.evaluate()、model.predict()）。
2. 它通过model.layers属性暴露了其内部层的列表。
3. 它提供了保存和序列化的API（save()、save_weights()等）。

实际上，Layer类对应于文献中提到的“层”（如“卷积层”或“循环层”）或“块”（如“ResNet块”或“Inception块”）。

与此同时，Model类对应于文献中提到的“模型”（如“深度学习模型”）或“网络”（如“深度神经网络”）。

所以，如果你正在思考“我应该使用Layer类还是Model类？”，那么你可以这样问自己：我是否需要在这个对象上调用fit()？我是否需要在这个对象上调用save()？如果需要，就选择Model。如果不是（要么是因为你的类只是更大系统中的一个块，要么是因为你自己编写了训练和保存的代码），就使用Layer。

例如，我们可以使用上面提到的mini-resnet示例，并使用它来构建一个可以使用fit()进行训练，并使用save_weights()保存的Model。

class ResNet(keras.Model):

    def __init__(self, num_classes=1000):
        super().__init__()
        self.block_1 = ResNetBlock()
        self.block_2 = ResNetBlock()
        self.global_pool = layers.GlobalAveragePooling2D()
        self.classifier = Dense(num_classes)

    def call(self, inputs):
        x = self.block_1(inputs)
        x = self.block_2(x)
        x = self.global_pool(x)
        return self.classifier(x)


resnet = ResNet()
dataset = ...
resnet.fit(dataset, epochs=10)
resnet.save(filepath.keras)

三、完整的用法示例
通过前面的探讨，我们认识了以下内容：

一个Layer封装了一个状态（在__init__()或build()中创建）和一些计算（在call()中定义）。
层可以递归嵌套以创建新的、更大的计算块。
只要层仅使用Keras API，那么它们就是与后端无关的。你可以使用特定后端的原生API（如jax.numpy、torch.nn或tf.nn），但这样你的层就只能在那个特定的后端上使用了。
层可以通过add_loss()创建和跟踪损失（通常是正则化损失）。
外部的容器，即你想要训练的东西，是一个Model，与Layer类似，但增加了训练和序列化的实用工具。

现在，让我们将这些知识整合到一个端到端的示例中：我们将实现一个变分自编码器（VAE），以一种与后端无关的方式——这样它就可以在TensorFlow、JAX和PyTorch上运行。我们将使用MNIST数字数据集来训练它。

我们的VAE将是Model的一个子类，由嵌套组合的、继承自Layer的层构建而成。它将包含一个正则化损失（KL散度）。

下面是一个简化版的VAE实现框架（不包括具体的训练和编码/解码层细节）：

class Sampling(layers.Layer):
    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.seed_generator = keras.random.SeedGenerator(1337)

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = ops.shape(z_mean)[0]
        dim = ops.shape(z_mean)[1]
        epsilon = keras.random.normal(shape=(batch, dim), seed=self.seed_generator)
        return z_mean + ops.exp(0.5 * z_log_var) * epsilon


class Encoder(layers.Layer):
    """Maps MNIST digits to a triplet (z_mean, z_log_var, z)."""

    def __init__(self, latent_dim=32, intermediate_dim=64, name="encoder", **kwargs):
        super().__init__(name=name, **kwargs)
        self.dense_proj = layers.Dense(intermediate_dim, activation="relu")
        self.dense_mean = layers.Dense(latent_dim)
        self.dense_log_var = layers.Dense(latent_dim)
        self.sampling = Sampling()

    def call(self, inputs):
        x = self.dense_proj(inputs)
        z_mean = self.dense_mean(x)
        z_log_var = self.dense_log_var(x)
        z = self.sampling((z_mean, z_log_var))
        return z_mean, z_log_var, z


class Decoder(layers.Layer):
    """Converts z, the encoded digit vector, back into a readable digit."""

    def __init__(self, original_dim, intermediate_dim=64, name="decoder", **kwargs):
        super().__init__(name=name, **kwargs)
        self.dense_proj = layers.Dense(intermediate_dim, activation="relu")
        self.dense_output = layers.Dense(original_dim, activation="sigmoid")

    def call(self, inputs):
        x = self.dense_proj(inputs)
        return self.dense_output(x)


class VariationalAutoEncoder(keras.Model):
    """Combines the encoder and decoder into an end-to-end model for training."""

    def __init__(
        self,
        original_dim,
        intermediate_dim=64,
        latent_dim=32,
        name="autoencoder",
        **kwargs
    ):
        super().__init__(name=name, **kwargs)
        self.original_dim = original_dim
        self.encoder = Encoder(latent_dim=latent_dim, intermediate_dim=intermediate_dim)
        self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim)

    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstructed = self.decoder(z)
        # Add KL divergence regularization loss.
        kl_loss = -0.5 * ops.mean(
            z_log_var - ops.square(z_mean) - ops.exp(z_log_var) + 1
        )
        self.add_loss(kl_loss)
        return reconstructed

然后使用fit() API在MNIST数据集上训练：

(x_train, _), _ = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype("float32") / 255

original_dim = 784
vae = VariationalAutoEncoder(784, 64, 32)

optimizer = keras.optimizers.Adam(learning_rate=1e-3)
vae.compile(optimizer, loss=keras.losses.MeanSquaredError())

vae.fit(x_train, x_train, epochs=2, batch_size=64)

Epoch 1/2
 938/938 ━━━━━━━━━━━━━━━━━━━━ 2s 1ms/step - loss: 0.0942
Epoch 2/2
 938/938 ━━━━━━━━━━━━━━━━━━━━ 1s 859us/step - loss: 0.0677

<keras.src.callbacks.history.History at 0x146fe62f0>