TensorFlow数据集拆分与批量配置引发OUT_OF_RANGE错误的求助

我刚接触TensorFlow的数据生成器和Dataset API，现在在批量大小、epoch数和步数的配置上遇到了麻烦，一直没法解决Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence这个错误。我试过用数据生成器的chunk大小、完整数据集大小、拆分后的数据集大小来配置，但都不管用。

我的数据生成器代码 #

def data_generator(df, chunk_size):

  total_number_sample = 10000
  
  for start_idx in range(1, total_number_sample , chunk_size):
    
    end_idx = start_idx + chunk_size-1 

    df_subset = df.where(col('idx').between(start_idx, end_idx))

    feature = np.array(df_subset.select("vector_features_scaled").rdd.map(lambda row: row[0].toArray()).collect())
    label = df_subset.select("ptype_s_l_m_v").toPandas().values.flatten()

    yield feature, label

构建TensorFlow Dataset #

dataset = tf.data.Dataset.from_generator(
    lambda: data_generator(df, chunk_size),
    output_signature=(
        tf.TensorSpec(shape=(None, 24), dtype=tf.float32),
        tf.TensorSpec(shape=(None, 4), dtype=tf.float32)
    ))

训练/验证集的拆分与批量处理 #

import math

batch_sz = 100
split_ratio = .9
split_size = math.floor((chunk_size*10) * split_ratio) 

train_dataset = dataset.take(split_size).batch(batch_sz)
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)
test_dataset = dataset.skip(split_size).batch(batch_sz)
test_dataset = test_dataset.prefetch(tf.data.experimental.AUTOTUNE)


steps_per_epoch=math.ceil((10000 * split_ratio) / batch_sz)
validation_steps=math.ceil((10000-split_size) / batch_sz)

model.fit(train_dataset, 
          steps_per_epoch=steps_per_epoch, 
          epochs=3, 
          validation_data=test_dataset, 
          validation_steps=validation_steps,
          verbose=2)

results = model.evaluate(dataset.batch(batch_sz))

遇到的错误 #

W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence
     [[{{node IteratorGetNext}}]]
/usr/lib/python3.11/contextlib.py:155: UserWarning: Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches. You may need to use the `.repeat()` function when building your dataset.
  self.gen.throw(typ, value, traceback)

注意：不进行批量处理时，model.fit()和model.evaluate()都能正常工作，但一加批量就报错。

问题根源分析 #

你遇到的OUT_OF_RANGE错误本质是数据生成的总量和你配置的steps_per_epoch、validation_steps完全不匹配，导致TensorFlow在训练过程中提前耗尽了数据。具体问题点有几个：

1. 数据集拆分逻辑完全错误
   你用dataset.take(split_size)和dataset.skip(split_size)拆分，但这里的split_size是基于chunk_size*10计算的，而Dataset.from_generator生成的每个元素是一整个chunk的样本（也就是data_generator每次yield的一组feature和label），不是单个样本。所以take(split_size)会直接取split_size个chunk，这和你想要的「按样本数拆分训练/验证集」逻辑完全不符。

2. 批量处理的顺序搞反了
   你先对「chunk级别的dataset」做了take/skip拆分，再做batch，这会把多个chunk打包成一个训练batch，每个batch的样本数是batch_sz * chunk_size，和你设置的batch_sz=100完全不是一回事。

3. steps_per_epoch计算逻辑混乱
   你用总样本数的90%除以batch size，但你的训练集根本不是按样本数拆分的，导致配置的步数远大于实际能生成的batch数量，训练到一半就没数据了。

修复方案 #

步骤1：先把数据集拆成单个样本（核心修正）

不管你用chunk式生成器还是单个样本生成器，首先要让Dataset的每个元素对应单个训练样本，这样后续的batch和拆分逻辑才会符合预期：

# 方式1：修改生成器，每次yield单个样本（适合小数据量）
def data_generator(df, total_samples=10000):
    for idx in range(1, total_samples+1):
        df_subset = df.where(col('idx') == idx)
        feature = np.array(df_subset.select("vector_features_scaled").rdd.map(lambda row: row[0].toArray()).collect())[0]
        label = df_subset.select("ptype_s_l_m_v").toPandas().values.flatten()[0]
        yield feature, label

# 方式2：保留原chunk生成器，用unbatch拆成单个样本（适合大数据量，保留Spark的chunk读取效率）
dataset = tf.data.Dataset.from_generator(
    lambda: data_generator(df, chunk_size),
    output_signature=(
        tf.TensorSpec(shape=(None, 24), dtype=tf.float32),
        tf.TensorSpec(shape=(None, 4), dtype=tf.float32)
    )).unbatch()  # 关键：把每个chunk拆成单个样本

步骤2：按样本数正确拆分训练/验证集

现在Dataset的每个元素是单个样本，就可以按总样本数的比例拆分了：

total_samples = 10000
split_ratio = 0.9
train_samples = int(total_samples * split_ratio)
val_samples = total_samples - train_samples
batch_sz = 100

# 先按样本数拆分，再做批量处理
train_dataset = dataset.take(train_samples).batch(batch_sz).prefetch(tf.data.AUTOTUNE)
val_dataset = dataset.skip(train_samples).batch(batch_sz).prefetch(tf.data.AUTOTUNE)

步骤3：正确配置训练参数（建议省略steps_per_epoch）

当你传入已经batch好的Dataset时，TensorFlow可以自动计算每个epoch的步数，完全不需要手动指定steps_per_epoch和validation_steps，这能彻底避免步数不匹配的问题：

model.fit(
    train_dataset,
    epochs=3,
    validation_data=val_dataset,
    verbose=2
)

如果一定要手动指定步数（比如和其他流程对齐），要按真实的batch数计算：

steps_per_epoch = train_samples // batch_sz
# 补上最后一个不足batch_sz的样本组
if train_samples % batch_sz != 0:
    steps_per_epoch += 1

validation_steps = val_samples // batch_sz
if val_samples % batch_sz != 0:
    validation_steps += 1

model.fit(
    train_dataset,
    steps_per_epoch=steps_per_epoch,
    epochs=3,
    validation_data=val_dataset,
    validation_steps=validation_steps,
    verbose=2
)

步骤4：多epoch训练的兜底处理

如果要训练多轮epoch，需要确保数据集可以重复生成，在构建Dataset时加上repeat()：

# 无限重复，由model.fit的epochs控制停止
train_dataset = dataset.take(train_samples).repeat().batch(batch_sz).prefetch(tf.data.AUTOTUNE)
# 此时必须指定steps_per_epoch，否则会无限训练
model.fit(train_dataset, steps_per_epoch=steps_per_epoch, epochs=3, ...)

简化后的完整可运行示例 #

import math
import tensorflow as tf
import numpy as np
from pyspark.sql.functions import col

# 保留原chunk生成器
def data_generator(df, chunk_size):
    total_number_sample = 10000
    for start_idx in range(1, total_number_sample , chunk_size):
        end_idx = start_idx + chunk_size-1 
        df_subset = df.where(col('idx').between(start_idx, end_idx))
        feature = np.array(df_subset.select("vector_features_scaled").rdd.map(lambda row: row[0].toArray()).collect())
        label = df_subset.select("ptype_s_l_m_v").toPandas().values.flatten()
        yield feature, label

# 构建Dataset并拆成单个样本
dataset = tf.data.Dataset.from_generator(
    lambda: data_generator(df, chunk_size=100),  # 这里chunk_size可以设为batch_sz，提升读取效率
    output_signature=(
        tf.TensorSpec(shape=(None, 24), dtype=tf.float32),
        tf.TensorSpec(shape=(None, 4), dtype=tf.float32)
    )).unbatch()

# 拆分训练/验证集
total_samples = 10000
split_ratio = 0.9
train_samples = int(total_samples * split_ratio)
batch_sz = 100

train_dataset = dataset.take(train_samples).batch(batch_sz).prefetch(tf.data.AUTOTUNE)
val_dataset = dataset.skip(train_samples).batch(batch_sz).prefetch(tf.data.AUTOTUNE)

# 训练（自动计算步数）
model.fit(
    train_dataset,
    epochs=3,
    validation_data=val_dataset,
    verbose=2
)

# 评估用验证集即可，不要用完整数据集
results = model.evaluate(val_dataset)

额外优化建议 #

• 你的生成器中用了rdd.map和collect()，在大数据量下会把数据拉到Spark Driver，性能很差，建议用Spark的向量化API（比如pyspark.ml.feature的相关转换）直接处理成TensorFlow可读取的格式，避免collect。
• 如果你的chunk_size设置成和batch_sz一致，可以直接在生成器中yield batch级别的数据，这样就不需要unbatch，进一步提升效率。

备注：内容来源于stack exchange，提问作者Jonathan Roy