在 zarr 数组上创建一个生成器，并为 pytorch 数据加载器提供开始和结束

Question

我正在开发一个 pytorch 项目，我的数据保存在zarr.

随机访问的zarr成本很高，但由于zarr使用了块式缓存，迭代速度非常快。为了利用这一事实，我与多个工作人员一起使用IterableDataset：

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会变得更大，创建生成器会在我的训练/验证过程中花费大量时间，因为仍然需要迭代不需要的元素，直到它到达。一旦为每个工作人员创建了一个生成器，这就像一个魅力，但实现这一目标需要很长时间。10e9self.startself.enditertools.islice(array, start, end)islicestart

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

Answer 1

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

我对 zarr 进行了一些深入研究，看起来这很容易从 zarr 内部启用。我在这里提出了一个问题，同时我制作了一个zarr 的分支来实现该功能array.islice(start, end)。

数据集__iter__方法如下所示：

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