Keras解决LSTM序列数据加载与处理的内存问题

问题背景 #

要基于TensorFlow & Keras训练LSTM模型，数据源是多个按时间排序的CSV文件，每个文件包含传感器数值、side（取值1/2）和目标值。需求是按每个CSV文件+side分组生成序列：特征为tₙ到tₙ+10的10步序列，目标为序列中间位置tₙ+5的数值。

现有代码在子集数据上可运行，但合并全量CSV为DataFrame后生成序列时，因内存不足失败，需要改用Keras的高效数据加载方案。

方案一：自定义Keras Sequence类 #

Keras的Sequence是专为批量加载数据设计的抽象类，训练时会逐批次生成数据，不会一次性把所有数据加载到内存，适合用Pandas处理数据的场景。

实现步骤 #

1. 预收集分组信息：遍历所有CSV文件，记录每个有效分组（CSV路径+ID+side）的有效数据范围（跳过前后60行），避免提前加载全量数据。
2. 实现自定义Sequence类：每次按批次读取对应分组的片段，生成训练所需的特征序列和目标值。

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

n_steps = 10
pred_step = 5
skip_rows = 60
feat_cols = ["sensor1", "sensor2"]  # 替换为你的特征列
feat_bin_cols = ["bin_feat1"]       # 替换为你的二值特征列
target = ["target_col"]             # 替换为你的目标列

# 第一步：收集所有有效分组的元信息
group_info = []
csv_paths = ["data1.csv", "data2.csv", ...]  # 替换为你的所有CSV文件路径

for csv_path in csv_paths:
    # 仅读取分组标识列，避免加载全量数据
    df_meta = pd.read_csv(csv_path, usecols=["ID", "SIDE"])
    for side in [1, 2]:
        df_side = df_meta[df_meta["SIDE"] == side]
        for id_val in df_side["ID"].unique():
            total_rows = len(df_meta[(df_meta["ID"] == id_val) & (df_meta["SIDE"] == side)])
            valid_rows = total_rows - 2 * skip_rows
            # 确保分组能生成至少一个序列
            if valid_rows > n_steps:
                group_info.append({
                    "csv_path": csv_path,
                    "id": id_val,
                    "side": side,
                    "valid_start": skip_rows,
                    "valid_end": total_rows - skip_rows,
                    "num_sequences": valid_rows - n_steps
                })

total_sequences = sum(g["num_sequences"] for g in group_info)

# 第二步：自定义Sequence类
class LSTMDataSequence(Sequence):
    def __init__(self, group_info, batch_size):
        self.group_info = group_info
        self.batch_size = batch_size
        # 预生成所有序列的索引映射：(分组索引, 序列起始行)
        self.sequence_indices = []
        for group_idx, g in enumerate(group_info):
            for seq_start in range(g["valid_start"], g["valid_end"] - n_steps):
                self.sequence_indices.append((group_idx, seq_start))
        np.random.shuffle(self.sequence_indices)  # 训练前打乱序列顺序

    def __len__(self):
        # 返回总批次数量
        return int(np.ceil(total_sequences / self.batch_size))

    def __getitem__(self, idx):
        # 生成第idx批次的数据
        batch_indices = self.sequence_indices[idx*self.batch_size : (idx+1)*self.batch_size]
        batch_x, batch_y = [], []
        for group_idx, seq_start in batch_indices:
            g = self.group_info[group_idx]
            # 仅读取当前序列需要的行，减少内存占用
            df = pd.read_csv(
                g["csv_path"],
                skiprows=seq_start,
                nrows=n_steps + 1,
                usecols=feat_cols + feat_bin_cols + target + ["ID", "SIDE"]
            )
            # 过滤当前分组的数据，避免跨分组读取
            df = df[(df["ID"] == g["id"]) & (df["SIDE"] == g["side"])]
            # 提取特征序列和目标值
            feat_seq = df[feat_cols + feat_bin_cols].values[:n_steps]
            target_val = df[target].values[pred_step]
            batch_x.append(feat_seq)
            batch_y.append(target_val)
        return np.array(batch_x), np.array(batch_y)

# 使用示例
batch_size = 32
train_sequence = LSTMDataSequence(group_info, batch_size)
model.fit(train_sequence, epochs=10, ...)

方案二：使用tf.data.Dataset（TensorFlow原生方案） #

如果数据量极大，推荐用tf.data.Dataset，它支持从CSV文件直接读取，配合窗口操作生成序列，全程无需加载全量数据到内存，性能更优。

实现步骤 #

1. 读取CSV并分组：用tf.data.experimental.make_csv_dataset读取所有CSV，按ID和SIDE分组。
2. 生成序列窗口：对每个分组的数据应用窗口操作，提取特征序列和对应目标值。
3. 批量预处理：打乱序列、设置批量大小并开启预取优化。

import tensorflow as tf

n_steps = 10
pred_step = 5
skip_rows = 60
feat_cols = ["sensor1", "sensor2"]  # 替换为你的特征列
feat_bin_cols = ["bin_feat1"]       # 替换为你的二值特征列
target_col = "target_col"           # 替换为你的目标列
csv_paths = ["data1.csv", "data2.csv", ...]  # 替换为你的所有CSV文件路径

def process_group(group):
    # 拼接特征列，提取目标列
    features = tf.concat([group[col] for col in feat_cols + feat_bin_cols], axis=-1)
    target = group[target_col]
    
    # 跳过前后60行无效数据
    features = features[skip_rows:-skip_rows]
    target = target[skip_rows:-skip_rows]
    
    # 生成滑动窗口，每个窗口包含n_steps+1个样本
    window_size = n_steps + 1
    feat_window = tf.data.Dataset.from_tensor_slices(features).window(window_size, shift=1, drop_remainder=True)
    target_window = tf.data.Dataset.from_tensor_slices(target).window(window_size, shift=1, drop_remainder=True)
    
    # 从窗口中提取特征序列和对应目标
    def extract_seq_target(feat_win, target_win):
        feat_seq = tf.stack(list(feat_win))[:n_steps]
        target_val = list(target_win)[pred_step]
        return feat_seq, target_val
    
    return tf.data.Dataset.zip((feat_window, target_window)).map(extract_seq_target)

# 读取CSV并按ID+SIDE分组
dataset = tf.data.experimental.make_csv_dataset(
    csv_paths,
    batch_size=1,
    label_name=None,
    select_columns=feat_cols + feat_bin_cols + [target_col, "ID", "SIDE"],
    shuffle=False  # 时间序列不能打乱分组顺序
).group_by_window(
    key_func=lambda x: (x["ID"], x["SIDE"]),
    reduce_func=lambda key, ds: ds.flat_map(process_group),
    window_size=tf.int64.max  # 处理每个分组的全部数据
)

# 打乱序列、设置批量大小并开启预取优化
dataset = dataset.shuffle(10000).batch(32).prefetch(tf.data.AUTOTUNE)

# 训练模型
model.fit(dataset, epochs=10, ...)

关键优化点 #

• 避免全量加载：两种方案都不会一次性加载所有CSV数据，而是按需读取或处理。
• 严格分组逻辑：完全按照CSV文件+ID+side的要求生成序列，符合业务规则。
• 内存高效：每次仅处理当前批次所需的序列数据，大幅降低内存占用。

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