custom_layers_and_models.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Cómo crear nuevas capas y modelos mediante subclases

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver el código fuente en GitHub

Descargar bloc de notas

Preparación

import tensorflow as tf
from tensorflow import keras

La clase Layer: la combinación de estado (ponderaciones) y algún cálculo

Una de las abstracciones centrales en Keras es la clase Layer. Una capa encapsula tanto un estado (las "ponderaciones" de la capa) como una transformación de entradas a salidas (una "llamada", el paso siguiente de la capa).

A continuación se muestra una capa densamente conectada. Tiene un estado: las variables w y b.

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

Para utilizar una capa, es necesario llamarla a una entrada o entradas de tensor, como si se tratara de una función de Python.

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

Tenga en cuenta que los pesos w y b son rastreados automáticamente por la capa al establecerse como atributos de la capa:

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

Tenga en cuenta que también tiene acceso a un atajo más rápido para agregar ponderación a una capa: el método add_weight():

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)

Las capas pueden tener ponderaciones que no se pueden entrenar

Además de las ponderaciones que se pueden entrenar, también se pueden agregar ponderaciones que no se pueden entrenar en una capa. Se supone que estas ponderaciones no se tendrán en cuenta durante la retropropagación, cuando la capa esté en entrenamiento.

A continuación se explica cómo agregar y utilizar una ponderación que no se puede entrenar:

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

Es parte de layer.weights, pero se clasifica como una ponderación que no se puede entrenar:

print("weights:", len(my_sum.weights))
print("non-trainable weights:", len(my_sum.non_trainable_weights))

# No se incluye en las ponderaciones que se pueden entrenar:
print("trainable_weights:", my_sum.trainable_weights)

Práctica recomendada: aplazar la creación de las ponderaciones hasta que se conozca la forma de las entradas.

Nuestra capa Linear anterior tomó un argumento input_dim que se utilizó para calcular la forma de las ponderaciones w y b en __init__():

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

En muchos casos, es posible que no conozca de antemano el tamaño de sus entradas, y le gustaría crear ponderaciones libremente cuando se conozca ese valor, un poco después de establecer la instancia de la capa.

En la API de Keras, recomendamos crear ponderaciones de capa en el método build(self, inputs_shape) de su capa. De la siguiente manera:

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

El método __call__() de su capa ejecutará automáticamente la creación la primera vez que se llame. Ahora tiene una capa que es lenta y por lo tanto más fácil de usar:

# En la creación de instancias, no sabemos en qué entradas va a recibir esta llamada
linear_layer = Linear(32)

# Las ponderaciones de la capa se crean dinámicamente la primera vez que llama a la capa
y = linear_layer(x)

Implementar build() por separado como se muestra arriba separa bien la creación de ponderaciones solo una vez de su uso en cada llamada. Sin embargo, para algunas capas personalizadas avanzadas, puede resultar poco práctico separar la creación del estado y el cálculo. Los implementadores de capas pueden diferir en la creación de las ponderaciones a la primera __call__(), pero deben tener cuidado de que las llamadas posteriores utilicen las mismas ponderaciones. Además, dado que __call__() es probable que se ejecute por primera vez dentro de una tf.function, cualquier creación de una variable que tenga lugar en __call__() debería estar envuelta en untf.init_scope.

Las capas se pueden componer recursivamente

Si asigna una instancia de una capa como atributo de otra capa, la capa externa comenzará a seguir las ponderaciones creadas por la capa interna.

Le recomendamos crear dichas subcapas en el método __init__() y dejar que sea el primer __call__() el que active la creación de sus ponderaciones.

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

El método add_loss()

Cuando escriba el método call() de una capa, puede crear tensores de pérdida que querrá utilizar más tarde, cuando escriba su bucle de entrenamiento. Esto se puede hacer llamando a self.add_loss(value):

# 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

Estas pérdidas (incluyendo las creadas por cualquier capa interna) pueden recuperarse mediante layer.losses. Esta propiedad se restablece al inicio de cada __call__() a la capa de nivel superior, de modo que layer.losses siempre contiene los valores de pérdidas creados durante el último pase siguiente.

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` se reinicia al comenzar cada __call__
_ = layer(tf.zeros(1, 1))
assert len(layer.losses) == 1  # This is the loss created during the call above

Además, la propiedad loss también contiene pérdidas obtenidas por regularización para las ponderaciones de cualquier capa interna:

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)`,
# creado por el `kernel_regularizer` anterior.
print(layer.losses)

Estas pérdidas deben tenerse en cuenta al escribir bucles de entrenamiento, como éste:

# Instancia de un optimizador.
optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Itera sobre los lotes de un conjunto de datos.
for x_batch_train, y_batch_train in train_dataset:
  with tf.GradientTape() as tape:
    logits = layer(x_batch_train)  # Logits for this minibatch
    # El valor de la pérdida para este minilote
    loss_value = loss_fn(y_batch_train, logits)
    # Agrega las pérdidas adicionales creadas durante este pase siguiente:
    loss_value += sum(model.losses)

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

Para obtener una guía detallada sobre la escritura de bucles de entrenamiento, consulte la guía para escribir un bucle de entrenamiento desde cero.

Estas pérdidas también funcionan perfectamente con fit() (se suman automáticamente y se agregan a la pérdida principal, si la hay):

import numpy as np

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

# Si pasa una pérdida en `compilar`, la regularización
# se le agregan las pérdidas
model.compile(optimizer="adam", loss="mse")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))

# También es posible no pasar ninguna pérdida en `compile`,
# ya que el modelo ya tiene una pérdida que minimizar, mediante la función `add_loss`
# ¡llame durante el pase siguiente!
model.compile(optimizer="adam")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))

El método add_metric()

De forma similar a add_loss(), las capas también tienen un método add_metric() para seguir el promedio móvil de una cantidad durante el entrenamiento.

Consideremos la siguiente capa: una capa de "punto final logístico". Toma como entradas predicciones y objetivos, calcula una pérdida que rastrea mediante add_loss(), y calcula un escalar de precisión, que se rastrea mediante add_metric().

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):
        # Calcule el valor de la pérdida en el periodo de entrenamiento y súmelo
        # a la capa utilizando `self.add_loss()`.
        loss = self.loss_fn(targets, logits, sample_weights)
        self.add_loss(loss)

        # Registre la precisión como métrica y agreguela
        # a la capa mediante `self.add_metric()`.
        acc = self.accuracy_fn(targets, logits, sample_weights)
        self.add_metric(acc, name="accuracy")

        # Devuelve el tensor de predicción en tiempo de inferencia (a `.predict()`).
        return tf.nn.softmax(logits)

Las métricas rastreadas de este modo son accesibles mediante 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()))

Al igual que en add_loss(), estas métricas se rastrean mediante fit():

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)

Puede activar opcionalmente la serialización en sus capas

Si es necesario que sus capas personalizadas se puedan serializar como parte de un {Functional model, puede implementar opcionalmente un método get_config():

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)

Tenga en cuenta que el método __init__() de la clase base Layer toma algunos argumentos de la palabra clave, en particular un name y un dtype. Es una práctica recomendada pasar estos argumentos a la clase padre en __init__() e incluirlos en la configuración de la capa:

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)

Si necesita mayor flexibilidad para deserializar la capa desde su configuración, también puede sobreescribir el método de la clase from_config(). Esta es la implementación base de from_config():

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

Para obtener más información sobre la serialización y el guardado, consulte la guía completa para guardar y serializar modelos.

Argumento training privilegiado en el método call()

Algunas capas, en particular la capa BatchNormalization y la capa Dropout, tienen comportamientos diferentes durante el entrenamiento y la inferencia. Para estas capas, es una práctica habitual como introducir un argumento training (booleano) en el método call().

Al exponer este argumento en call(), permite que los bucles de entrenamiento y evaluación incorporados (por ejemplo fit()) utilicen correctamente la capa en el entrenamiento y la inferencia.

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

Argumento mask privilegiado en el método call()

El otro argumento privilegiado soportado por call() es el argumento mask.

Lo encontrará en todas las capas RNN de Keras. Una máscara es un tensor booleano (un valor booleano por paso de tiempo en la entrada) utilizado para omitir ciertos pasos de tiempo de entrada al procesar datos de series temporales.

Keras pasará automáticamente el argumento mask correcto a __call__() para las capas que lo admitan, cuando una capa anterior genere una máscara. Las capas que generan máscaras son la capa Embedding configurada con mask_zero=True, y la capa Masking.

Para obtener más información sobre el enmascaramiento y cómo escribir capas con enmascaramiento, consulte la guía "cómo comprender el relleno y el enmascaramiento".

La clase Model

En general, utilizará la clase Layer para definir los bloques del cálculo interno y utilizará la clase Model para definir el modelo externo, el objeto que entrenará.

Por ejemplo, en un modelo ResNet50, tendría varios bloques ResNet que subclasificarían Layer, y un único Model que abarcaría toda la red ResNet50.

La clase Model tiene la misma API que Layer, con las siguientes diferencias:

• Expone bucles de entrenamiento, evaluación y predicción incorporados (model.fit(), model.evaluate(), model.predict()).

• Expone la lista de sus capas internas, mediante la propiedad model.layers.

• Expone las API de guardado y serialización (save(), save_weights()...)

En efecto, la clase Layer corresponde a lo que en la literatura se denomina "capa" (como en "capa de convolución" o "capa recurrente") o "bloque" (como en "bloque ResNet" o "bloque Inception").

Por otra parte, la clase Model corresponde a lo que en la literatura se denomina "modelo" (como en "modelo de aprendizaje profundo") o "red" (como en "red neuronal profunda").

Así que si se pregunta "¿debería usar la clase Layer o la clase Model?", pregúntese: ¿necesitaré llamar a fit() en ella? ¿Tendré que llamar a save()? Si es así, utilice Model. Si no es así (ya sea porque su clase es sólo un bloque en un sistema más grande, o porque está escribiendo el código de entrenamiento y guardado usted mismo), utilice Layer.

Por ejemplo, podríamos tomar nuestro ejemplo de mini-resnet anterior, y utilizarlo para construir un Model que podríamos entrenar con fit(), y que podríamos guardar con save_weights():

class ResNet(tf.keras.Model):

    def __init__(self, num_classes=1000):
        super(ResNet, self).__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)

Todo en uno: un ejemplo de principio a fin

Esto es lo que ha aprendido hasta ahora:

• A Layer encapsulate a state (created in __init__() or build()) and some computation (defined in call()).

• Las capas pueden anidarse recursivamente para crear nuevos bloques de cálculo más grandes.

• Las capas pueden crear y rastrear pérdidas (normalmente pérdidas de regularización), así como métricas, mediante add_loss() y add_metric().

• El contenedor externo, lo que desea entrenar, es un Model. Un Model es igual que un Layer, pero con utilidades adicionales de entrenamiento y serialización.

Pongamos todas estas cosas juntas en un ejemplo de principio a fin: vamos a implementar un Autoencoder Variacional (VAE). Lo entrenaremos con dígitos MNIST.

Nuestra VAE será una subclase de Model, construida como una composición anidada de capas que subclasifican a Layer. Incluirá una pérdida de regularización (divergencia 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

Escriba un bucle de entrenamiento simple en 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()))

Tenga en cuenta que, dado que la VAE es una subclase de Model, incorpora bucles de entrenamiento. Así que también podría haberlo entrenado así:

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)

Más allá del desarrollo orientado a objetos: la API funcional

¿Le resultó este ejemplo demasiado orientado a objetos? También puede construir modelos utilizando la API Funcional. Es importante destacar que la elección de un estilo u otro no le impide aprovechar los componentes escritos en el otro estilo: siempre puede mezclar y hacer las combinaciones que desee.

Por ejemplo, el siguiente ejemplo de API funcional reutiliza la misma capa Sampling que definimos en el ejemplo anterior:

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)

Para obtener más información, asegúrese de leer la Guía funcional de la API.