迭代大型 npz 文件的有效方法

Question

我有 50000+ 个 npz 文件（文件大小 15MB），每个文件都有三个 numpy 数组 A、B、C。

A 是一个形状为 (33000,) 的 numpy 数组； B 是一个形状为 (33000,) 的 numpy 数组，C 是一个形状为 (33000,224,224) 的 uint8 类型的位数组。

npz_file_dict = np.load(npz_file_path)
num_records = len(npz_file_dict['A'])

A = npz_file_dict['A']
B = npz_file_dict['B']
C = npz_file_dict['C']


def get_items(num_records):
   for i in range(num_records)
      yield A[i], B[i], C[i]
for a,b,c in get_items(num_records):
   # do some work with (a,b,c)
   # store the result in a json format with {'a' : a, 'b' : b, 'c' : c}

for 循环需要很长时间才能完成（对于 30000 条记录，有时需要超过 35 分钟）。有没有一种有效的方法来解析记录？

编辑：我尝试做内存映射的 numpy 数组，但它似乎并没有提高性能，因为它们只对 npy 文件有影响。

编辑：需要对我提出的方法进行批评。

我已经求助于使用内存映射来加速我的迭代，对于我读取的每个文件，我都会花费一些时间（在我所谓的设置阶段）将这个巨大的 numpy 数组映射到文件指针。

num_records = len(npz_file_dict['np_cids'])
print(num_records)


start_memmap = time()
data = npz_file_dict['C']

print('Time to read data {}'.format(time() - start_memmap))


filename = path.join(mkdtemp(), 'newfile.dat')

print('Path to file',  filename)
fp = np.memmap(filename, dtype='float32', mode='w+', shape=(num_records, 224,224))
fp[:] = data[:]

finish_memmap = time()

print('Time to create memmap {}'.format(finish_memmap - start_memmap))

完成此设置后，我只需在 for 循环中遍历 A 和 B 数组并使用 fp 获取相应的 i^th 记录。这种方法的优点是迭代非常快。缺点是设置时间介于单进程模式下的 50 秒和多进程模式下的 100 秒之间。

如果您将设置时间摊销到 30000 条记录的集合中，则大约是 3.3 毫秒。在迭代过程中，读取速度非常快——几百微秒。这使得每条记录所花费的总时间大约为 5 毫秒。与我之前的方法（每条记录需要 150 毫秒）相比，这是 96% 的速度。

Answer 1

我认为你可以清理代码：

npz_file_dict = np.load(npz_file_path)
A = npz_file_dict['A']   # (33000,) shape array
B = npz_file_dict['B']
C = npz_file_dict['C']

for a,b,c in zip(A,B,C):
   # do some work with (a,b,c)
   # store the result in a json format with {'a' : a, 'b' : b, 'c' : c}

这清楚地表明您正在使用从文件中加载的数组。我不知道它是否会加快速度；这取决于npz 加载器是否使用某种具有重复npz_file_dict['A] 访问的缓存。

但无论它们是如何创建或加载的，通过 numpy 数组的迭代都很慢 - 比对它们使用编译的“整个数组”操作要慢得多。列表的迭代速度更快。事实上，这样做可能会有所帮助

for a,b,c in zip(A.tolist(), B.tolist(), list(C):
    ...

A.tolist() 是一种将数组转换为列表的快速方法。我不在C 上使用它，因为它是3d 的，我假设您想将C[i,:,:] 用作数组，而不是嵌套列表。虽然由于您正在将c 写入json，但您可能希望它是一个列表列表。

===

迭代memmap 数组比迭代内存中的数组要慢很多：

In [1]: C=np.ones((3000,224,224))                                               
In [2]: np.savez('bigC.npz',C=C)                                                
In [3]: fp = np.memmap('bigC.dat', dtype=C.dtype, mode='w+', shape=C.shape)     
In [4]: fp[:] = C[:]                                                            
In [5]: %%timeit  
   ...: for i in range(3000): 
   ...:     c = C[i] 
   ...:                                                                         
566 µs ± 2.55 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [7]: %%timeit  
   ...: for i in range(3000): 
   ...:     c = fp[i] 
   ...:      
   ...:                                                                         
9.74 ms ± 94.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

迭代 npz 负载确实很慢 - 我的测试仍在运行：

In [8]: %%timeit d = np.load('bigC.npz') 
   ...: for i in range(3000): 
   ...:     c = d['C'][i] 

使用一个初始负载进行测试：

In [243]: d = np.load('bigC.npz')                                               
In [244]: %%timeit  
     ...: D = d['C'] 
     ...: for i in range(3000): 
     ...:     c = D[i] 
     ...:                                                                       
2.17 s ± 27.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [245]: %%timeit  
     ...: D = d['C'] 
     ...:  
     ...:                                                                       
2.14 s ± 6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

因此，从npz 一次性加载大型数组是可以避免的，但您当然不想为每一行重复加载它。如果数组太大而无法放入内存，则将其转储到 memmap 可能会有所帮助，这并不是对内存数组的改进。