TensorFlow通过子类化创建新的层和模型

通过子类化创建新的层和模型

通过本篇文章，你可以

学会通过继承tf.keras.layers.Layer类自定义层，和继承tf.keras.Model类自定神经网络模型

通过子类化创建新的层

创建步骤

# 导入相应的工具包
import tensorflow as tf
from tensorflow import keras

# 新创建的层需要继承keras.layers.Layer类
class Linear(keras.layers.Layer):
	# 在初始化函数里面定义该层所需要的权重，权重的形状、初始化、类型、是否可训练都可以定义
	def __init__(self, units=32, input_dim=32):
		super(Linear, self).__init__()  # 继承需要添加的语句super
		# 可以自定义权重的初始化
		w_init = tf.random_normal_initializer()
		# 可以自定义权重的形状，类型，是否可训练
		self.w = tf.Variable(
			initial_value = w_init(shape=(input_dim, units), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
		# 可以自定义权重的初始化
		b_init = tf.zeros_initializer()
		self.b = tf.Variable(
			initial_value = b_init(shape=( units, ), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
	# 通过call函数来定义层的运算逻辑,即前向计算逻辑
	def call(self, inputs):
		y = tf.matmul(inputs, self.w) + self.b
		return y
		
		

创建的层的权重如何提取？

linear_layer = Linear(units=8, input_dim=4)
# 获取层所有可训练的权重，结果是个列表
trainable_weights = linear.trainable_weights
# 获取层所有不可训练的权重，结果是个列表
non_trainable_weights = linear.non_trainable_weights
# 获取层所有权重，不管是否可以训练，结果是个列表
weights = linear.weights

层实例化之后再创建权重

上面的创建层的步骤需要事先知道输入的形状，其实在大多数情况中，你可能事先不知道输入的形状大小。所以我们希望可以在实例化层之后，再创建层的权重。该怎么做呢？你可以在层的build(self, inputs_shape)方法中创建层的权重，inputs_shape这个参数表示的就是输入的形状。

修改上面的例子，变成如下

# 新创建的层需要继承keras.layers.Layer类
class Linear(keras.layers.Layer):
	# 在初始化函数里面定义该层所需要的权重，权重的形状、初始化、类型、是否可训练都可以定义
	def __init__(self, units=32):
		super(Linear, self).__init__()  # 继承需要添加的语句super
		self.units = units
    def build(self, inputs_shape):
        # 可以自定义权重的初始化
		w_init = tf.random_normal_initializer()
		# 可以自定义权重的形状，类型，是否可训练
		self.w = tf.Variable(
			initial_value = w_init(shape=(inputs_shape[-1], self.units), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
		# 可以自定义权重的初始化
		b_init = tf.zeros_initializer()
		self.b = tf.Variable(
			initial_value = b_init(shape=( self.units, ), dtype='float32'), # 权重的形状，类型
			trainable = True # 是否可训练
		)
	# 通过call函数来定义层的运算逻辑,即前向计算逻辑
	def call(self, inputs):
		y = tf.matmul(inputs, self.w) + self.b
		return y

这是一种比较实际的做法，因为你只需创建一次实例化层，然后每次使用该层的时候，权重都会重新创建，所以权重都不一样。这是一种比较合理的做法

层可以递归组合，即层里面可以嵌套层

如果将层实例作为另一个层的组件，则外部层会跟踪内部层创建的权重

一般在外部层的初始化方法中创建子层，并在call方法中进行逻辑计算

下面创建一个外部层，以MLP模块进行示例

class MLPBlock(keras.layers.Layer):
	# 在初始化函数里面实例化子层
	def __init__(self):
        super(MLPBlock, self).__init__() 
        # 下面每个层的权重都不一样
        self.linear1 = Linear(32)  
        self.linear2 = Linear(32)
        self.linear3 = Linear(1)
    # 在call方法中进行前向逻辑计算
    def call(self, inputs):
        x = self.linear_1(inputs)
        x = tf.nn.relu(x)
        x = self.linear_2(x)
        x = tf.nn.relu(x)
        y = self.linear_3(x)
        return y
mlp = MLPBlock()
y = mlp(tf.ones(shape=(3, 64)))  # 第一次调用层的时候才会创建权重
print("weights:", len(mlp.weights))  # 有六个权重参数
print("trainable weights:", len(mlp.trainable_weights)) # 所有权重参数都可以训练

给层添加自定义的损失张量，add_loss

在自定义层的时候，可以在call方法中添加在训练时可以用上的损失张量，该操作可以通过self.add_loss(value)来实现。这些损失，包括由任何内部层创建的损失）可通过layer.losses取回，layer.losses属性还包含任何内部层的权重创建的正则化损失。

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

    def call(self, inputs):
        self.add_loss(self.rate * tf.reduce_sum(inputs))  # 内部层自定义的损失，第一个损失
        return inputs

class OuterLayerWithKernelRegularizer(keras.layers.Layer):
    def __init__(self):
        super(OuterLayerWithKernelRegularizer, self).__init__()
        self.dense = keras.layers.Dense(
            32, kernel_regularizer=tf.keras.regularizers.l2(1e-3)  # 内部层权重定义的正则化损失，第二个损失
        )
        self.activity_reg = ActivityRegularizationLayer(1e-2)

    def call(self, inputs):
        self.add_loss(0.01*tf.reduce_sum(inputs))  # 外部层自定义的损失，第三个损失
        x = self.dense(inputs)
        return self.activity_reg(x)

layer = OuterLayerWithKernelRegularizer()
_ = layer(tf.ones((1, 1)))

print(len(layer.losses)  # 总共有三个损失值

给层添加自定义的评价指标，add_metric

本质、原理和属性跟add_loss一样，请自行验证

给call方法指定training参数

某些层，尤其是 BatchNormalization 层和 Dropout 层，在训练和推断期间具有不同的行为。对于此类层，标准做法是在 call() 方法中公开 training（布尔）参数。

通过在 call() 中公开此参数，可以启用内置的训练和评估循环（例如 fit()）以在训练和推断中正确使用层。

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

    def call(self, inputs, training=None):
        if training:
            return tf.nn.dropout(inputs, rate=self.rate)
        return inputs
    

通过子类化创建新的模型

什么是Model类

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

例如，在 ResNet50 模型中，您会有几个子类化 Layer 的 ResNet 块，以及一个包含整个 ResNet50 网络的 Model。

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

- 它会公开内置训练、评估和预测循环（model.fit()、model.evaluate()、model.predict()）。

- 它会通过 model.layers 属性公开其内部层的列表。

- 它会公开保存和序列化 API（save()、save_weights()…）

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

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

因此，如果您想知道“我应该用 Layer 类还是 Model 类？”，请问自己：我是否需要在它上面调用 fit()？我是否需要在它上面调用 save()？如果是，则使用 Model。如果不是（要么因为您的类只是更大系统中的一个块，要么因为您正在自己编写训练和保存代码），则使用 Layer。

如何通过子类化来创建新的模型

# 导入相应的工具包
import tensorflow as tf
from tensorflow import keras

# 继承tf.keras.Model
class ResNet(tf.keras.Model):
    # 初始化的时候定义后续要使用的层和块
    def __init__(self, num_classes=1000)
        super(RetNet, self).__init__()
        # 下面都是实现定义好的层
        self.block_1 = ResNetBlock()
        self.block_2 = ResNetBlock()
        self.global_pool = layers.GlobalAveragePooling2D()
        self.classifier = Dense(num_classes)
	# 在call方法中定义前向传播逻辑
    def call(self, inputs):
        x = self.block_1(inputs)
        x = self.block_2(x)
        x = self.global_pool(x)
        return self.classifier(x)

通过自定义层和模型来创建模型并进行训练的完整流程

到目前为止，您已学习以下内容：

• Layer 封装了状态（在 __init__() 或 build() 中创建）和一些计算（在 call() 中定义）。
• 层可以递归嵌套以创建新的更大的计算块。
• 层可以通过 add_loss() 和 add_metric() 创建并跟踪损失（通常是正则化损失）以及指标。
• 您要训练的外部容器是 Model。Model 就像 Layer，但是添加了训练和序列化实用工具。

让我们将这些内容全部汇总到一个端到端示例：我们将实现一个变分自动编码器 (VAE)，并用 MNIST 数字对其进行训练。

我们的 VAE 将是 Model 的一个子类，它是作为子类化 Layer 的嵌套组合层进行构建的。它将具有正则化损失（KL 散度）。

创建模型

from tensorflow.keras import layers


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

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.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(Encoder, self).__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(Decoder, self).__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(VariationalAutoEncoder, self).__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 * tf.reduce_mean(
            z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1
        )
        self.add_loss(kl_loss)
        return reconstructed

训练模型

让我们在 MNIST 上编写一个简单的训练循环：

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

# 定义优化器
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3) 
# 定义损失函数
mse_loss_fn = tf.keras.losses.MeanSquaredError()

# 定义平均指标
loss_metric = tf.keras.metrics.Mean()

# 加载数据
(x_train, _), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype("float32") / 255

train_dataset = tf.data.Dataset.from_tensor_slices(x_train)
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64)

# 训练两个循环
epochs = 2

# Iterate over epochs.
for epoch in range(epochs):
    print("Start of epoch %d" % (epoch,))

    # 通过批次训练模型
    for step, x_batch_train in enumerate(train_dataset):
        # 创建梯度收集器
        with tf.GradientTape() as tape:
            reconstructed = vae(x_batch_train) # 通过模型前向传播计算结果
            # 通过模型结果和输入值计算损失值
            loss = mse_loss_fn(x_batch_train, reconstructed)
            loss += sum(vae.losses)  # 自定义层的损失值添加到模型损失值里面
		
        grads = tape.gradient(loss, vae.trainable_weights)  # 通过该损失值和所有权重值计算所有权重的梯度
        optimizer.apply_gradients(zip(grads, vae.trainable_weights))  # 通过权重的梯度值更新权重值

        loss_metric(loss)  # 记录损失值和平均指标

        if step % 100 == 0:
            print("step %d: mean loss = %.4f" % (step, loss_metric.result()))

Start of epoch 0
step 0: mean loss = 0.3435
step 100: mean loss = 0.1252
step 200: mean loss = 0.0991
step 300: mean loss = 0.0891
step 400: mean loss = 0.0842
step 500: mean loss = 0.0808
step 600: mean loss = 0.0787
step 700: mean loss = 0.0771
step 800: mean loss = 0.0759
step 900: mean loss = 0.0749
Start of epoch 1
step 0: mean loss = 0.0746
step 100: mean loss = 0.0740
step 200: mean loss = 0.0735
step 300: mean loss = 0.0730
step 400: mean loss = 0.0727
step 500: mean loss = 0.0723
step 600: mean loss = 0.0720
step 700: mean loss = 0.0717
step 800: mean loss = 0.0715
step 900: mean loss = 0.0712

文章参考官方链接：https://www.tensorflow.org/guide/keras/custom_layers_and_models