用于各种维度输入的高效 PyTorch DataLoader collate_fn 函数

Posted 2023-03-12

技术标签:

【中文标题】用于各种维度输入的高效 PyTorch DataLoader collate_fn 函数【英文标题】：Efficient PyTorch DataLoader collate_fn function for inputs of various dimensions 【发布时间】：2019-06-02 15:31:28 【问题描述】：

我在为 PyTorch DataLoader 类编写自定义 collate_fn 函数时遇到问题。我需要自定义函数，因为我的输入具有不同的维度。

我目前正在尝试编写Stanford MURA paper 的基线实现。数据集有一组带标签的研究。一项研究可能包含多个图像。我创建了一个自定义 Datasetclass，它使用 torch.stack 堆叠这些多个图像。

然后将堆叠张量作为输入提供给模型，并对输出列表进行平均以获得单个输出。当batch_size=1 时，此实现适用于DataLoader。但是，当我尝试将 batch_size 设置为 8 时，就像原始论文中的情况一样，DataLoader 失败，因为它使用 torch.stack 来堆叠批次并且我的批次中的输入具有可变尺寸（因为每个研究可以有多个图像）。

为了解决这个问题，我尝试实现我的自定义 collate_fn 函数。

def collate_fn(batch):
    imgs = [item['images'] for item in batch]
    targets = [item['label'] for item in batch]
    targets = torch.LongTensor(targets)
    return imgs, targets

然后在我的训练周期循环中，我像这样循环每个批次：

for image, label in zip(*batch):
    label = label.type(torch.FloatTensor)
    # wrap them in Variable
    image = Variable(image).cuda()  
    label = Variable(label).cuda()
    # forward
    output = model(image)
    output = torch.mean(output)
    loss = criterion(output, label, phase)

但是，这并没有给我任何改进的时代时间，并且仍然需要与只有 1 的批量大小一样长。我还尝试将批量大小设置为 32，但这并没有改善时间要么。

我做错了吗？有更好的方法吗？

【问题讨论】：

“这并没有给我任何改进的时间”为什么它应该更快？ @Ivan 我希望增加批次大小以缩短总训练时间，因为批次中将处理更多样本。 【参考方案1】：

非常有趣的问题！如果我对您的理解正确（并且还检查了论文摘要），您有来自 14,863 项研究的 40,561 张图像，其中每项研究都由放射科医生手动标记为正常或异常。

我相信您遇到问题的原因是，例如，您创建了一个堆栈，

研究 A - 12 张图片研究 B - 13 张图片研究 C - 7 张图片研究 D - 1 张图片等

并且您尝试在训练期间使用 8 的批大小，这在学习 D 时会失败。

因此，我们是否有理由要对研究中的输出列表进行平均以适合单个标签？否则，我将简单地收集所有 40,561 张图像，为来自同一研究的所有图像分配相同的标签（例如将 A 中的输出列表与 12 个标签的列表进行比较）。

因此，使用单个数据加载器，您可以在研究中随机播放（如果需要）并在训练期间使用所需的批量大小。

我看到这个问题已经存在了一段时间，我希望它对未来的人有所帮助:)

【讨论】：