Google Colab加载大型.npy文件内存溢出问题求解决方案

问题根源 #

你使用allow_pickle=True加载的npy文件存储的不是纯numpy基础数值类型（比如int/float），而是包含了Python对象（比如字符串、列表、自定义类实例等）。序列化后的800MB文件在反序列化到内存时会大幅膨胀，最终占用内存超过12GB的实例上限导致崩溃；同时因为存在Python对象，numpy的mmap内存映射机制不支持这类数组，所以会抛出对应错误。

可行解决方案 #

方案1：分块加载原始npy文件 #

先在本地设备或高内存的Colab实例上拆分原始文件，再传到Colab分块读取即可：

• 拆分代码（本地/高内存环境运行）：

import numpy as np

# 先完整加载原始文件
full_arr = np.load('oddata.npy', allow_pickle=True)
chunk_size = 10000 # 可根据单块能承受的内存大小调整
total_length = len(full_arr)

for i in range(0, total_length, chunk_size):
    chunk = full_arr[i:i+chunk_size]
    np.save(f'oddata_chunk_{i//chunk_size}.npy', chunk, allow_pickle=True)

• Colab端按需加载对应块，不需要一次性加载全部内容：

import numpy as np
# 示例：加载第0个块
chunk_0 = np.load('oddata_chunk_0.npy', allow_pickle=True)

方案2：转换为支持混合类型的列式存储格式 #

如果需要频繁随机访问不同部分的数据，推荐转成Parquet格式，该格式支持懒加载、按需查询，不会一次性加载全部内容到内存：

• 格式转换代码（本地/高内存环境运行）：

import numpy as np
import pandas as pd

full_arr = np.load('oddata.npy', allow_pickle=True)
# 先转成pandas DataFrame适配Parquet的结构化存储
df = pd.DataFrame(full_arr)
# 存储为Parquet，压缩率高，支持按行/按列按需读取
df.to_parquet('oddata.parquet')

• Colab端按需读取：

import pandas as pd
# 可以只读取指定列/指定行范围，不会加载全量数据
# 示例：读取前1000行的指定列
part_df = pd.read_parquet('oddata.parquet', columns=['col1', 'col2'], filters=[('index', '<', 1000)])
# 转成numpy数组使用
part_arr = part_df.to_numpy()

方案3：升级Colab运行时配置 #

如果不想修改文件格式或者拆分，也可以切换到高内存运行时：在Colab界面点击「修改」->「笔记本设置」->「运行时类型」，选择高内存选项，高内存基础版实例内存为25GB，足够承载反序列化后的数组。

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