Pandas读取大数据集OOM问题及CNN训练内存优化咨询

1. 逐文件分批处理，避免加载全量数据 #

核心思路是不将所有文件拼接成单个大DataFrame，而是逐个读取文件、生成样本后立即释放内存，仅保留最终的训练样本集合：

先统一类别编码映射 #

提前固定类别与编码的对应关系，避免不同文件因factorize()导致编码不一致：

# 手动定义已知的5类映射（无需从数据中推导）
siniflar_map = {'WK':0, 'N1':1, 'N2':2, 'N3':3, 'REM':4}

逐文件生成样本 #

import pandas as pd
import numpy as np
import os
import gc

data_dir = r'/content/drive/MyDrive/uyku_parquet'
# 获取所有Parquet文件列表
files = [f.split('.')[0] for f in os.listdir(data_dir) if f.endswith('.parquet')]

# 初始化样本存储容器
X_total = []
y_total = []

for file in files:
    # 读取单个Parquet文件
    df = pd.read_parquet(f'{data_dir}/{file}.parquet')
    # 统一编码类别列，过滤无效类别
    df['SleepStaging'] = df['SleepStaging'].map(siniflar_map)
    df = df[df['SleepStaging'].notna()]
    
    # 按类别切割生成6000×19的样本
    for class_code in siniflar_map.values():
        class_data = df[df['SleepStaging'] == class_code].to_numpy()
        # 计算可生成的完整样本数
        num_valid_samples = class_data.shape[0] // 6000
        for i in range(num_valid_samples):
            start_idx = i * 6000
            end_idx = start_idx + 6000
            # 提取特征列（第1到19列）
            sample = class_data[start_idx:end_idx, 1:20]
            X_total.append(sample)
            y_total.append(class_code)
    
    # 立即释放当前文件的内存
    del df
    gc.collect()

# 转换为训练用数组并保存到磁盘（避免重复处理）
X = np.array(X_total)
y = np.array(y_total)
np.savez('/content/drive/MyDrive/sleep_training_data.npz', X=X, y=y)

2. 优化数据类型，降低单文件内存占用 #

读取Parquet时指定更小的数据类型，可将单文件内存占用降低约50%：

# 根据实际列定义数据类型映射
dtype_config = {
    'SleepStaging': 'int8',  # 类别编码仅需0-4，用int8足够
    # 所有特征列从float64转为float32
    **{col: 'float32' for col in df.columns if col != 'SleepStaging'}
}

# 读取时指定类型
df = pd.read_parquet(f'{data_dir}/{file}.parquet', dtype=dtype_config)

3. 针对类别不平衡的补充优化 #

除了增加数据量，还可通过以下方式提升模型对少数类的识别能力：

• 过采样少数类别：对样本量不足的类别（如N1、REM），可重复利用剩余不足6000行的数据生成样本，或对时序数据做轻微噪声注入、时间平移等增强操作。
• 设置类别权重：在模型训练时给少数类别更高的权重：

from sklearn.utils.class_weight import compute_class_weight

class_weights = compute_class_weight('balanced', classes=np.unique(y), y=y)
class_weight_dict = dict(zip(np.unique(y), class_weights))

# Keras训练时传入权重
model.fit(X, y, class_weight=class_weight_dict, epochs=10, batch_size=32)

4. 环境层面的优化 #

• Google Colab：升级到Colab Pro/Pro+，获取最高52GB的内存配额，同时使用GPU/TPU加载数据到显存，减少内存占用。
• 流式训练：用TensorFlow Dataset加载预生成的样本，无需将全量数据放入内存：

import tensorflow as tf

# 从磁盘加载数据并构建数据集
data = np.load('/content/drive/MyDrive/sleep_training_data.npz')
dataset = tf.data.Dataset.from_tensor_slices((data['X'], data['y']))
# 打乱、分批、预取
dataset = dataset.shuffle(len(data['X'])).batch(32).prefetch(tf.data.AUTOTUNE)

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

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