Keras 3中移除add_loss()后的变分自编码器损失处理方案求助

问题背景 #

正在学习一门使用Keras 2构建变分自编码器（VAE）的课程，尝试将代码迁移至Keras 3时，多数问题已解决，但在损失处理环节受阻。

原迁移后的损失计算代码（已改用ops） #

reconstruction_loss = ops.binary_crossentropy(inputs, outputs)
reconstruction_loss *= 784
kl_loss = 0.5 * (ops.exp(z_log_var) - (1 + z_log_var) + ops.square(z_mean))
kl_loss = ops.sum(kl_loss, axis=-1)

total_vae_loss = ops.mean(reconstruction_loss + kl_loss)

原训练代码 #

vae_model.add_loss(total_vae_loss)
vae_model.compile(optimizer='rmsprop', metrics=['accuracy'])

vae_model.fit(x_train_flat, epochs=epochs, batch_size=batch_size)

遇到的错误 #

1. 移除add_loss行后报错：

ValueError: No loss to compute. Provide a `loss` argument in `compile()`.

2. 将损失逻辑封装为函数传入compile()的loss参数后，训练时报错：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[52], line 1
----> 1 vae_model.fit(x_train_flat,
      2         epochs=epochs,
      3         batch_size=batch_size)

File ~/code/python/mls/mlsenv/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:122, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    119     filtered_tb = _process_traceback_frames(e.__traceback__)
    120     # To get the full stack trace, call:
    121     # `keras.config.disable_traceback_filtering()`
--> 122     raise e.with_traceback(filtered_tb) from None
    123 finally:
    124     del filtered_tb

File ~/code/python/mls/mlsenv/lib/python3.10/site-packages/optree/ops.py:594, in tree_map(func, tree, is_leaf, none_is_leaf, namespace, *rests)
    592 leaves, treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace)
    593 flat_args = [leaves] + [treespec.flatten_up_to(r) for r in rests]
--> 594 return treespec.unflatten(map(func, *flat_args))

ValueError: None values not supported.

解决办法 #

Keras 3中确实移除了符号式的Layer.add_loss()，只能在模型或层的call()方法内部调用add_loss()，这不仅适用于自定义层，也适用于自定义模型。正确的VAE实现方式是把损失计算逻辑嵌入自定义VAE模型的call方法中：

1. 自定义VAE模型类 #

import keras
from keras import ops

class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def call(self, inputs):
        # 编码器输出z_mean和z_log_var
        z_mean, z_log_var = self.encoder(inputs)
        # 重参数化采样
        z = z_mean + ops.exp(0.5 * z_log_var) * keras.random.normal(shape=ops.shape(z_mean))
        # 解码器输出重构结果
        outputs = self.decoder(z)

        # 计算损失
        reconstruction_loss = ops.binary_crossentropy(inputs, outputs)
        reconstruction_loss *= 784  # 匹配MNIST输入维度
        kl_loss = 0.5 * (ops.exp(z_log_var) - (1 + z_log_var) + ops.square(z_mean))
        kl_loss = ops.sum(kl_loss, axis=-1)
        total_loss = ops.mean(reconstruction_loss + kl_loss)

        # 在call方法内添加损失
        self.add_loss(total_loss)
        # 返回重构结果（可选，用于后续生成或评估）
        return outputs

2. 编译与训练 #

# 假设encoder和decoder已提前定义完成
vae_model = VAE(encoder, decoder)
# 编译时无需传入loss参数，损失已在call中定义
vae_model.compile(optimizer='rmsprop')
# 自监督训练仅需传入输入数据
vae_model.fit(x_train_flat, epochs=epochs, batch_size=batch_size)

错误原因说明 #

之前尝试将损失函数传入compile()的方式不可行，因为VAE的损失依赖编码器输出的z_mean和z_log_var等中间变量，普通损失函数无法直接获取这些内部张量，导致计算过程中出现None值触发报错。而在call方法内计算损失可以直接访问模型内部的中间张量，完全符合Keras 3的设计要求。

内容的提问来源于stack exchange，提问作者Weaver