基于Keras Sequence API实现跨CSV文件的行级随机训练测试拆分

需求说明 #

我使用Keras训练神经网络，通过Sequence API实现数据生成器处理大规模CSV数据。数据存储在多个大小不一的CSV文件中，其中A-D列为特征变量，E列为标签变量。当前采用按文件拆分训练/测试集的方式，但希望改为按所有数据行拆分，让单个CSV文件中同时包含训练行和测试行，提升数据随机性。由于数据总量过大无法全量加载到内存，需保持小批量处理模式，且训练、测试生成器需基于同一类实现，确保测试生成器能“记住”哪些行已用于训练，仅使用剩余行。

以下是当前按文件拆分的基础代码：

import keras
import numpy as np
import pandas as pd
import glob

class DataGenerator(keras.utils.Sequence):
     
    def __init__(self, list_of_csv_files, predictor_column_names, predicted_column_name, shuffle=True, batch_size=5):
        self.list_of_csv_files = list_of_csv_files
        self.predictor_column_names = predictor_column_names
        self.predicted_column_name = predicted_column_name
        self.shuffle = shuffle
        self.batch_size = batch_size
        self.on_epoch_end()

    def on_epoch_end(self): # Updates/shuffles indexes after each epoch
        if self.shuffle:
            np.random.shuffle(self.list_of_csv_files)
            
    def __len__(self) :
        return np.ceil(len(self.list_of_csv_files) / float(self.batch_size)).astype(int)
    
    def __getitem__(self, idx): # Generates one batch of data
        files_current_batch = self.list_of_csv_files[idx*self.batch_size : idx*self.batch_size + self.batch_size]
        X = pd.concat([pd.read_csv(s, usecols=self.predictor_column_names) for s in files_current_batch], axis=0, ignore_index=True)
        y = pd.concat([pd.read_csv(s, usecols=self.predicted_column_name) for s in files_current_batch], axis=0, ignore_index=True)
        X = np.array(X)
        y = np.array(y)
        return X, y

# define parameters for data generator
list_of_csv_files_train = glob.glob("/train_directory_name/" + "*.csv")
list_of_csv_files_test = glob.glob("/test_directory_name/" + "*.csv")
predictor_column_names = ["A", "B", "C", "D"]
predicted_column_name = ["E"]
batch_size = 5
# build model    
model = keras.models.Sequential()
model.add(keras.layers.Dense(64, activation="relu"))
model.add(keras.layers.Dense(1, activation='linear'))
model.compile()    
# use data generators
train_generator = DataGenerator(list_of_csv_files_train, predictor_column_names, predicted_column_name, shuffle=True, batch_size=5)
test_generator = DataGenerator(list_of_csv_files_test, predictor_column_names, predicted_column_name, shuffle=True, batch_size=5)
# train model
model.fit(train_generator, validation_data=test_generator, epochs=num_epochs)

解决方案 #

核心思路 #

通过预先生成每个CSV文件的训练/测试行掩码（轻量级二进制文件存储），让生成器根据掩码读取对应行的数据，既实现按行拆分，又避免全量加载数据。

步骤1：预生成训练/测试行掩码 #

运行一次该脚本，为每个CSV文件生成训练行和测试行的掩码，存储为.npy文件（占用空间极小）：

import numpy as np
import pandas as pd
import glob
import os

# 参数配置
test_size = 0.2  # 测试集比例
all_csv_files = glob.glob("/path/to/all/csv/files/*.csv")  # 所有CSV文件路径
mask_save_dir = "./train_test_masks"  # 掩码存储目录

# 创建掩码存储目录
os.makedirs(mask_save_dir, exist_ok=True)

for csv_path in all_csv_files:
    # 获取文件名（不含后缀）
    file_base_name = os.path.basename(csv_path).replace(".csv", "")
    # 高效计算文件行数（不加载全量数据）
    with open(csv_path, 'r') as f:
        total_rows = sum(1 for _ in f) - 1  # 减去表头行
    # 生成训练行掩码：True表示该行用于训练，False用于测试
    train_mask = np.random.choice([True, False], size=total_rows, p=[1-test_size, test_size])
    # 保存训练掩码
    np.save(os.path.join(mask_save_dir, f"{file_base_name}_train_mask.npy"), train_mask)
    # 生成并保存测试掩码（训练掩码的取反）
    test_mask = ~train_mask
    np.save(os.path.join(mask_save_dir, f"{file_base_name}_test_mask.npy"), test_mask)

步骤2：修改DataGenerator类 #

更新后的生成器支持加载掩码，根据is_train参数读取训练或测试行：

import keras
import numpy as np
import pandas as pd
import glob
import os

class DataGenerator(keras.utils.Sequence):
    def __init__(self, csv_files, predictor_cols, target_col, mask_dir, is_train=True, shuffle=True, batch_size=32):
        self.csv_files = csv_files
        self.predictor_cols = predictor_cols
        self.target_col = target_col
        self.mask_dir = mask_dir
        self.is_train = is_train
        self.shuffle = shuffle
        self.batch_size = batch_size
        
        # 预加载所有文件的掩码（仅加载掩码，不加载数据）
        self.file_masks = {}
        for csv_path in self.csv_files:
            file_base = os.path.basename(csv_path).replace(".csv", "")
            mask_filename = f"{file_base}_train_mask.npy" if is_train else f"{file_base}_test_mask.npy"
            self.file_masks[csv_path] = np.load(os.path.join(mask_dir, mask_filename))
        
        # 初始化文件索引列表
        self.file_indices = np.arange(len(self.csv_files))
        self.on_epoch_end()

    def on_epoch_end(self):
        # 每个epoch后打乱文件顺序（仅训练生成器生效）
        if self.shuffle:
            np.random.shuffle(self.file_indices)

    def __len__(self):
        # 计算总批次数量
        return np.ceil(len(self.file_indices) / self.batch_size).astype(int)

    def __getitem__(self, idx):
        # 获取当前批次的文件索引
        batch_file_indices = self.file_indices[idx*self.batch_size : (idx+1)*self.batch_size]
        batch_files = [self.csv_files[i] for i in batch_file_indices]
        
        X_batch = []
        y_batch = []
        for csv_path in batch_files:
            # 读取当前文件的指定列
            df = pd.read_csv(csv_path, usecols=self.predictor_cols + self.target_col)
            # 应用掩码筛选行
            mask = self.file_masks[csv_path]
            selected_df = df[mask]
            # 分离特征和标签
            X_batch.append(selected_df[self.predictor_cols].values)
            y_batch.append(selected_df[self.target_col].values)
        
        # 合并批次数据
        X = np.concatenate(X_batch, axis=0)
        y = np.concatenate(y_batch, axis=0)
        return X, y

步骤3：使用生成器训练模型 #

# 参数配置
all_csv_files = glob.glob("/path/to/all/csv/files/*.csv")
predictor_cols = ["A", "B", "C", "D"]
target_col = ["E"]
mask_dir = "./train_test_masks"
batch_size = 32
num_epochs = 10

# 创建训练和测试生成器
train_generator = DataGenerator(
    csv_files=all_csv_files,
    predictor_cols=predictor_cols,
    target_col=target_col,
    mask_dir=mask_dir,
    is_train=True,
    shuffle=True,
    batch_size=batch_size
)

test_generator = DataGenerator(
    csv_files=all_csv_files,
    predictor_cols=predictor_cols,
    target_col=target_col,
    mask_dir=mask_dir,
    is_train=False,
    shuffle=False,
    batch_size=batch_size
)

# 构建并训练模型
model = keras.models.Sequential()
model.add(keras.layers.Dense(64, activation="relu"))
model.add(keras.layers.Dense(1, activation='linear'))
model.compile(optimizer='adam', loss='mean_squared_error')

model.fit(train_generator, validation_data=test_generator, epochs=num_epochs)

关键说明 #

• 掩码复用：掩码仅需预生成一次，后续训练可直接复用，确保训练/测试行的一致性
• 内存效率：每次仅加载一个批次的CSV文件，并筛选对应掩码的行，无需全量加载数据
• 随机性控制：训练生成器每个epoch会打乱文件顺序，提升训练数据的随机性；测试生成器不打乱，保证评估稳定性

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