PyTorch：我可以按长度对批次进行分组吗？

Question

我正在从事一个 ASR 项目，我在其中使用了来自 HuggingFace (wav2vec2) 的模型。我现在的目标是将训练过程转移到 PyTorch，因此我正在尝试重新创建 HuggingFace 的 Trainer() class 提供的所有内容。

这些实用程序之一是能够按长度对批次进行分组并将其与动态填充相结合（通过数据整理器）。然而老实说，我什至不确定如何在 PyTorch 中开始这个。

在我的案例中，输入是代表 .wav 文件原始波形的一维数组。所以在训练之前，我需要确保将相似大小的数组一起批处理。我是否需要创建一个自定义 Dataloader class 并更改它，以便每次它给我的批处理大小都尽可能接近长度？

我的一个想法是以某种方式将数据从最短到最长（或相反）排序，并且每次都从中提取 batch_size 个样本。这样，第一批将包含最大长度的样本，第二批将具有第二大长度，依此类推。

不过，我不确定如何处理这个实现。任何建议将不胜感激。

提前致谢。

Answer 1

解决此问题的一种可能方法是使用 批处理采样器 并为您的数据加载器实现 collate_fn，它将对您的批处理元素执行动态填充。

使用这个基本数据集：

class DS(Dataset):
    def __init__(self, files):
        super().__init__()
        self.len = len(files)
        self.files = files

    def __getitem__(self, index):
        return self.files[index]

    def __len__(self):
        return self.len

用一些随机数据初始化：

>>> file_len = np.random.randint(0, 100, (16*6))
>>> files = [np.random.rand(s) for s in file_len]
>>> ds = DS(files)

首先定义您的批次采样器，这本质上是一个可迭代的返回批次的索引，数据加载器将使用这些批次从数据集中检索元素。正如您所解释的，我们可以只对长度进行排序并根据这种排序构造不同的批次：

>>> batch_size = 16
>>> batches = np.split(file_len.argsort()[::-1], batch_size)

我们应该有长度彼此接近的元素。

我们可以对 assemble 批处理元素实施 collate_fn 函数并集成动态填充。这基本上是在数据集和数据加载器之间放置一个额外的用户定义层。目标是找到批次中最长的元素，并用正确数量的 0s:

填充所有其他元素

def collate_fn(batch):
    longest = max([len(x) for x in batch])
    s = np.stack([np.pad(x, (0, longest - len(x))) for x in batch])
    return torch.from_numpy(s)

然后你可以初始化一个数据加载器：

>>> dl = DataLoader(dataset=ds, batch_sampler=batches, collate_fn=collate_fn)

并尝试迭代，如您所见，我们得到了长度递减的批次：

>>> for x in dl:
...   print(x.shape)
torch.Size([6, 99])
torch.Size([6, 93])
torch.Size([6, 83])
torch.Size([6, 76])
torch.Size([6, 71])
torch.Size([6, 66])
torch.Size([6, 57])
...

虽然这种方法有一些缺陷，例如，元素的分布总是相同的。这意味着您将始终以相同的出现顺序获得相同的批次。这是因为此方法基于数据集中元素的长度排序，批次的创建没有可变性。您可以通过打乱批次来减少这种影响（例如，通过将 batches 包装在 RandomSampler 中）。然而，正如我所说，批次的内容在整个培训过程中将保持不变，这可能会导致一些问题。

请注意，在您的数据加载器中使用 batch_sampler 是互斥选项 batch_size、shuffle 和 sampler!

PyTorch：我可以按长度对批次进行分组吗？

PyTorch: Can I group batches by length?

pytorch

pytorch-dataloader

huggingface-datasets