如何向Tensorflow fit()方法传入预先计算好的数据批次？

完全可以实现，不需要手写基于GradientTape的自定义训练循环，有两种Keras原生支持的方案可以直接传入预构建好的固定批次给model.fit()，两种方案都不会自动随机拆分重组批次内的样本，完全满足自定义批次属性的要求。

适合所有批次可以提前一次性构建完成、内存可以装下的场景。
操作步骤：

• 按照你的损失计算要求，提前把所有训练样本切分成若干个满足属性约束的批次，每个批次对应匹配模型输入输出形状的特征张量、标签张量，不要留随机拆分的逻辑。
• 用生成器把预构建好的批次封装成tf.data.Dataset，注意不要在数据集流水线中调用.batch()、.shuffle()这类会重组样本的操作。
• 调用model.fit()时直接传入这个数据集，不需要额外指定batch_size参数，Keras会逐批读取你预构建好的样本完成训练。

示例代码：

import tensorflow as tf

# 预构建批次逻辑：此处替换成你自己的批次生成代码，保证每个批次满足约束
# 示例为单输入分类任务，每个批次x形状为(batch_size, 特征维度)，y形状为(batch_size, 类别数)
pre_batched_train = [build_one_batch(i) for i in range(total_batch_num)]

def batch_generator():
    for x, y in pre_batched_train:
        yield x, y

# 定义张量签名，形状和你的批次匹配，None表示可变批次大小（如果你的批次大小固定可以写死数值）
train_ds = tf.data.Dataset.from_generator(
    batch_generator,
    output_signature=(
        tf.TensorSpec(shape=(None, feature_dim), dtype=tf.float32),
        tf.TensorSpec(shape=(None, num_classes), dtype=tf.float32)
    )
)

# 开始训练，不要传batch_size参数
model.fit(
    train_ds,
    epochs=10,
    # 如果要固定批次训练顺序，加shuffle=False；开shuffle只会打乱批次先后顺序，不会改批次内样本
    shuffle=False
)

适合数据量较大、无法一次性把所有预构建批次加载到内存的场景，支持多线程预加载，训练效率更高。
操作步骤：

• 继承tf.keras.utils.Sequence类，重写核心方法：
  • __len__：返回一个epoch训练需要的总批次数
  • __getitem__：传入批次索引，返回对应索引下满足要求的预构建批次（可以实时构建，不需要提前存在内存里）
  • （可选）on_epoch_end：每个epoch结束后的自定义逻辑，比如重新生成新的满足约束的批次
• 初始化这个类的实例，直接传给model.fit()即可，同样不需要指定batch_size参数。

示例代码：

from tensorflow.keras.utils import Sequence

class CustomBatchSequence(Sequence):
    def __init__(self, raw_x, raw_y, batch_size, build_batch_func):
        self.x = raw_x
        self.y = raw_y
        self.batch_size = batch_size
        self.build_batch = build_batch_func # 你自己的批次构建逻辑，生成满足约束的批次
        self.n_batches = len(raw_x) // batch_size

    def __len__(self):
        return self.n_batches

    def __getitem__(self, idx):
        # 按照索引返回对应批次，内部逻辑完全自定义，保证返回的批次满足损失计算要求
        return self.build_batch(self.x, self.y, idx, self.batch_size)

    def on_epoch_end(self):
        # 如果需要每个epoch重新生成批次，可以在这里实现逻辑，不需要就留空
        pass

# 初始化序列
train_seq = CustomBatchSequence(train_x, train_y, batch_size=32, build_batch_func=your_batch_logic)

# 传入fit训练
model.fit(
    train_seq,
    epochs=10,
    shuffle=False
)

• 两种方案都完全兼容model.fit()的原生功能，回调函数、验证集传入、指标计算、混合精度训练等特性都可以正常使用，不需要修改模型或者损失函数的外层逻辑。
• 训练时传入的验证集如果也有固定批次要求，可以用同样的方式构建后传给validation_data参数。
• 如果开启shuffle=True，Keras只会打乱不同批次之间的训练顺序，不会拆分、重组单个批次内部的样本，不会破坏你预设的批次属性。

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