在带有 pytorch 数据加载器的开始和结束的 zarr 数组上创建生成器

Question

我正在做一个 pytorch 项目，我的数据保存在 zarr。

zarr 上的随机访问代价高昂，但由于 zarr 使用块缓存，迭代非常快。为了利用这个事实，我将 IterableDataset 与多个 worker 一起使用：

class Data(IterableDataset):
    def __init__(self, path, start=None, end=None):
        super(Data, self).__init__()
        store = zarr.DirectoryStore(path)
        self.array = zarr.open(store, mode='r')

        if start is None:
            start = 0
        if end is None:
            end = self.array.shape[0]

        assert end > start

        self.start = start
        self.end = end

    def __iter__(self):
        return islice(self.array, self.start, self.end)

问题是 self.array 有大约 10e9 行和连续的工作人员，因为 self.start 和 self.end 自然变大，创建生成器像itertools.islice(array, start, end) 在我的 training/validation 过程中占用了大量时间，因为 islice 仍然必须遍历不需要的元素，直到到达 start。一旦为每个工人创建了一个生成器，这就像一个魅力，但到达那里需要太长时间。

有没有更好的方法来创建这样的生成器？或者也许有更聪明的方法在 pytorch 中使用 zarr?

Answer 1

更新：截至 2021 年 3 月，此更改已合并到 zarr。

我对 zarr 进行了深入研究，看起来这最容易从 zarr 内部启用。我已经打开了一个实现功能 array.islice(start, end).

的问题 here, in the meantime I made a fork of zarr

数据集 __iter__ 方法如下所示：

def __iter__(self):
    return self.array.islice(self.start, self.end)

在带有 pytorch 数据加载器的开始和结束的 zarr 数组上创建生成器

Creating a generator over a zarr array with start and end for pytorch dataloader

python

arrays

pytorch

zarr