Jupyter Notebook执行train_test_split时触发内存错误求助

我之前也碰到过类似的坑——明明数据集文件才2.5MB，加载后操作却爆内存，咱们一步步拆解排查：

1. 先搞清楚数据在内存里的真实占用 #

文件大小和内存加载后的大小完全是两回事！比如：

• 如果你的X是Pandas DataFrame，大量object类型的列会占用巨量内存（每个元素都是指针）；
• 默认用float64/int64这类大数值类型，其实很多场景下换成float32/int32甚至int8就足够；
• 要是X原本是稀疏矩阵，不小心被转成了稠密数组，内存会直接暴涨数倍。

你可以先运行这行代码查看真实内存占用：

X.info(memory_usage='deep')

如果发现内存远超预期，先做数据类型优化：

• 把低基数的object列转成category类型：X['col_name'] = X['col_name'].astype('category')
• 缩小数值类型：X = X.astype('float32')（精度足够的前提下）

2. 清空Jupyter残留的内存占用 #

Jupyter会保留之前运行的所有变量，哪怕你重新运行单元格，旧变量可能还悄悄占着内存。可以先执行：

%reset -f

清空所有变量后，重新加载数据再运行train_test_split试试。

3. 避开train_test_split的潜在内存开销 #

有时候train_test_split内部的shuffle操作会临时占用额外内存，你可以手动实现拆分来测试：

import numpy as np

# 生成打乱的索引
shuffle_idx = np.random.permutation(len(X))
train_size = int(len(X) * 0.9)

# 手动拆分（兼容DataFrame和numpy数组）
X_train = X.iloc[shuffle_idx[:train_size]] if isinstance(X, pd.DataFrame) else X[shuffle_idx[:train_size]]
y_train = y.iloc[shuffle_idx[:train_size]] if isinstance(y, pd.Series) else y[shuffle_idx[:train_size]]
X_test = X.iloc[shuffle_idx[train_size:]] if isinstance(X, pd.DataFrame) else X[shuffle_idx[train_size:]]
y_test = y.iloc[shuffle_idx[train_size:]] if isinstance(y, pd.Series) else y[shuffle_idx[train_size:]]

如果手动拆分没问题，那大概率是train_test_split的内部处理导致的内存问题，你可以尝试更新scikit-learn：pip install --upgrade scikit-learn

4. 增大swap空间缓解内存压力 #

从你的free -m输出看，swap只有255MB，太小了。当物理内存不够时，系统会用swap当虚拟内存，但swap不足会直接触发内存错误。你可以临时创建swap文件：

# 创建1GB的swap文件
sudo fallocate -l 1G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile

用完后可以关闭swap：sudo swapoff /swapfile && sudo rm /swapfile

5. 排查隐性内存泄漏 #

比如你之前运行的单元格有没有加载过其他大变量？或者Jupyter内核本身存在内存泄漏？可以重启Jupyter内核后再重新运行代码，排除内核残留的问题。

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