代码之家 › 专栏 › 技术社区 › Bram Vanroy

在PyTorch的数据加载器中使用带有酸洗数据的生成器

pytorch pickle generator dataset python

Bram Vanroy · 技术社区 · 6 年前

我以前做过一些预处理和特性选择,我有一个pickle培训输入数据,它由列表列表组成,例如(但是pickle)

[[1,5,45,13], [23,256,4,2], [1,12,88,78], [-1]]
[[12,45,77,325], [23,257,5,28], [3,7,48,178], [12,77,89,99]]
[[13,22,78,89], [12,33,97], [-1], [-1]]

[-1] 是个填充标记,但我不认为这很重要。

因为这个文件有很多千兆字节,所以我希望有足够的内存,并使用一个生成器在pickle中逐行读取(按列表列出)。我已经找到 this answer 这可能会有所帮助。如下所示:

def yield_from_pickle(pfin):
    with open(pfin, 'rb') as fhin:
        while True:
            try:
                yield pickle.load(fhin)
            except EOFError:
                break

接下来,我希望在pytorch(1.0.1)中使用此数据 Dataloader . 根据我在其他答案中的发现,我必须给它一个 Dataset 它可以是子集,但必须包含 __len__ 和 __getitem__ . 它可以是这样的:

class TextDataset(Dataset):
    def __init__(self, pfin):
        self.pfin = pfin

    def __len__(self):
        # memory-lenient way but exhaust generator?
        return sum(1 for _ in self.yield_from_pickle())

    def __getitem__(self, index):
        # ???
        pass

    def yield_from_pickle(self):
        with open(self.pfin, 'rb') as fhin:
            while True:
                try:
                    yield pickle.load(fhin)
                except EOFError:
                    break

但我一点也不确定这是否可能。我如何实施 爱伦尼亚 和 第二章 以合理的方式?我不认为我在做什么 爱伦尼亚 是个好主意,因为那样会耗尽发电机的能量,而我完全不知道如何安全地实施 第二章 同时固定发电机。

有更好的方法吗?总的来说,我想构建一个数据集,它可以被输入到PyTorch的数据加载器(因为它的多处理能力),但是以一种内存高效的方式,我不需要将整个文件读取到内存中。