如何在TensorFlow中导入大型数据集训练 解决序列生成内存溢出问题

方案1：实时滑动窗口生成器（无需额外存储磁盘，内存占用最低） #

直接替换原有的全量生成序列逻辑，每次训练时按需生成对应批次的时序窗口，全程不会把全量序列加载到内存：
首先自定义时序数据生成类（适配Keras训练逻辑，支持多进程加载）：

import numpy as np
from tensorflow.keras.utils import Sequence

class TimeSeriesSeqGenerator(Sequence):
    def __init__(self, df, seq_params, batch_size, is_train=False):
        self.df = df.values # 提前转成numpy数组加快索引速度
        self.seq_len = seq_params['seq_len'] # 对应你SEQ_DATA里的序列长度参数
        self.pred_len = seq_params.get('pred_len', 1) # 对应预测步长，根据自己需求改
        self.batch_size = batch_size
        self.is_train = is_train # 训练集可传True开启人工平衡逻辑
        # 计算总有效样本数
        self.total_samples = len(self.df) - self.seq_len - self.pred_len + 1

    def __len__(self):
        # 返回每轮epoch的总步数
        return int(np.ceil(self.total_samples / self.batch_size))

    def __getitem__(self, idx):
        # 计算当前batch的样本索引范围
        start = idx * self.batch_size
        end = min((idx + 1) * self.batch_size, self.total_samples)
        batch_x = []
        batch_y = []
        for i in range(start, end):
            # 生成单条时序窗口和对应标签
            x = self.df[i:i+self.seq_len]
            y = self.df[i+self.seq_len:i+self.seq_len+self.pred_len, 0] # 这里标签列索引根据你自己的数据集改
            batch_x.append(x)
            batch_y.append(y)
        # 如果你原来的SequenceCollector有人工平衡逻辑，在这里加对应处理即可
        return np.array(batch_x), np.array(batch_y)

适配你的训练流程 #

1. 去掉原有裁剪数据集的代码，直接用全量DF做拆分、归一化，原有拆分和归一化逻辑不需要修改
2. 实例化对应数据集的生成器：

# 替换原有调用SequenceCollector生成全量train_x/train_y的逻辑
train_gen = TimeSeriesSeqGenerator(main_df, SEQ_DATA, batch_size=BATCH, is_train=True)
val_gen = TimeSeriesSeqGenerator(validation_main_df, SEQ_DATA, batch_size=BATCH, is_train=False)
test_gen = TimeSeriesSeqGenerator(test_df, SEQ_DATA, batch_size=BATCH, is_train=False)

3. 修改训练调用代码：

history = model.fit(
    train_gen,
    epochs=EPOCHS,
    validation_data=val_gen,
    callbacks=[tensorboard],
    workers=4, # 根据你的CPU核心数调整，设为核心数的一半左右即可
    use_multiprocessing=True # 开启多进程加载加快速度
)

方案2：预生成序列存磁盘按需加载（适合实时生成计算量过高的场景） #

如果你的序列生成逻辑包含复杂计算，实时生成拖慢训练速度，可以提前把序列分块存到磁盘，训练时按需加载：

第一步：预生成序列分块存储 #

import os
chunk_size = 10000 # 每块存10000条序列，可根据你的磁盘空间调整
save_dir = './preprocessed_seqs/'
os.makedirs(save_dir, exist_ok=True)

# 以训练集为例，验证集、测试集逻辑相同
total_train_samples = len(main_df) - SEQ_DATA['seq_len'] - 1 + 1
chunk_idx = 0
current_chunk_x = []
current_chunk_y = []
for i in range(total_train_samples):
    x = main_df.values[i:i+SEQ_DATA['seq_len']]
    y = main_df.values[i+SEQ_DATA['seq_len'], 0] # 标签列自行修改
    current_chunk_x.append(x)
    current_chunk_y.append(y)
    if len(current_chunk_x) == chunk_size or i == total_train_samples -1:
        np.save(f'{save_dir}/train_x_chunk_{chunk_idx}.npy', np.array(current_chunk_x))
        np.save(f'{save_dir}/train_y_chunk_{chunk_idx}.npy', np.array(current_chunk_y))
        current_chunk_x = []
        current_chunk_y = []
        chunk_idx +=1

第二步：写加载生成器训练 #

def disk_seq_generator(chunk_dir, data_type, batch_size):
    chunk_idx = 0
    while True:
        # 加载当前块
        x = np.load(f'{chunk_dir}/{data_type}_x_chunk_{chunk_idx}.npy')
        y = np.load(f'{chunk_dir}/{data_type}_y_chunk_{chunk_idx}.npy')
        # 按batch输出
        for i in range(0, len(x), batch_size):
            batch_x = x[i:i+batch_size]
            batch_y = y[i:i+batch_size]
            yield batch_x, batch_y
        chunk_idx +=1
        # 所有块遍历完后重置索引，开启下一轮epoch
        if chunk_idx == len([f for f in os.listdir(chunk_dir) if f.startswith(f'{data_type}_x_chunk')]):
            chunk_idx = 0

# 实例化生成器
train_gen = disk_seq_generator('./preprocessed_seqs/', 'train', BATCH)
val_gen = disk_seq_generator('./preprocessed_seqs/', 'val', BATCH)
# 计算每轮步数
train_steps = total_train_samples // BATCH
val_steps = len(validation_main_df) // BATCH
# 训练调用
history = model.fit(
    train_gen,
    steps_per_epoch=train_steps,
    epochs=EPOCHS,
    validation_data=val_gen,
    validation_steps=val_steps,
    callbacks=[tensorboard]
)

注意事项 #

• 时序数据拆分时必须保持shuffle=False，不能打乱时间顺序，避免数据泄露
• 归一化的scaler必须仅用训练集拟合，验证集、测试集直接复用训练集的scaler，你的原有逻辑是正确的，无需修改
• 如果你的任务允许打乱训练序列，可在每轮epoch结束后调整预生成块的加载顺序，进一步提升模型泛化性

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