如何以内存高效方式将大型数据集采样为正负样本数量相等

针对你需要从大型数据集中高效获取等量正负样本（从4:2调整为2:2）的需求，这里有两个实用的方案，既能解决内存占用高的问题，又能精准达到采样目标：

方法一：分块读取采样（适合超大型数据集） #

如果你的数据集大到无法一次性加载到内存（比如GB/TB级的CSV/Parquet文件），分块处理是最优选择——我们只加载部分数据到内存，逐块采集样本直到达到目标数量：

import pandas as pd

# 第一步：快速统计各类样本总数（只加载class列，内存占用极低）
class_counts = pd.read_csv('your_dataset.csv', usecols=['class']).value_counts()
target_count = class_counts.min()  # 这里对应你需要的每类样本数：2

# 第二步：分块读取并按需采样
sampled_chunks = []
chunk_size = 10_000  # 根据你的内存情况调整，比如内存小就设成5000

for chunk in pd.read_csv('your_dataset.csv', chunksize=chunk_size):
    # 对每个类别，计算还需要采集多少样本
    for cls in class_counts.index:
        collected = sum(len(c) for c in sampled_chunks if c['class'].eq(cls).any())
        remaining = target_count - collected
        if remaining <= 0:
            continue
        
        # 提取当前块中的该类样本，按需采样
        cls_samples = chunk[chunk['class'] == cls]
        take_num = min(remaining, len(cls_samples))
        sampled_chunks.append(cls_samples.sample(n=take_num, random_state=42))

# 第三步：合并所有采样块
final_df = pd.concat(sampled_chunks, ignore_index=True)
# 验证结果
print(final_df['class'].value_counts())

这个方法的核心是绝不加载全量数据，每一步只处理一小块，完美解决大数据集的内存瓶颈。

方法二：内存中高效分组采样（适合能勉强加载的数据集） #

如果你的数据集能完整加载到内存，但之前用value_counts的方式内存开销大，直接用分组采样更高效——不需要额外存储统计结果，一步完成采样：

# 假设df是已经加载好的数据集
target_count = df['class'].value_counts().min()
final_df = df.groupby('class').apply(lambda group: group.sample(n=target_count, random_state=42)).reset_index(drop=True)

这个方法利用Pandas的分组机制，直接对每个类别进行采样，中间变量极少，比先统计再筛选的方式更省内存。

为什么这两个方法更内存友好？ #

• 分块方法：仅加载数据集的一小部分到内存，避免了全量加载的内存爆炸问题，适合超大型数据集。
• 分组采样：不需要单独存储类别统计结果，直接在分组过程中完成采样，减少了不必要的内存占用。

小提示：如果你的数据集是数据库、HDF5等其他格式，只需要调整读取方式（比如用SQL的随机采样+LIMIT，或者HDF5的分块读取），核心思路都是按需加载，按需采样。

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