记一次 Pandas 内存溢出的问题

最近在使用 Jupyter Notebook 进行数据科学任务的时候，发现经常跑崩，明明数据可以读取，计算，可是跑着跑着就内核就崩溃了

我的设备是 MacBook Pro M3 + 16G 内存，理论上不应该在大约 500MB 的数据上翻车

在检查代码、重启等操作都不管用以及怀疑是不是设备比较垃圾的同时，进一步研究发现，原来 Mac 真的不背这个锅

排查思路

因为我的开发环境是在 Docker 中构建的，所以肯定是先 docker stats 监控下崩溃过程中的状态

很明显，内存快要兜不住了，那就检查一下运行崩溃的变量的内存占用

print(df.info(memory_usage='deep'))

发现仅仅这一个变量就占用了 4 个 G

<class 'pandas.core.frame.DataFrame'>
Int64Index: 316793 entries, 0 to 116792
Columns: 1654 entries, 编号 to y
dtypes: float64(1651), int64(3)
memory usage: 3.9 GB
None

再用下面的代码检查当前环境中的全部变量及大小

import sys

local_vars = list(globals().items())
memory_info = []

for var_name, var_val in local_vars:
    memory_info.append((var_name, sys.getsizeof(var_val)))

mem_df = pd.DataFrame(memory_info, columns=['Variable', 'Memory Usage (Bytes)'])
mem_df = mem_df.sort_values(by='Memory Usage (Bytes)', ascending=False).reset_index(drop=True)

print(mem_df)

可以看到仅仅三个变量，就占用了约 7 个 G 的内存，所有存储在内存中的变量，无论是否在当前计算中使用，都会占用内存。

        Variable  Memory Usage (Bytes)
         data            4194339336
  data_sample            1603200144
    data_test             936212832
       mem_df                 16183
         _i14                  2037
..           ...                   ...
        cs1                    48
RANDOM_SEED                    28
   __spec__                    16
 __loader__                    16
__package__                    16

这些 _7 和 _5 等类似的符号是 Jupyter Notebook 或 IPython 环境中的自动变量，它们表示之前执行单元的输出结果，并且它们也会占用内存。

而我后面又会基于这个变量，通过合并、计算等操作衍生其他的变量，过不了几个步骤内存就飙到了十几个 G

难怪到最后可能只是执行一个数据的抽取，内核就会挂掉，所以确实是数据量比较大导致的，那要怎么解决呢？

解决方案

既然内存兜不住了，那解决方案也就很自然的是如何优化，我搜了一下，大概有下面几种方案

清理不必要的变量

比较自然的方式就在原始/中间变量计算完之后，就删除并释放内存

del data

# 强制垃圾回收
import gc
gc.collect()

通过使用 del 删除不再需要的中间变量，并调用 gc.collect() 强制进行垃圾回收，可以有效释放内存。

Tips

del 语句用于删除对象的引用，但它并不会立即释放内存。当对象的引用计数降为零时，Python 的垃圾回收器会回收这些对象并释放相应的内存。

然而，垃圾回收器不一定会立即运行，因此在处理大规模数据时，即使使用 del 删除了变量，内存也可能不会立即释放，所以我们还需要手动释放

gc.collect() 是手动触发垃圾回收的一个方法，它可以强制 Python 的垃圾回收器运行并回收所有可以回收的对象，释放内存。这在处理大规模数据时尤其有用，因为它可以帮助及时释放内存，避免内存不足的问题。

封装在函数中

或者将部分代码封装在函数中，这样在函数执行完毕后就会自动释放局部变量：

def process_data():
    # 这里定义你的数据处理逻辑

    return processed_data

result = process_data()

这样仅在需要时调用函数，结果存储在一个新变量中，但是要注意在执行函数的过程中，这些变量仍然会占用内存，所以函数中间如果变量占用过大，可能需要手动删除变量、释放内存

修改变量类型

减少数据类型的内存占用也是一个不错的办法，将数据类型转换为更节省内存的类型

df = df.astype({
    'column1': 'float32',
    'column2': 'int32',
})

这个也是网上比较推荐的一种方案，但是我的数据几千列，虽然可以批量操作，但是还是感觉有点复杂

分批次处理数据

虑分批次处理数据，避免一次性加载所有数据，也是网上比较主流的办法，下面就是分批读取数据的函数

def read_csv_in_chunks(file, encoding='gbk', chunksize=10000):
    chunks = pd.read_csv(file, encoding=encoding, chunksize=chunksize)
    return pd.concat(chunk for chunk in chunks)

但对我来说，个人不是很喜欢，感觉让整个过程变复杂了

使用更高效的库

对于特别大的数据集，可以考虑使用如 Dask 这样的库，这些库可以处理大规模数据而不需要一次性加载到内存中：

import dask.dataframe as dd
ddf = dd.read_csv('large_dataset.csv')

ddf = ddf.merge(xxx)

ddf.to_csv('processed_large_dataset.csv', single_file=True)

但是考虑到后面的代码在其他机器上的复现，多安装一个库可能不是我的选择，感兴趣的可以尝试一下，还有一个 Vaex，貌似也可以实现。

小结

没什么好总结的，踩了一个坑，以后知道了，挺好的(^▽^ )

记一次 Pandas 内存溢出的问题

排查思路​

解决方案​

清理不必要的变量​

封装在函数中​

修改变量类型​

分批次处理数据​

使用更高效的库​

小结​