作者: fchollet
创建时间 2019/03/01
上次修改时间 2023/06/25
描述:从头开始编写 Layer
和 Model
对象的完整指南。
本指南将涵盖构建您自己的子类化层和模型所需的所有内容。特别是,您将了解以下功能
Layer
类add_weight()
方法build()
方法add_loss()
方法call()
中的 training
参数call()
中的 mask
参数让我们开始吧。
import numpy as np
import keras
from keras import ops
from keras import layers
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
在被设置为层属性时会自动被层跟踪。
assert linear_layer.weights == [linear_layer.w, linear_layer.b]
除了可训练的权重之外,您还可以向层添加不可训练的权重。此类权重旨在在反向传播期间不予考虑,当您训练层时。
以下是添加和使用不可训练权重的方法
class ComputeSum(keras.layers.Layer):
def __init__(self, input_dim):
super().__init__()
self.total = self.add_weight(
initializer="zeros", shape=(input_dim,), trainable=False
)
def call(self, inputs):
self.total.assign_add(ops.sum(inputs, axis=0))
return self.total
x = ops.ones((2, 2))
my_sum = ComputeSum(2)
y = my_sum(x)
print(y.numpy())
y = my_sum(x)
print(y.numpy())
[2. 2.]
[4. 4.]
它是 layer.weights
的一部分,但它被归类为不可训练的权重。
print("weights:", len(my_sum.weights))
print("non-trainable weights:", len(my_sum.non_trainable_weights))
# It's not included in the trainable weights:
print("trainable_weights:", my_sum.trainable_weights)
weights: 1
non-trainable weights: 1
trainable_weights: []
我们上面的 Linear
层接受了一个 input_dim
参数,该参数用于在 __init__()
中计算权重 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
在许多情况下,您可能事先不知道输入的大小,并且希望在该值已知时延迟创建权重,即在实例化层一段时间后。
在 Keras API 中,我们建议在层的 build(self, inputs_shape)
方法中创建层权重。像这样
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
您的层的 __call__()
方法将在第一次调用时自动运行构建。您现在拥有一个惰性层,因此更易于使用。
# 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)
如上所示,单独实现 build()
可以很好地将仅创建一次权重与每次调用中使用权重区分开来。
如果您将 Layer 实例分配为另一个 Layer 的属性,则外部层将开始跟踪内部层创建的权重。
我们建议在 __init__()
方法中创建此类子层,并将其留给第一次 __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))
weights: 6
trainable weights: 6
只要层仅使用 keras.ops
命名空间(或其他 Keras 命名空间,如 keras.activations
、keras.random
或 keras.layers
)中的 API,那么它就可以与任何后端一起使用——TensorFlow、JAX 或 PyTorch。
本指南中您迄今为止看到的全部层都适用于所有 Keras 后端。
keras.ops
命名空间允许您访问
ops.matmul
、ops.sum
、ops.reshape
、ops.stack
等。ops.softmax
、ops
.conv、
ops.binary_crossentropy、
ops.relu` 等。您也可以在您的层中使用特定于后端的 API(例如 tf.nn
函数),但是如果您这样做,则您的层将只能与相关后端一起使用。例如,您可以使用 jax.numpy
编写以下特定于 JAX 的层
import jax
class Linear(keras.layers.Layer):
...
def call(self, inputs):
return jax.numpy.matmul(inputs, self.w) + self.b
这将是等效的特定于 TensorFlow 的层
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
由于跨后端兼容性是一个非常有用的特性,因此我们强烈建议您始终努力通过仅利用 Keras API 来使您的层与后端无关。
add_loss()
方法在编写层的 call()
方法时,您可以创建稍后在编写训练循环时想要使用的损失张量。这可以通过调用 self.add_loss(value)
来实现。
# 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
始终包含上次正向传递期间创建的损失值。
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
属性还包含为任何内部层的权重创建的正则化损失。
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>
如果您需要将自定义层作为 函数式模型 的一部分进行序列化,则可以选择实现 get_config()
方法。
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}
如果您在从其配置反序列化层时需要更多灵活性,您还可以覆盖 from_config()
类方法。这是 from_config()
的基本实现。
def from_config(cls, config):
return cls(**config)
要了解有关序列化和保存的更多信息,请参阅完整的 模型保存和序列化指南。
call()
方法中的特权 training
参数某些层,特别是 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
call()
方法中的特权 mask
参数call()
支持的另一个特权参数是 mask
参数。
您将在所有 Keras RNN 层中找到它。掩码是一个布尔张量(输入中每个时间步长一个布尔值),用于在处理时间序列数据时跳过某些输入时间步长。
当先前层生成掩码时,Keras 会自动将正确的 mask
参数传递给支持它的层的 __call__()
。生成掩码的层是配置为 mask_zero=True
的 Embedding
层和 Masking
层。
Model
类通常,您将使用 Layer
类定义内部计算块,并使用 Model
类定义外部模型——您将训练的对象。
例如,在 ResNet50 模型中,您将有几个继承 Layer
的 ResNet 块,以及一个包含整个 ResNet50 网络的 Model
。
Model
类与 Layer
具有相同的 API,但存在以下差异
model.fit()
、model.evaluate()
、model.predict()
)。model.layers
属性公开了其内部层的列表。save()
、save_weights()
...)。实际上,Layer
类对应于我们在文献中称为“层”(如“卷积层”或“循环层”)或“块”(如“ResNet 块”或“Inception 块”)。
同时,Model
类对应于我们在文献中称为“模型”(如“深度学习模型”)或“网络”(如“深度神经网络”)。
因此,如果您想知道“我应该使用 Layer
类还是 Model
类?”,请自问:我需要对其调用 fit()
吗?我需要对其调用 save()
吗?如果是,则使用 Model
。如果不是(要么是因为您的类只是更大系统中的一个块,要么是因为您自己编写了训练和保存代码),则使用 Layer
。
例如,我们可以采用上面的迷你 ResNet 示例,并用它来构建一个 Model
,我们可以用 fit()
训练它,并用 save_weights()
保存它。
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()
中定义)。jax.numpy
、torch.nn
或 tf.nn
),但随后您的层将只能与该特定后端一起使用。add_loss()
创建和跟踪损失(通常是正则化损失)。Model
。Model
就像一个Layer
,但增加了训练和序列化工具。让我们将所有这些内容整合到一个端到端的示例中:我们将以一种与后端无关的方式实现一个变分自动编码器(VAE)——以便它可以在 TensorFlow、JAX 和 PyTorch 上运行相同。我们将用 MNIST 数字对其进行训练。
我们的 VAE 将是Model
的一个子类,构建为Layer
子类的层的嵌套组合。它将包含一个正则化损失(KL 散度)。
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>