使用model.train_on_batch训练GAN时64批次被拆分为2个32批次的原因

核心原因分析 #

你遇到的train_on_batch将64大小批次拆分为两个32大小批次的现象，主要有以下几种可能：

1. steps_per_execution 参数设置问题 #

TensorFlow的模型编译参数steps_per_execution控制每次调用模型时执行的梯度更新步数。如果该参数被设置为2（可能是环境默认值或之前的代码修改了全局配置），那么每次train_on_batch调用会自动将输入批次拆分为2个小批次执行。

2. 分布式训练策略自动拆分批次 #

如果你的环境中默认启用了分布式训练策略（比如MirroredStrategy，常见于多GPU环境），TensorFlow会自动将批次均匀分配到各个设备上。若有2个设备，64的批次就会被拆分为两个32的批次分别处理。

3. 旧版本TensorFlow的bug #

部分早期TensorFlow版本在处理包含BatchNormalization层的模型时，可能存在train_on_batch错误拆分批次的bug，升级到较新版本可解决。

解决方案 #

针对上述原因，你可以尝试以下步骤修复：

1. 显式设置steps_per_execution=1 #

在编译判别器时添加该参数，强制每次调用只处理一个批次：

discriminator.compile(
    optimizer="adam",
    loss="binary_crossentropy",
    metrics="accuracy",
    steps_per_execution=1
)

2. 检查并禁用不必要的分布式策略 #

如果你的代码不需要分布式训练，确保没有默认启用相关策略。可以在代码开头添加以下代码验证：

print("当前使用的策略:", tf.distribute.get_strategy())

如果输出不是DefaultDistributionStrategy，则需要显式使用默认策略：

strategy = tf.distribute.get_strategy()
with strategy.scope():
    # 在这里重新创建和编译你的模型
    generator = create_generator(latent_dim)
    discriminator = create_discriminator((H, W, 1))
    # ...后续编译和训练代码

3. 升级TensorFlow到稳定版本 #

如果使用的是较旧的TensorFlow版本，建议升级到最新稳定版（如2.15+），以修复已知的批次处理bug。

额外代码问题提示 #

你的代码中存在一个关键错误，会导致生成器训练失败：
组合模型combined_model的输出应该是判别器的预测结果（discriminator_output），而不是生成器的输出图像。当前代码中：

combined_model = tf.keras.Model(i, generator_output)

应该修改为：

combined_model = tf.keras.Model(i, discriminator_output)

因为生成器的训练目标是让判别器将生成图像判定为真实（即输出1），所以组合模型需要以判别器的输出作为损失计算的依据，否则会出现形状不匹配的错误（生成图像是(64,28,28,1)，而目标标签是(64,)）。

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