numpy 数组的最快保存和加载选项

Fastest save and load options for a numpy array

我有一个生成二维 numpy 数组的脚本,其中 dtype=float 的形状约为 (1e3, 1e6)。现在我正在使用 np.savenp.load 对数组执行 IO 操作。但是,这些函数对每个数组都需要几秒钟的时间。是否有更快的方法来保存和加载整个数组(即不假设其内容并减少它们)?只要数据被准确保留,我愿意在保存之前将数组转换为另一种类型。

对于非常大的数组,我听说过几种解决方案,它们大多是在 I/O 上偷懒:

  • NumPy.memmap,将大数组映射为二进制形式
    • 优点:
      • 除 Numpy 外没有依赖项
      • ndarray 的透明替换(任何 class 接受 ndarray 接受 memmap
    • 缺点:
      • 您的数组块限制为 2.5G
      • 仍然受到 Numpy 吞吐量的限制
  • 对 HDF5 使用 Python 绑定,这是一种大数据就绪的文件格式,如 PyTables or h5py

    • 优点:
      • 格式支持压缩、索引等超棒的功能
      • 显然是最终的 PetaByte 大文件格式
    • 缺点:
      • 分层格式的学习曲线?
      • 必须定义您的性能需求(见下文)
  • Python's pickling系统(退出比赛,提到Pythonicity而不是速度)

    • 优点:
      • 这是 Pythonic ! (哈哈)
      • 支持各种对象
    • 缺点:
      • 可能比其他人慢(因为针对任何对象而不是数组)

Numpy.memmap

来自 NumPy.memmap 的文档:

Create a memory-map to an array stored in a binary file on disk.

Memory-mapped files are used for accessing small segments of large files on disk, without reading the entire file into memory

The memmap object can be used anywhere an ndarray is accepted. Given any memmap fp , isinstance(fp, numpy.ndarray) returns True.


HDF5 阵列

来自h5py doc

Lets you store huge amounts of numerical data, and easily manipulate that data from NumPy. For example, you can slice into multi-terabyte datasets stored on disk, as if they were real NumPy arrays. Thousands of datasets can be stored in a single file, categorized and tagged however you want.

该格式支持以各种方式压缩数据(为相同的 I/O 读取加载更多位),但这意味着数据变得不那么容易单独查询,但在您的情况下(纯粹加载/转储数组)它可能是有效的

这里是与 PyTables 的比较。

由于内存限制,我无法达到 (int(1e3), int(1e6)。 因此,我使用了一个较小的数组:

data = np.random.random((int(1e3), int(1e5)))

NumPy save:

%timeit np.save('array.npy', data)
1 loops, best of 3: 4.26 s per loop

NumPy load:

%timeit data2 = np.load('array.npy')
1 loops, best of 3: 3.43 s per loop

PyTables 写作:

%%timeit
with tables.open_file('array.tbl', 'w') as h5_file:
    h5_file.create_array('/', 'data', data)

1 loops, best of 3: 4.16 s per loop

PyTables 阅读:

 %%timeit
 with tables.open_file('array.tbl', 'r') as h5_file:
      data2 = h5_file.root.data.read()

 1 loops, best of 3: 3.51 s per loop

数字非常相似。所以在这里 PyTables 没有真正的收获。 但是我们已经非常接近我的 SSD 的最大写入和读取速率了。

写作:

Maximum write speed: 241.6 MB/s
PyTables write speed: 183.4 MB/s

阅读:

Maximum read speed: 250.2
PyTables read speed: 217.4

由于数据的随机性,压缩并没有真正帮助:

%%timeit
FILTERS = tables.Filters(complib='blosc', complevel=5)
with tables.open_file('array.tbl', mode='w', filters=FILTERS) as h5_file:
    h5_file.create_carray('/', 'data', obj=data)
1 loops, best of 3: 4.08 s per loop

读取压缩数据变得有点慢:

%%timeit
with tables.open_file('array.tbl', 'r') as h5_file:
    data2 = h5_file.root.data.read()

1 loops, best of 3: 4.01 s per loop

这与常规数据不同:

 reg_data = np.ones((int(1e3), int(1e5)))

写入速度明显加快:

%%timeit
FILTERS = tables.Filters(complib='blosc', complevel=5)
with tables.open_file('array.tbl', mode='w', filters=FILTERS) as h5_file:
    h5_file.create_carray('/', 'reg_data', obj=reg_data)

1 个循环,3 个循环中的最佳:每个循环 849 毫秒

阅读也是如此:

%%timeit
with tables.open_file('array.tbl', 'r') as h5_file:
    reg_data2 = h5_file.root.reg_data.read()

1 loops, best of 3: 1.7 s per loop

结论:你的数据越规则,使用 PyTables 应该越快。

根据我的经验,np.save()&np.load() 是迄今为止在硬盘和内存之间传输数据时最快的解决方案。 在我意识到这个结论之前,我的数据加载非常依赖于数据库和 HDFS 系统。 我的测试表明: 数据库数据加载(从硬盘到内存)带宽可能在50MBps左右(Byets/Second),但是np.load()带宽几乎与我的硬盘最大带宽相同:2GBps(Byets/Second).两个测试环境都使用了最简单的数据结构。

而且我认为用几秒钟来加载形状为 (1e3, 1e6) 的数组不是问题。例如。 你的数组形状是(1000, 1000000),它的数据类型是float128,那么纯数据大小是(128/8)*1000*1,000,000=16,000,000,000=16GBytes 如果需要 4 秒, 那么你的数据加载带宽就是16GBytes/4Seconds = 4GBps。 SATA3最大带宽是600MBps=0.6GBps,你的数据加载带宽已经是它的6倍,你的数据加载性能都快可以媲美DDR's maximum bandwidth,你还想要什么?

所以我最后的结论是:

不要使用python的Pickle,不要使用任何数据库,不要使用任何大数据系统将你的数据存储到硬盘中,如果你可以使用np.save() 和 np.load()。这两个函数是目前硬盘和内存之间传输数据最快的解决方案。

我也测试了 HDF5 ,发现它比 np.load() 和 np.save() 慢得多,所以使用 np.save()& np.load() 如果你的平台有足够的 DDR 内存。

我比较了一些使用 perfplot 的方法(我的一个项目)。以下是结果:

写作

对于大型数组,所有方法的速度都差不多。文件大小也相等,这是可以预料的,因为输入数组是随机双精度数,因此很难压缩。

重现情节的代码:

import perfplot
import pickle
import numpy
import h5py
import tables
import zarr


def npy_write(data):
    numpy.save("npy.npy", data)


def hdf5_write(data):
    f = h5py.File("hdf5.h5", "w")
    f.create_dataset("data", data=data)


def pickle_write(data):
    with open("test.pkl", "wb") as f:
        pickle.dump(data, f)


def pytables_write(data):
    f = tables.open_file("pytables.h5", mode="w")
    gcolumns = f.create_group(f.root, "columns", "data")
    f.create_array(gcolumns, "data", data, "data")
    f.close()


def zarr_write(data):
    zarr.save("out.zarr", data)


perfplot.save(
    "write.png",
    setup=numpy.random.rand,
    kernels=[npy_write, hdf5_write, pickle_write, pytables_write, zarr_write],
    n_range=[2 ** k for k in range(28)],
    xlabel="len(data)",
    equality_check=None,
)

阅读

pickles、pytables 和 hdf5 的速度大致相同; pickles 和 zarr 对于大型阵列来说速度较慢。

重现情节的代码:

import perfplot
import pickle
import numpy
import h5py
import tables
import zarr


def setup(n):
    data = numpy.random.rand(n)
    # write all files
    #
    numpy.save("out.npy", data)
    #
    f = h5py.File("out.h5", "w")
    f.create_dataset("data", data=data)
    f.close()
    #
    with open("test.pkl", "wb") as f:
        pickle.dump(data, f)
    #
    f = tables.open_file("pytables.h5", mode="w")
    gcolumns = f.create_group(f.root, "columns", "data")
    f.create_array(gcolumns, "data", data, "data")
    f.close()
    #
    zarr.save("out.zip", data)


def npy_read(data):
    return numpy.load("out.npy")


def hdf5_read(data):
    f = h5py.File("out.h5", "r")
    out = f["data"][()]
    f.close()
    return out


def pickle_read(data):
    with open("test.pkl", "rb") as f:
        out = pickle.load(f)
    return out


def pytables_read(data):
    f = tables.open_file("pytables.h5", mode="r")
    out = f.root.columns.data[()]
    f.close()
    return out


def zarr_read(data):
    return zarr.load("out.zip")


b = perfplot.bench(
    setup=setup,
    kernels=[
        npy_read,
        hdf5_read,
        pickle_read,
        pytables_read,
        zarr_read,
    ],
    n_range=[2 ** k for k in range(27)],
    xlabel="len(data)",
)
b.save("out2.png")
b.show()

我创建了一个基准测试工具,并使用 python 3.9 生成了各种 loading/saving 方法的基准。我 运行 它在快速 NVMe 上(具有 >6GB/s t运行sfer 速率,因此此处的测量不受磁盘 I/O 限制)。测试的 numpy 数组的大小从微小到 16GB 不等。可以看到结果here。 该工具的 github 存储库是 here

结果略有不同,受数组大小影响;有些方法执行数据压缩,因此需要权衡这些方法。这是 I/O 比率的想法(更多结果来自上面的 link):

图例(扑救): np: np.save(), npz: np.savez(), npzc: np.savez_compressed(), hdf5: h5py.File().create_dataset(), pickle: pickle.dump(), zarr_zip: zarr.save_array() w/ .zip 扩展,zarr_zip:zarr.save_array() w/ .zarr 扩展,pytables:tables.open_file().create_array().

我很惊讶地看到 torch.load 和 torch.save 被认为是最佳的或 near-optimal 根据这里的基准,但我发现它的速度很慢它应该做的事情。所以我试了一下,想出了一个更快的替代方案:fastnumpyio

运行 3x64x64 浮点数组的 100000 save/load 次迭代(计算机视觉中的常见场景)我在 numpy.save 和 numpy.load 上实现了以下加速(我想 numpy.load 这么慢是因为它必须先解析文本数据?):

Windows11,Python3.9.5,Numpy 1.22.0,英特尔酷睿 i7-9750H:

numpy.save: 0:00:01.656569
fast_numpy_save: 0:00:00.398236
numpy.load: 0:00:16.281941
fast_numpy_load: 0:00:00.308100

Ubuntu 20.04,Python 3.9.7,Numpy 1.21.4,英特尔酷睿 i7-9750H:

numpy.save: 0:00:01.887152
fast_numpy_save: 0:00:00.745052
numpy.load: 0:00:16.368871
fast_numpy_load: 0:00:00.381135

macOS 12.0.1,Python3.9.5,Numpy 1.21.2,Apple M1:

numpy.save: 0:00:01.268598
fast_numpy_save: 0:00:00.449448
numpy.load: 0:00:11.303569
fast_numpy_load: 0:00:00.318216

对于更大的数组 (3x512x512),fastnumpyio 的保存速度仍然稍快,加载速度快 2 倍。