使用ImageDataGenerator的Keras变分自编码器报Graph execution error

DevOps 小满

2026-7-13

修复变分自编码器(VAE)训练时的形状不匹配错误

问题背景

构建图像生成VAE时，调用vae.fit(image_generator)触发InvalidArgumentError，提示required broadcastable shapes，核心原因是损失计算中重建损失与KL损失的形状不匹配，无法直接相加。

原模型代码

def sampling(args):
    mu, log_var = args
    batch_size = K.shape(mu)[0]
    dim = K.int_shape(mu)[1]
    epsilon = K.random_normal(shape=(batch_size, dim), mean=0., stddev=1.0)
    return K.exp(0.5 * log_var) * epsilon + mu

def build_vae(input_count,neuron_count_per_hidden_layer,encoded_dim,hidden_activation,output_activation):
    #Encoder
    encoder_input = layers.Input(shape=(80,60,1), name='encoder_input')
    prev_layer=encoder_input

    prev_layer = layers.Flatten()(prev_layer)
    prev_layer = layers.Dense(256,activation=hidden_activation)(prev_layer)
    for neuron_count in neuron_count_per_hidden_layer:
      hidden_layer=layers.Dense(neuron_count,activation=hidden_activation)(prev_layer)
      prev_layer=hidden_layer
    
    mu = layers.Dense(encoded_dim, name='mu')(prev_layer)
    log_var = layers.Dense(encoded_dim, name='log_var')(prev_layer)
    
    encoder = keras.Model(encoder_input, [mu, log_var], name='encoder')

    #Decoder
    decoder_input = layers.Input(shape=(encoded_dim,), name='decoder_input')

    prev_layer=decoder_input
    for neuron_count in reversed(neuron_count_per_hidden_layer):
      hidden_layer=layers.Dense(neuron_count,activation=hidden_activation)(prev_layer)
      prev_layer=hidden_layer

    prev_layer = layers.Dense(4800,activation=hidden_activation)(prev_layer)
    decoder_output_layer=layers.Reshape((80,60,1), name='decoder_output')(prev_layer)

    decoder = keras.Model(decoder_input, decoder_output_layer, name='decoder')

    #Sampling layer
    s = layers.Lambda(sampling, output_shape=(encoded_dim,), name='s')([mu, log_var])

    #VAE
    vae=keras.Model(encoder.input, decoder(s),name='vae')
    
    return vae,encoder,decoder

def vae_loss(vae_input,vae_ouput,mu,log_var,kl_coefficient, input_count):
  #Reconstruction loss
  reconstruction_loss = keras.losses.mean_squared_error(vae_input,vae_ouput) * input_count

  #Regularization loss
  kl_loss = 0.5 * K.sum(K.square(mu) + K.exp(log_var) - log_var - 1, axis = -1)

  #Combined loss
  return reconstruction_loss + kl_coefficient*kl_loss

def getVAE(shape):
   vae, vae_encoder, vae_decoder=build_vae(shape, [128,64], 16,'sigmoid','sigmoid')
   vae.summary()
   return vae, vae_encoder, vae_decoder

原训练代码

datagen = ImageDataGenerator(
    rescale=1.0 / 255.0,  # Scale pixel values between 0 and 1
) 

data_dir = paths.BW_IMG_FOLDER
batch_size = 128

image_generator = datagen.flow_from_directory(
    data_dir,
    color_mode="grayscale",
    target_size=(80, 60),  # Set your desired image dimensions
    batch_size=batch_size,
    class_mode='input',  # No class labels, unsupervised learning
    shuffle=True  # Shuffle the data
)

importlib.reload(v1) # the module "v1" was described in previous code sample
vae, vae_encoder, vae_decoder = v1.getVAE(4800)

kl_coefficient=1

#Information needed to compute the loss function
vae_input=vae.input
vae_output=vae.output
mu=vae.get_layer('mu').output
log_var=vae.get_layer('log_var').output

vae.add_loss(v1.vae_loss(vae_input,vae_output,mu,log_var,kl_coefficient,4800))
vae.compile(optimizer='adam')

epoch_count = 100
patience=5

early_stop = keras.callbacks.EarlyStopping(monitor='val_loss', patience=patience, restore_best_weights=True)

history = vae.fit(image_generator,epochs=epoch_count,callbacks=[early_stop])

错误信息

InvalidArgumentError                      Traceback (most recent call last)
c:\Users\dengr\Desktop\Fashion-Product-Images-Generation\src\main.py in line 8
      133 patience=5
      135 early_stop = keras.callbacks.EarlyStopping(monitor='val_loss', patience=patience, restore_best_weights=True)
----> 137 history = vae.fit(image_generator,epochs=epoch_count,callbacks=[early_stop])

File c:\Users\dengr\anaconda3\envs\DL-10\lib\site-packages\keras\utils\traceback_utils.py:70, in filter_traceback..error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File c:\Users\dengr\anaconda3\envs\DL-10\lib\site-packages\tensorflow\python\eager\execute.py:54, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52 try:
     53   ctx.ensure_initialized()
---> 54   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                       inputs, attrs, num_outputs)
     56 except core._NotOkStatusException as e:
     57   if name is not None:

InvalidArgumentError: Graph execution error:
Detected at node 'vae/tf.__operators__.add_1/AddV2' defined at (most recent call last):
    File "c:\Users\user\anaconda3\envs\DL-10\lib\runpy.py", line 196, in _run_module_as_main
      return _run_code(code, main_globals, None,
    File "c:\Users\user\anaconda3\envs\DL-10\lib\runpy.py", line 86, in _run_code
      exec(code, run_globals)
...
...
...
   File "c:\Users\user\anaconda3\envs\DL-10\lib\site-packages\keras\utils\traceback_utils.py", line 96, in error_handler
      return fn(*args, **kwargs)
    File "c:\Users\user\anaconda3\envs\DL-10\lib\site-packages\keras\layers\core\tf_op_layer.py", line 242, in _call_wrapper
      return self._call_wrapper(*args, **kwargs)
    File "c:\Users\user\anaconda3\envs\DL-10\lib\site-packages\keras\layers\core\tf_op_layer.py", line 279, in _call_wrapper
      result = self.function(*args, **kwargs)
Node: 'vae/tf.__operators__.add_1/AddV2'
required broadcastable shapes
     [[{{node vae/tf.__operators__.add_1/AddV2}}]] [Op:__inference_train_function_1412]

问题分析

损失形状不匹配：原vae_loss中，mean_squared_error返回的是形状为(batch_size,80,60,1)的逐像素损失，乘以input_count后仍保留空间维度；而kl_loss是形状为(batch_size,)的逐样本损失，两者无法广播相加，触发错误。
冗余参数：build_vae的input_count参数未实际使用，传入的4800属于冗余信息，易造成混淆。

修复方案

1. 修改损失函数，统一损失维度

将重建损失压缩到逐样本维度，确保与KL损失形状一致：

def vae_loss(vae_input, vae_output, mu, log_var, kl_coefficient):
    # 对每个样本的所有像素MSE求和，得到(batch_size,)的张量
    reconstruction_loss = K.sum(keras.losses.mean_squared_error(vae_input, vae_output), axis=(1,2,3))
    # KL损失已为(batch_size,)维度
    kl_loss = 0.5 * K.sum(K.square(mu) + K.exp(log_var) - log_var - 1, axis=-1)
    # 合并损失
    return reconstruction_loss + kl_coefficient * kl_loss

2. 移除冗余参数，简化模型构建函数

def build_vae(neuron_count_per_hidden_layer, encoded_dim, hidden_activation, output_activation):
    #Encoder
    encoder_input = layers.Input(shape=(80,60,1), name='encoder_input')
    prev_layer=encoder_input

    prev_layer = layers.Flatten()(prev_layer)
    prev_layer = layers.Dense(256,activation=hidden_activation)(prev_layer)
    for neuron_count in neuron_count_per_hidden_layer:
        hidden_layer=layers.Dense(neuron_count,activation=hidden_activation)(prev_layer)
        prev_layer=hidden_layer
    
    mu = layers.Dense(encoded_dim, name='mu')(prev_layer)
    log_var = layers.Dense(encoded_dim, name='log_var')(prev_layer)
    
    encoder = keras.Model(encoder_input, [mu, log_var], name='encoder')

    #Decoder
    decoder_input = layers.Input(shape=(encoded_dim,), name='decoder_input')

    prev_layer=decoder_input
    for neuron_count in reversed(neuron_count_per_hidden_layer):
        hidden_layer=layers.Dense(neuron_count,activation=hidden_activation)(prev_layer)
        prev_layer=hidden_layer

    prev_layer = layers.Dense(4800,activation=hidden_activation)(prev_layer)
    decoder_output_layer=layers.Reshape((80,60,1), name='decoder_output')(prev_layer)

    decoder = keras.Model(decoder_input, decoder_output_layer, name='decoder')

    #Sampling layer
    s = layers.Lambda(sampling, output_shape=(encoded_dim,), name='s')([mu, log_var])

    #VAE
    vae=keras.Model(encoder.input, decoder(s),name='vae')
    
    return vae,encoder,decoder

def getVAE():
    vae, vae_encoder, vae_decoder=build_vae([128,64], 16,'sigmoid','sigmoid')
    vae.summary()
    return vae, vae_encoder, vae_decoder

3. 调整训练代码调用

vae, vae_encoder, vae_decoder = v1.getVAE()

kl_coefficient=1

#Information needed to compute the loss function
vae_input=vae.input
vae_output=vae.output
mu=vae.get_layer('mu').output
log_var=vae.get_layer('log_var').output

vae.add_loss(v1.vae_loss(vae_input,vae_output,mu,log_var,kl_coefficient))
vae.compile(optimizer='adam')

修复说明

通过K.sum对重建损失的空间维度求和，将其从(batch_size,80,60,1)转换为(batch_size,)，与KL损失形状匹配，解决广播错误。
移除冗余的input_count参数，避免参数混淆，模型输入形状已固定为(80,60,1)，总像素数可通过输入推导。

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

使用ImageDataGenerator的Keras变分自编码器报Graph execution error

问题背景 #

原模型代码 #

原训练代码 #

错误信息 #

问题分析 #

修复方案 #

1. 修改损失函数，统一损失维度 #

2. 移除冗余参数，简化模型构建函数 #

3. 调整训练代码调用 #

修复说明 #