通过子类来创建新的层和模型 making new layers and models via subclassing

通过子类来创建新的层和模型 making new layers and models via subclassing

• 翻译者 anakin
• 原文参考连接 Making new layers and models via subclassing
• 原文更新日期 2020/04/13
• 原始网页 https://keras.io/guides/making_new_layers_and_models_via_subclassing/

1. 准备

import tensorflow as tf
from tensorflow import keras

2. 层类：将状态（权重参数）和一些计算进行组合，即为层类

Keras的核心抽象对象之一就是层类。一个层类封装了状态（权重参数）和输入到输出的转换。

这是一个全连接层，它的状态：变量$\boxed{w}$ 和$\boxed{h}$。

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__init__()
        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
        )

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

使用一些张量作为入参调用该层对象进行计算，类似调用python函数。

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

打印结果

tf.Tensor(
[[ 0.01013444 -0.01070027 -0.01888977  0.05208318]
 [ 0.01013444 -0.01070027 -0.01888977  0.05208318]], shape=(2, 4), dtype=float32)

Note 权重$\boxed{w}$和$\boxed{h}$被设置为层的属性后，层将自动记录跟踪。方法如下

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

Note 更快捷的方法来添加权重到一个层：$\boxed{add\_weight\left(\right)}$方法

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__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 tf.matmul(inputs, self.w) + self.b


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

打印结果

tf.Tensor(
[[-0.01331179 -0.00605625 -0.01042787  0.17160884]
 [-0.01331179 -0.00605625 -0.01042787  0.17160884]], shape=(2, 4), dtype=float32)

3. 层可具有不可训练的权重

除了可训练的权重，您还可以向层添加不可训练的权重。当您训练该层时，在反向传播期间不应考虑此类权重。

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

class ComputeSum(keras.layers.Layer):
    def __init__(self, input_dim):
        super(ComputeSum, self).__init__()
        self.total = tf.Variable(initial_value=tf.zeros((input_dim,)), trainable=False)

    def call(self, inputs):
        self.total.assign_add(tf.reduce_sum(inputs, axis=0))
        return self.total


x = tf.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: []

4. 最佳实践：延时创建权重，延时至知道输入的形状

我们Linear上面的层采用了一个$\boxed{input\_dim}$参数，用于计算权重$\boxed{w}$和$\boxed{bin}$的形状$\boxed{\_\_init\_\_\left(\right)}$：

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__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 tf.matmul(inputs, self.w) + self.b

在很多情景下，您可能无法提前获知输入的形状；或者层实例化后，您已经知道输入的形状，就希望延迟创建权重。

因此在Keras API中，推荐在KaTeX parse error: Expected '}', got '_' at position 46: …ld(self, inputs_̲shape)}方法中创建您的层的权重，如下所示：

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__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 tf.matmul(inputs, self.w) + self.b

该类的$\boxed{\_\_call\_\_\left(\right)}$方法将在第一次调用时自动运行$\boxed{build}$函数。现在有了一个可延迟创建的层，使用起来也更容易。

# 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)

5. 层可递归组合

如果内部层实例作为外部层的属性，那么外部层将跟踪内部层的权重参数。

推荐在$\boxed{\_\_init\_\_\left(\right)}$方法中创建子层实例（因为子层有构建（$\boxed{build}$方法），在外层进行构建时，子层也进行构建）

个人理解：

1. ++推荐操作主要是方便了延迟构建，运行程序自行推导参数，类似CPP STL的思想。++
2. ++层递归组合后的类，一般起名为xxxBlock,意思为几个层组成一个块状结构。++

# Let's assume we are reusing the Linear class
# with a `build` method that we defined above.


class MLPBlock(keras.layers.Layer):
    def __init__(self):
        super(MLPBlock, self).__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 = tf.nn.relu(x)
        x = self.linear_2(x)
        x = tf.nn.relu(x)
        return self.linear_3(x)


mlp = MLPBlock()
y = mlp(tf.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

6. $\boxed{add\_loss\left(\right)}$ 方法

在编写层的$\boxed{call\left(\right)}$时，您可以创建稍后用于训练循环的损失张量。这可以通过调用$\boxed{self.add\_loss\left(value\right)}$实现：

个人理解：

1. 在构建VAE，GAN等复杂模型时，将非常方便

# 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

这些损失（包括任何内层造成的损失）可以通过 layer.losses. 此属性在每个$\boxed{\_\_call\_\_\left(\right)}$顶层层的开始时重置，因此它layer.losses始终包含在上次前向传递期间创建的损失值。

class OuterLayer(keras.layers.Layer):
    def __init__(self):
        super(OuterLayer, self).__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(tf.zeros(1, 1))
assert len(layer.losses) == 1  # We created one loss value

# `layer.losses` gets reset at the start of each __call__
_ = layer(tf.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(OuterLayerWithKernelRegularizer, self).__init__()
        self.dense = keras.layers.Dense(
            32, kernel_regularizer=tf.keras.regularizers.l2(1e-3)
        )

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


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

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

打印输出

[<tf.Tensor: shape=(), dtype=float32, numpy=0.0018842274>]

在编写训练循环时要考虑这些损失，如下所示：

# Instantiate an optimizer.
optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Iterate over the batches of a dataset.
for x_batch_train, y_batch_train in train_dataset:
  with tf.GradientTape() as tape:
    logits = layer(x_batch_train)  # Logits for this minibatch
    # Loss value for this minibatch
    loss_value = loss_fn(y_batch_train, logits)
    # Add extra losses created during this forward pass:
    loss_value += sum(model.losses)

  grads = tape.gradient(loss_value, model.trainable_weights)
  optimizer.apply_gradients(zip(grads, model.trainable_weights))

有关编写训练循环的详细指南，请参阅从头开始编写训练循环的 指南。

这些损失也可以无缝地与$\boxed{fit\left(\right)}$（它们会自动求和并添加到主要损失中，如果有的话）：

import numpy as np

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 1ms/step - loss: 0.1555
1/1 [==============================] - 0s 927us/step - loss: 0.0336

<tensorflow.python.keras.callbacks.History at 0x145bca6d0>

7. $\boxed{add\_metric\left(\right)}$ 方法

类似于$\boxed{add\_loss\left(\right)}$方法，层也有一个$\boxed{add\_metric\left(\right)}$方法，用来跟踪训练期间的质量的移动均值。metric是用来评价模型性能的函数

例如：LogisticEndpoint层，它将预测和目标作为输入，通过$\boxed{add\_loss\left(\right)}$方法来计算损失，通过$\boxed{add\_metric\left(\right)}$来计算准确度标量。

class LogisticEndpoint(keras.layers.Layer):
    def __init__(self, name=None):
        super(LogisticEndpoint, self).__init__(name=name)
        self.loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)
        self.accuracy_fn = keras.metrics.BinaryAccuracy()

    def call(self, targets, logits, sample_weights=None):
        # Compute the training-time loss value and add it
        # to the layer using `self.add_loss()`.
        loss = self.loss_fn(targets, logits, sample_weights)
        self.add_loss(loss)

        # Log accuracy as a metric and add it
        # to the layer using `self.add_metric()`.
        acc = self.accuracy_fn(targets, logits, sample_weights)
        self.add_metric(acc, name="accuracy")

        # Return the inference-time prediction tensor (for `.predict()`).
        return tf.nn.softmax(logits)

以读取layer.metrics方式来跟踪指标。

layer = LogisticEndpoint()

targets = tf.ones((2, 2))
logits = tf.ones((2, 2))
y = layer(targets, logits)

print("layer.metrics:", layer.metrics)
print("current accuracy value:", float(layer.metrics[0].result()))

打印输出

layer.metrics: [<tensorflow.python.keras.metrics.BinaryAccuracy object at 0x145bccdd0>]
current accuracy value: 1.0

类似$\boxed{add\_loss\left(\right)}$，$\boxed{fit\left(\right)}$方法中将跟踪这些指标

inputs = keras.Input(shape=(3,), name="inputs")
targets = keras.Input(shape=(10,), name="targets")
logits = keras.layers.Dense(10)(inputs)
predictions = LogisticEndpoint(name="predictions")(logits, targets)

model = keras.Model(inputs=[inputs, targets], outputs=predictions)
model.compile(optimizer="adam")

data = {
    "inputs": np.random.random((3, 3)),
    "targets": np.random.random((3, 10)),
}
model.fit(data)

打印输出

1/1 [==============================] - 0s 999us/step - loss: 1.0366 - binary_accuracy: 0.0000e+00

<tensorflow.python.keras.callbacks.History at 0x1452c7650>

8. 可选，启用层的序列化

如果您需要把自定义层进行序列化，则需要实现$\boxed{get\_config\left(\right)}$方法

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__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 tf.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}

Note 基类的$\boxed{\_\_init\_\_\left(\right)}$采用的一些关键参数，如$\boxed{name}$参数和$\boxed{dtype}$参数。最佳实践是传递参数给父类并且将它们包含在层的配置中：

class Linear(keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super(Linear, self).__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 tf.matmul(inputs, self.w) + self.b

    def get_config(self):
        config = super(Linear, self).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_8', 'trainable': True, 'dtype': 'float32', 'units': 64}

如果在反序列化时您需要更大自由，则可以重写$\boxed{from\_config\left(\right)}$方法。下面时$\boxed{from\_config\left(\right)}$的基本实现

def from_config(cls, config):
  return cls(**config)

要了解有关序列化和保存的更多信息，请参阅保存和序列化模型指南。

9. $\boxed{call\left(\right)}$方法中的特权参数$\boxed{training}$

某些层，尤其是BatchNormalization层和Dropout层，在训练和推理过程中具有不同的行为。对于此类层，标准做法是在$\boxed{call\left(\right)}$方法中公开$\boxed{training}$（布尔）参数。

通过在在$\boxed{call\left(\right)}$方法中公开$\boxed{training}$参数，您可使能内置的训练和评估循环以实现在训练和推理中正确使用该层。

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

10. $\boxed{call\left(\right)}$方法中的特权参数$\boxed{mask}$

$\boxed{call\left(\right)}$方法中另一个特权参数是$\boxed{mask}$。

在所有Keras RNN层中都有该参数。一个$\boxed{mask}$是一个布尔张量（一个布尔值对应输入中的一个时间步长），用于处理时间序列数据时跳过某些输入时间步长。

当前一层生成掩码时，Keras会自动将正确的$\boxed{mask}$参数传递给$\boxed{\_\_call\_\_\left(\right)}$支持它的层。配置了$\boxed{mask\_zero=True\left(\right)}$的Embedding层和Masking层都是遮罩产生层。

进一步理解遮罩和如何编写遮罩使能层，请参考指南理解填充和遮罩

11. 模型类

通常，您使用层类来定义内部计算块，使用模型类来定义外部将要训练的对象。

例如，在ResNet50模型中，您有多个ResNet块继承自Layer然而一个Model就包含了整个ResNet50网络。

该模型类的多个API和Layer类中相似，以下是差别部分：

• 它公开了内置的训练、评估和预测循环 ( $\boxed{model.fit\left(\right)}$, $\boxed{model.evaluate\left(\right)}$, $\boxed{model.predict\left(\right)}$)。
• 它通过$\boxed{model.layers}$属性公开其内层的列表。
• 它公开了保存和序列化 API ( $\boxed{save\left(\right)}$, $\boxed{save\_weights\left(\right)}$,…)

实际上，Layer类对应我们在文献中所说的“层”（卷积层，循环层）或者“块”（残差块，初始块）；与此同时，Model类对应我们在文献中所说的“模型”（深度学习模型）或“网络”（深度神经网络）

如果你想知道，“我应该使用Layer类还是Model类？”，问问自己：我需要调用$\boxed{fit\left(\right)}$吗？我需要打调用$\boxed{save\left(\right)}$吗？如果是这样，那就使用Model。如果不是（或者因为您的类只是更大系统中的一个块，或者因为您自己编写培训和保存代码），请使用Layer.

12. 归纳整理: 一个完整的例子

以下是您目前为止所学到的：

• 一个层封装一个状态(created in $\boxed{\_\_init\_\_\left(\right)}$ or $\boxed{build\left(\right)}$)和一些计算(defined in $\boxed{call\left(\right)}$)
• 层可以递归嵌套来创建新的、更大的计算块
• 层可以通过$\boxed{add\_loss\left(\right)}$创建和跟踪损失（通常是正则化损失），同理可以通过$\boxed{add\_metric\left(\right)}$创建和跟踪准确度
• 外层容器，你要训练的东西，是一个模型。模型就像一个层，但增加了训练和序列化等功能

归纳整理到一个完整的例子：我们实现一个变分自动编码器(VAE)，然后在MNIST数字集上进行训练。

VariationalAutoEncoder继承自Model，构建了Sampling，Encoder，Decoder层，具有正则化损失（KL divergence）

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,))

    # Iterate over the batches of the dataset.
    for step, x_batch_train in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            reconstructed = vae(x_batch_train)
            # Compute reconstruction loss
            loss = mse_loss_fn(x_batch_train, reconstructed)
            loss += sum(vae.losses)  # Add KLD regularization loss

        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.3577
step 100: mean loss = 0.1258
step 200: mean loss = 0.0994
step 300: mean loss = 0.0893
step 400: mean loss = 0.0843
step 500: mean loss = 0.0809
step 600: mean loss = 0.0788
step 700: mean loss = 0.0772
step 800: mean loss = 0.0760
step 900: mean loss = 0.0750
Start of epoch 1
step 0: mean loss = 0.0747
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

请注意，由于 VAE 是子类化Model，它具有内置的训练循环。所以你也可以这样训练它：

vae = VariationalAutoEncoder(784, 64, 32)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())
vae.fit(x_train, x_train, epochs=2, batch_size=64)

打印输出

Epoch 1/2
938/938 [==============================] - 1s 1ms/step - loss: 0.0745
Epoch 2/2
938/938 [==============================] - 1s 1ms/step - loss: 0.0676

<tensorflow.python.keras.callbacks.History at 0x15f10e150>

13. 对比面向对象开发: the Functional API

上述是一个面向对象开发形式，你也可以采用函数式API来完成开发。

则例子程序可以如下所写：

original_dim = 784
intermediate_dim = 64
latent_dim = 32

# Define encoder model.
original_inputs = tf.keras.Input(shape=(original_dim,), name="encoder_input")
x = layers.Dense(intermediate_dim, activation="relu")(original_inputs)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
z = Sampling()((z_mean, z_log_var))
encoder = tf.keras.Model(inputs=original_inputs, outputs=z, name="encoder")

# Define decoder model.
latent_inputs = tf.keras.Input(shape=(latent_dim,), name="z_sampling")
x = layers.Dense(intermediate_dim, activation="relu")(latent_inputs)
outputs = layers.Dense(original_dim, activation="sigmoid")(x)
decoder = tf.keras.Model(inputs=latent_inputs, outputs=outputs, name="decoder")

# Define VAE model.
outputs = decoder(z)
vae = tf.keras.Model(inputs=original_inputs, outputs=outputs, name="vae")

# 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)
vae.add_loss(kl_loss)

# Train.
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())
vae.fit(x_train, x_train, epochs=3, batch_size=64)

打印输出

Epoch 1/3
938/938 [==============================] - 1s 1ms/step - loss: 0.0747
Epoch 2/3
938/938 [==============================] - 1s 1ms/step - loss: 0.0676
Epoch 3/3
938/938 [==============================] - 1s 1ms/step - loss: 0.0676

<tensorflow.python.keras.callbacks.History at 0x15f3240d0>