使用zarr处理训练大型数据集

Zarr是一种存储分组、压缩的N维数组的格式。

zarr官方教程：https://zarr.readthedocs.io/en/stable/tutorial.html

---

对于处理大型数据集，np.array往往是力不从心的，因此需要使用磁盘到内存映射的技术（如memmap、zarr等）。

• 对于一整个较大的.npy或.dat文件，可以直接使用np.memmap读取。
• 而对于大量文件的情况下（如38万个1M小文件），想全部读入内存，并且生成一个array数组，建议使用zarr。

速度比较

• 对于适应内存的小型数据集，速度性能上：numpy>zarr>memmap
• 对于超出内存的大型数据集，速度性能上：zarr>memmap

zarr使用方法

1、创建一个空的zarr数组

import zarr

path = '.'
zarr_arr = zarr.open(f'{path}/zarr_data.zarr',
                     mode='w',
                     shape=(10000, 10, 100),
                     chunks=(100, -1, -1),  # chunks=(100, -1, -1) == (100, None, None)
                     dtype='float32')

创建zarr数组时，我们需要使用mode=‘a’（如果文件存在，则读取，否则就写入）或mode=‘w’（始终写入或覆盖）。

zarr的核心关键就是分块存储映射，块的大小可以通过参数chunks来设置。

• chunks=(100, -1, -1)表示在第一个维度上分块，每个块大小为100，而在第二第三维度上不分块。相当于每个小块的shape=(100, 10, 100)。
• 若设置chunks=False，则表示不分块，直接整个数组都载入内存
• 若不设置chunks，zarr会默认使用特定的优化算法来设置块大小，但不一定是最优的。
• chunks的大小会影响zarr的速度。

2、把你的大型数据放入zarr数组中去

zarr_arr是一个类似np.array的对象，因此可以直接像使用np.array那样使用zarr_arr。

下面举一个小例子：

import numpy as np
from tqdm import tqdm

for i in tqdm(range(10000)):
    zarr_arr[i,:,:]=np.random.rand(10,100)

3、读取.zarr文件时注意事项

如果只是读取不修改，最好使用mode='r'模式读取。如下：

import zarr

path = '.'
zarr_arr = zarr.open(f'{path}/train_data.zarr', mode='r')

zarr使用注意事项

zarr不需要close()。每次使用数据会自动刷新到磁盘，并且每当修改数组时都会自动关闭文件。

zarr特性

• zarr的核心就是分块从磁盘加载到内存

• zarr数组是压缩过的，所以占的内存小

• 可以使用多线程