内存放不下？使用PyTables存储大型numpy数组

最近需要提取100万张图像的深度特征，按照特征是4096维 float 格式计算，至少需要 $\frac{1000000\ast 4096\ast 32}{8\ast 1024^3}\approx 15GB$ 的存储空间，内存存放是不太合理的。当我们需要处理这类内存存放不下的 numpy 数组时，分成多个文件存储在硬盘上未免有些麻烦。PyTables 基于 HDF5，PyTables 的 EArray (Extendable Array) 可以让我们在硬盘上用单个文件压缩存储大型 numpy 数组的同时，像存放在内存中一样访问它们。

安装 PyTables

pip install tables 即可。

更多方式详见[官方文档]。

创建 EArray 并添加数据

import tables
import numpy as np

# 生成随机数据
data1 = np.random.random((1000, 4096))

hdf5_path = "test_data.hdf5"
# 和普通文件操作类似，'w'、'r'、'a' 分别表示写数据、读数据、追加数据
hdf5_file = tables.open_file(hdf5_path, mode='w')
# 设定压缩级别和压缩方法
filters = tables.Filters(complevel=5, complib='blosc')
earray = hdf5_file.create_earray(
	hdf5_file.root,
	'data', # 数据名称，之后需要通过它来访问数据
	tables.Atom.from_dtype(data1.dtype), # 设定数据格式（和data1格式相同）
	shape=(0, 4096), # 第一维的 0 表示数据可沿行扩展
	filters=filters,
	expectedrows=1000000 # 完整数据大约规模，可以帮助程序提高时空利用效率
)
# 将 data1 添加进 earray
earray.append(data1)
# 写完之后记得关闭文件
hdf5_file.close()

读取 EArray 数据

hdf5_path = "test_data.hdf5"
hdf5_file = tables.open_file(hdf5_path, mode='r')
# 数据名称为 'data'，我们可以通过 .root.data 访问到它
hdf5_data = hdf5_file.root.data
print(hdf5_data.shape) # (1000, 4096)
# 像在内存中一样自由读取数据切片！
print(hdf5_data[:10])
hdf5_file.close()

向先前创建的 EArray 中添加数据

# 再生成一组随机数据
data2 = np.random.random((1000, 4096))

hdf5_path = "test_data.hdf5"
# 注意这里 mode 应为 'a'
hdf5_file = tables.open_file(hdf5_path, mode='a')
hdf5_data = hdf5_file.root.data
print(hdf5_data.shape) # (1000, 4096)
hdf5_data.append(data2)
hdf5_file.close()

# 重新以 'r' 方式打开文件，查看数据规模是否改变
hdf5_file = tables.open_file(hdf5_path, mode='r')
print(hdf5_file.root.data.shape) # (2000, 4096)
hdf5_file.close()

Conlusion

有了 PyTables，我们就可以将小规模的 numpy 数组不断 append 进 EArray，需要取用的时候只要像全部数据都在内存中一样，取出内存可以负荷大小的数据切片进行处理即可。