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.