TensorFlow中fit是否传入完整batch size？自定义损失函数报错求解

问题背景 #

使用TensorFlow训练模型时，设置batch size为1960，但发现fit似乎每次仅向自定义损失函数传入单个值，同时出现以下形状不匹配错误：

ValueError: Shapes must be equal rank, but are 0 and 2
From merging shape 0 with other shapes. for '{{node generator_loss/weighted_loss/packed}} = Pack[N=2, T=DT_FLOAT, axis=0](generator_loss/truediv, generator_loss/Mul)' with input shapes: [], [?,1].

原代码如下：

import tensorflow as tf
from tensorflow import keras

import math

import numpy as np


from tensorflow.keras.layers import Dense

from tensorflow.keras.layers import Dropout

chunksize = 40

stepsize = 1961



x = tf.data.Dataset.from_tensor_slices(np.random.rand(9000, 40, 7))
y = tf.data.Dataset.from_tensor_slices(np.random.rand(9000, 40, 7))
dataset = tf.data.Dataset.zip((x, y)).batch(stepsize - 1)

def Generator():
    model_m = keras.models.Sequential()
    

    model_m.add(tf.keras.layers.Conv1D(3, 1, input_shape=(40, 7)))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50, return_sequences=True))
    model_m.add(Dropout(0.2))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50,return_sequences=True))
    model_m.add(Dropout(0.2))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50,return_sequences=True))
    model_m.add(Dropout(0.2))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50))
    model_m.add(Dropout(0.2))
    model_m.add(Dense(units=1, activation="sigmoid"))
    
    return model_m

model = Generator()


def generator_loss(target, genor_output1):
    
    dat = tf.cast(tf.reshape(target, (len(target), 1)), tf.float32)
    
    mult = tf.reduce_sum(dat)
    dat = tf.math.multiply(genor_output1, dat)    
    dat2 = tf.reduce_sum(dat)
    dat2 = tf.math.divide(tf.math.add(tf.math.add(tf.math.abs(mult), 1), mult), tf.math.add(tf.math.add(tf.math.abs(mult), 1), dat2))    
    return dat2, dat
generator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.00007, epsilon=2e-4, beta_1=0.5)
model.compile(loss=generator_loss, optimizer="Adam")
model.fit(dataset, epochs=100, batch_size=(stepsize - 1))

错误原因 #

1. 损失函数返回值违规：Keras自定义损失函数必须仅返回单个标量损失值，但你的generator_loss返回了两个值——标量dat2（shape []）和张量dat（shape [?,1]），导致TensorFlow在打包损失数据时因形状秩不匹配报错。
2. batch参数冲突：你已经通过dataset.batch(stepsize - 1)设置了批量大小，却又在model.fit中重复指定batch_size，这会干扰数据加载逻辑，造成“仅传入单个值”的错觉。
3. 静态形状获取不可靠：用len(target)获取batch大小在TensorFlow动态图模式下可能失效，因为target的静态形状可能未完全确定。

解决方案 #

修正后的代码 #

import tensorflow as tf
from tensorflow import keras
import numpy as np
from tensorflow.keras.layers import Dense, Dropout

chunksize = 40
stepsize = 1961

# 构建数据集，已指定batch大小
x = tf.data.Dataset.from_tensor_slices(np.random.rand(9000, 40, 7))
y = tf.data.Dataset.from_tensor_slices(np.random.rand(9000, 40, 7))
dataset = tf.data.Dataset.zip((x, y)).batch(stepsize - 1)

def Generator():
    model_m = keras.models.Sequential()
    model_m.add(tf.keras.layers.Conv1D(3, 1, input_shape=(40, 7)))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50, return_sequences=True))
    model_m.add(Dropout(0.2))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50, return_sequences=True))
    model_m.add(Dropout(0.2))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50, return_sequences=True))
    model_m.add(Dropout(0.2))
    model_m.add(tf.compat.v1.keras.layers.CuDNNLSTM(units=50))
    model_m.add(Dropout(0.2))
    model_m.add(Dense(units=1, activation="sigmoid"))
    return model_m

model = Generator()

def generator_loss(target, genor_output1):
    # 用tf.shape获取动态batch大小，替代静态的len(target)
    batch_size = tf.shape(target)[0]
    dat = tf.cast(tf.reshape(target, (batch_size, 1)), tf.float32)
    
    mult = tf.reduce_sum(dat)
    dat = tf.math.multiply(genor_output1, dat)    
    dat2 = tf.reduce_sum(dat)
    # 仅返回单个标量损失值
    return tf.math.divide(tf.math.add(tf.math.add(tf.math.abs(mult), 1), mult), 
                         tf.math.add(tf.math.add(tf.math.abs(mult), 1), dat2))

# 使用定义好的自定义优化器（原代码中未正确使用）
generator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.00007, epsilon=2e-4, beta_1=0.5)
model.compile(loss=generator_loss, optimizer=generator_optimizer)
# 移除fit中的batch_size参数，避免冲突
model.fit(dataset, epochs=100)

关键调整说明 #

• 移除generator_loss的多余返回值，确保仅返回标量损失。
• 删除model.fit中的batch_size参数，依赖数据集已设置的批量大小。
• 用tf.shape(target)[0]获取动态batch大小，避免静态形状导致的错误。
• 修正优化器的使用，将定义好的generator_optimizer传入model.compile，替代原代码中的字符串"Adam"。

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