TensorFlow构建CVAE：非Binary Crossentropy损失无法收敛求助

用Python+TensorFlow搭建卷积变分自编码器（CVAE），训练数据为自制的64×64像素图像。当前遇到以下问题：

• 仅采用Binary Crossentropy作为重建损失时，模型可正常收敛
• 切换为MSE、MAE、KLDiv等损失函数后，模型完全无法收敛，各轮次损失值固定不变，预测结果全为NaN
• 已尝试调整各类超参数，问题仍未解决

实现代码 #

import tensorflow as tf
from tensorflow.keras import layers, models, losses, optimizers, callbacks
import numpy as np

# Define input dimensions
input_shape = (64, 64, 3)  # Change to (64, 64, 1) if grayscale

# Encoder
def build_encoder(input_shape, latent_dim):
    encoder_inputs = layers.Input(shape=input_shape)
    x = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(encoder_inputs)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)
    x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)
    x = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)
    x = layers.Conv2D(256, (3, 3), activation='relu', padding='same')(x)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)
    x = layers.Flatten()(x)
    #x = layers.Dense(512, activation='relu')(x)
    #x = layers.Dense(128, activation='relu')(x)
    z_mean = layers.Dense(latent_dim, name='z_mean')(x)
    z_log_var = layers.Dense(latent_dim, name='z_log_var')(x)

    def sampling(args):
        z_mean, z_log_var = args
        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

    z = layers.Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])

    encoder = models.Model(encoder_inputs, [z_mean, z_log_var, z], name='encoder')
    return encoder

# Decoder
def build_decoder(latent_dim):
    latent_inputs = layers.Input(shape=(latent_dim,))
    #x = layers.Dense(512, activation='relu')(latent_inputs)
    #x = layers.Dense(256, activation='relu')(x)
    x = layers.Dense(4 * 4 * 256, activation='relu')(latent_inputs)
    x = layers.Reshape((4, 4, 256))(x)
    x = layers.UpSampling2D((2, 2))(x)
    x = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = layers.UpSampling2D((2, 2))(x)
    x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = layers.UpSampling2D((2, 2))(x)
    x = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(x)
    x = layers.UpSampling2D((2, 2))(x)
    decoder_outputs = layers.Conv2D(3, (3, 3), activation='sigmoid', padding='same')(x)  # 3 for color, change to 1 for grayscale

    decoder = models.Model(latent_inputs, decoder_outputs, name='decoder')
    return decoder

# Define the Variational Autoencoder (VAE) model
class VAE(models.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstructed = self.decoder(z)
        kl_loss = -0.5 * tf.reduce_mean(
            z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
        kl_loss *= 0.0 #Sæt vægten her
        self.add_loss(kl_loss)
        return reconstructed

# Parameters
latent_dim = 50  # Adjust as needed

# Build encoder and decoder
encoder = build_encoder(input_shape, latent_dim)
decoder = build_decoder(latent_dim)

# Build VAE
vae = VAE(encoder, decoder)
optimizer = optimizers.Adam(learning_rate=0.001)
#vae.compile(optimizer=optimizer, loss=losses.MeanSquaredError())
vae.compile(optimizer=optimizer, loss=losses.BinaryCrossentropy())

optimizer.SGD 
# Print model summaries
encoder.summary()
decoder.summary()

# Ensure data shapes are correct before training
print(f"Shape of X_train: {X_train.shape}")
print(f"Shape of X_val: {X_val.shape}")

early_stopping = callbacks.EarlyStopping(monitor='val_loss', patience=6, restore_best_weights=True)
reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=3, min_lr=1e-4)

vae.fit(X_train, X_train, epochs=30, batch_size=32, validation_data=(X_val, X_val), callbacks=[early_stopping, reduce_lr])
vae.summary()

核心原因 #

1. 输出激活与损失不匹配：解码器最后一层用sigmoid激活，输出被压缩到(0,1)区间，但输入图像若未归一化到相同区间，MSE/MAE计算的损失值会异常巨大，引发梯度爆炸或消失，最终出现NaN。而Binary Crossentropy天然适配0-1区间的输出，因此能正常工作。
2. KL损失被完全关闭：代码中kl_loss *= 0.0移除了KL正则约束，导致编码器失去对潜在空间的约束，容易引发潜在空间崩溃，在非交叉熵损失下更易出现数值不稳定。
3. 无效代码干扰：optimizer.SGD是无意义的代码，可能导致潜在的配置冲突。

具体修复步骤 #

• 归一化输入数据：
  将训练/验证集图像归一化到(0,1)区间，与解码器输出尺度对齐：

  X_train = X_train.astype('float32') / 255.0
  X_val = X_val.astype('float32') / 255.0
  

• 恢复KL损失权重：
  设置合理的KL损失权重（如0.1，可根据效果调整），避免潜在空间崩溃：

  kl_loss *= 0.1  # 替换原有的0.0
  

• 优化数值稳定性：
  在采样函数中限制z_log_var的范围，防止指数运算溢出：

  def sampling(args):
      z_mean, z_log_var = args
      batch = tf.shape(z_mean)[0]
      dim = tf.shape(z_mean)[1]
      # 限制z_log_var范围，避免exp溢出
      z_log_var = tf.clip_by_value(z_log_var, -10, 10)
      epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
      return z_mean + tf.exp(0.5 * z_log_var) * epsilon
  

• 清理无效代码：
  删除无意义的optimizer.SGD行，避免干扰优化器配置。
• 重新编译模型：
  以MSE为例，修改后重新编译：

  vae.compile(optimizer=optimizer, loss=losses.MeanSquaredError())
  

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