将异构 numpy 数组写入二进制文件

Question

我有大量 n 的 3x3 矩阵、长度为 3 的向量以及需要以给定二进制格式写入文件的整数。我可以很容易地使用for 循环到fh.write() 一个接一个的项目，但这很慢。另一种方法是将数据复制到具有特殊 dtype 的数组中。这要快得多，但会在内存中创建一个令人望而却步的大副本：

import numpy as np

n = 100  # a large number
A = np.random.rand(n, 3, 3)
b = np.random.rand(n, 3)
c = np.ones(n, dtype=int)

# slow
with open("out.dat", "wb") as fh:
    for a_, b_, c_ in zip(A, b, c):
        fh.write(a_)
        fh.write(b_)
        fh.write(c_)

# memory-consuming
dtype = np.dtype([
  ('A', ('<f', (3, 3))),
  ('b', ('<f', 3)),
  ('c', '<H'),
])
data = np.empty(n, dtype=dtype)
data["A"] = A
data["b"] = b
data["c"] = c
with open("out.dat", "wb") as fh:
    data.tofile(fh)

这里有没有快速、节省内存的替代方案？

Answer 1

逐块写入文件

请考虑您的第一个和第二个版本会导致不同的结果。我将在这里重点介绍第二个版本。与逐块写入相比，此版本确实不仅内存开销很大，而且比将进程拆分为多个块要慢。

示例

def write_method_2(file_name,A,b,c):
    n=A.shape[0]

    dtype = np.dtype([
      ('A', ('<f', (3, 3))),
      ('b', ('<f', 3)),
      ('c', '<H'),
    ])

    data = np.empty(n, dtype=dtype)
    data["A"] = A
    data["b"] = b
    data["c"] = c
    with open(file_name, "wb") as fh:
        data.tofile(fh)

唯一的缺点是代码较长...使用生成器函数也应该可以将其推广到多个 IO 操作。

def write_method_3(file_name,A,b,c):
    n=A.shape[0]
    blk_size=10_000

    dtype = np.dtype([
      ('A', ('<f', (3, 3))),
      ('b', ('<f', 3)),
      ('c', '<H'),
    ])

    data = np.empty(blk_size, dtype=dtype)
    with open(file_name, "wb") as fh:
        #write block-wise
        n_full_blocks=n//blk_size
        for i in range(n_full_blocks):
            data["A"] = A[i*blk_size:i*blk_size+blk_size]
            data["b"] = b[i*blk_size:i*blk_size+blk_size]
            data["c"] = c[i*blk_size:i*blk_size+blk_size]
            data.tofile(fh)
        #write remainder
        n_full_blocks=n//blk_size
        data=data[:n-n_full_blocks*blk_size]
        data["A"] = A[n_full_blocks*blk_size:]
        data["b"] = b[n_full_blocks*blk_size:]
        data["c"] = c[n_full_blocks*blk_size:]
        data.tofile(fh)

编辑

这是一种使用非简单数据类型将数据从多个 nd 数组写入文件的更通用方法。

def write_method_3_gen(fh,dtype,tuple_of_arr,blk_size=500_000):
    """
    fh             file-handle
    dtype          some non-simple dtype
    tuple_of_arr   tuple of arrays
    blk_size       size of a block, default 0.5MB
    """
    n=tuple_of_arr[0].shape[0]
    blk_size=blk_size//dtype.itemsize
    data = np.empty(blk_size, dtype=dtype)

    #write block-wise
    n_full_blocks=n//blk_size
    for i in range(n_full_blocks):
        for j in range(len(tuple_of_arr)):
            data[keys[j]] = tuple_of_arr[j][i*blk_size:i*blk_size+blk_size]
        data.tofile(fh)

    #write remainder
    n_full_blocks=n//blk_size
    data=data[:n-n_full_blocks*blk_size]
    for j in range(len(tuple_of_arr)):
        data[keys[j]] = tuple_of_arr[j][n_full_blocks*blk_size:]
    data.tofile(fh)

时间安排

import numpy as np
import time

n = 10_000_000  # a large number
A = np.random.rand(n, 3, 3)
b = np.random.rand(n, 3)
c = np.ones(n, dtype=int)

t1=time.time()
write_method_2("out_2.dat",A,b,c)
print(time.time()-t1)
#3.7440097332000732

#with blk_size=10_000 this has only 0.5MB memory overhead, 
#which stays constant, even on much larger examples
t1=time.time()
write_method_3("out_3.dat",A,b,c)
print(time.time()-t1)
#0.8538124561309814