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.

Criando novas camadas e modelos via subclasses

View on TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Configuração

import tensorflow as tf
from tensorflow import keras

A classe Layer: a combinação de estado (pesos) e alguma computação

Uma das abstrações centrais do Keras é a classe Layer (camada). Uma camada encapsula um estado (os "pesos" da camada) e uma transformação de entradas para saídas (uma "chamada" ou "call", o passo seguinte da camada).

Aqui está uma camada densamente conectada. Ela tem um estado: as variáveis w e 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

Você usa uma camada ao chamá-la em algumas entradas de tensor, de forma semelhante como se chama uma função Python.

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

Observe que os pesos w e b são rastreados automaticamente pela camada ao serem definidos como atributos da camada:

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

Observe que você também tem acesso a um atalho mais rápido para adicionar peso a uma camada: o 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)

As camadas podem ter pesos não treináveis

Além dos pesos treináveis, você também pode adicionar pesos não treináveis ​​a uma camada. Esses pesos não devem ser levados em consideração durante a retropropagação, quando você está treinando a camada.

Veja como adicionar e usar um peso não treinável:

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

Ele faz parte de layer.weights, mas é categorizado como um peso não treinável:

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)

Melhor prática: adiar a criação do peso até que o formato das entradas seja conhecido

Nossa camada Linear acima recebeu um argumento input_dim que foi usado para calcular o formato dos pesos w e b em __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

Em muitos casos, você pode não saber com antecedência o tamanho de suas entradas e pode querer criar pesos de forma postergada (lazy) quando esse valor se tornar conhecido, algum tempo depois de instanciar a camada.

Na API Keras, recomendamos criar pesos de camada no método build(self, inputs_shape) de sua camada. Da forma mostrada a seguir:

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

O método __call__() de sua camada será executado automaticamente na primeira vez que for chamado. Agora você tem uma camada lazy e, portanto, mais fácil de usar:

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

A implementação separada de build(), como mostrado acima, separa bem a criação de pesos para uso único, do uso de pesos em cada chamada. No entanto, para algumas camadas personalizadas avançadas, pode se tornar impraticável separar a criação e a computação do estado. Os implementadores de camada podem adiar a criação de pesos para a primeira __call__(), mas precisam tomar cuidado para que chamadas posteriores usem os mesmos pesos. Além disso, como __call__() provavelmente será executada pela primeira vez dentro de uma tf.function, qualquer criação de variável que ocorrer em __call__() deve ser empacotada num tf.init_scope.

As camadas podem ser compostas recursivamente

Se você atribuir uma instância de Layer como atributo de outra Layer, a camada externa começará a rastrear os pesos criados pela camada interna.

Recomendamos criar tais subcamadas no método __init__() e deixar para que a primeira chamada __call__() cuide da construção de seus pesos.

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

O método add_loss()

Ao escrever o método call() de uma camada, você pode criar tensores de perda que deseja usar mais tarde, quando for escrever seu loop de treinamento. Isto é possível chamando 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

Essas perdas (incluindo aquelas criadas por qualquer camada interna) podem ser recuperadas via layer.losses. Esta propriedade é reiniciada no início de cada __call__() para a camada de nível superior, de modo que layer.losses sempre irá conter os valores de perda criados durante o último passo adiante.

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

Além disso, a propriedade loss também contém perdas de regularização criadas para os pesos de qualquer camada 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)`,
# created by the `kernel_regularizer` above.
print(layer.losses)

Essas perdas devem ser levadas em consideração ao escrever loops de treinamento, como este:

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

Se quiser um guia detalhado sobre como escrever loops de treinamento, consulte o guia para escrever um loop de treinamento do zero .

Essas perdas também funcionam perfeitamente com fit() (elas são automaticamente somadas e adicionadas à perda principal, se houver):

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

O método add_metric()

Assim como add_loss(), as camadas também possuem um método add_metric() para rastrear a média móvel de uma quantidade durante o treinamento.

Considere a seguinte camada: uma camada de "endpoint logístico". Ele recebe como entradas previsões e alvos, calcula uma perda que rastreia via add_loss(), e calcula um escalar de exatidão, que ele rastreia via 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):
        # 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)

As métricas rastreadas dessa maneira podem ser acessadas via 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()))

Assim como para add_loss(), essas métricas são rastreadas por 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)

Você pode, opcionalmente, ativar a serialização em suas camadas

Se você precisa que suas camadas personalizadas sejam serializáveis ​​como parte de um modelo Functional, você pode, opcionalmente, implementar um 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)

Observe que o método __init__() da classe Layer base recebe alguns argumentos de palavra-chave, em particular um name e um dtype. É boa prática passar esses argumentos para a classe pai em __init__() e incluí-los na configuração da camada:

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)

Se você precisar de mais flexibilidade ao desserializar a camada de sua configuração, também poderá sobrepor o método de classe from_config(). Esta é a implementação básica de from_config():

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

Para saber mais sobre serialização e salvamento, consulte o guia de salvamento e serialização de modelos.

Argumento training privilegiado no método call()

Algumas camadas, em particular as camadas BatchNormalization e Dropout, têm comportamentos diferentes durante o treinamento e a inferência. Para tais camadas, é prática padrão expor um argumento training (booleano) no método call().

Ao expor esse argumento em call(), você permite que os loops integrados de treinamento e avaliação (por exemplo, fit()) usem corretamente a camada em treinamento e inferência.

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 no método call()

O outro argumento privilegiado suportado por call() é o argumento mask.

Você vai encontrá-lo em todas as camadas Keras RNN. Uma máscara é um tensor booleano (um valor booleano por timestep na entrada) usado para ignorar determinados timestep de entrada ao processar dados de série temporal.

O Keras passará automaticamente o argumento mask correto para __call__() para as camadas que o suportarem, quando uma máscara for gerada por uma camada anterior. As camadas geradoras de máscara são a camada Embedding configurada com mask_zero=True e a camada Masking.

Para saber mais sobre mascaramento e como escrever camadas habilitadas para mascaramento, consulte o guia "entendendo preenchimento e mascaramento".

A classe Model

Em geral, você usará a classe Layer para definir blocos de computação internos e usará a classe Model para definir o modelo externo: o objeto que você vai treinar.

Por exemplo, num modelo ResNet50, você teria vários blocos ResNet criando subclasses de Layer, e um único Model abrangendo toda a rede ResNet50.

A classe Model tem a mesma API que Layer, com as seguintes diferenças:

• Ela expõe loops integrados de treinamento, avaliação e previsão (model.fit(), model.evaluate(), model.predict()).

• Ela expõe a lista de suas camadas internas, através da propriedade model.layers.

• Ela expõe APIs de salvamento e serialização (save(), save_weights()...)

Efetivamente, a classe Layer corresponde ao que nos referimos na literatura como uma "camada" (como em "camada de convolução" ou "camada recorrente") ou como um "bloco" (como em "bloco ResNet" ou "bloco Inception").

Já a classe Model corresponde ao que é referido na literatura como "modelo" (como em "modelo de aprendizado profundo") ou como "rede" (como em "rede neural profunda").

Então, se você estiver se perguntando "devo usar a classe Layer ou a classe Model?", pergunte-se: precisarei chamar fit() nela? Precisarei chamar save() nela? Se assim for, vá com Model. Caso contrário (porque sua classe é apenas um bloco em um sistema maior ou porque você mesmo está escrevendo treinamento e salvando código), use Layer.

Por exemplo, poderíamos pegar nosso exemplo de mini-resnet acima e usá-lo para construir um Model que poderíamos treinar com fit() e que poderíamos salvar com 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)

Juntando tudo: um exemplo completo

Aqui está o que você aprendeu até agora:

• Uma Layer encapsula um estado (criado em __init__() ou build()) e alguma computação (definida em call()).

• As camadas podem ser aninhadas recursivamente para criar novos blocos de computação maiores.

• As camadas podem criar e rastrear perdas (normalmente perdas de regularização), bem como métricas, via add_loss() e add_metric()

• O container externo, a coisa que você deseja treinar, é um Model. Um Model é como uma Layer, mas com utilitários adicionais de treinamento e serialização.

Vamos juntar todas essas coisas num exemplo completo: vamos implementar um Variational AutoEncoder (VAE). Vamos treiná-lo em dígitos MNIST.

Nosso VAE será uma subclasse de Model, construída como uma composição aninhada de camadas que são subclasses de Layer. Ela contará com uma perda de regularização (divergência 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

Vamos escrever um loop de treinamento simples no 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()))

Observe que, como o VAE é uma subclasse de Model, ele apresenta loops de treinamento integrados. Então você também poderia tê-lo treinado assim:

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)

Além do desenvolvimento orientado a objetos: a API Functional

Você achou que esse exemplo foi demasiado orientado a objetos para você? Você também pode criar modelos usando a API Functional. É importante ressaltar que escolher um estilo ou outro não impede que você aproveite os componentes escritos no outro estilo: você sempre pode misturar e combinar.

Por exemplo, o exemplo abaixo usando a API Functional reutiliza a mesma camada Sampling que definimos no exemplo acima:

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 mais informações, não deixe de ler o guia da API Functional.