小编Use*_*non的帖子

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

我目前正在尝试编写斯坦福 MURA 论文的基线实现。该数据集有一组标记的研究。一项研究可能包含多个图像。我创建了一个自定义Dataset类，使用torch.stack.

然后将堆叠张量作为输入提供给模型，并对输出列表进行平均以获得单个输出。此实现适用于DataLoaderwhen batch_size=1。但是，当我尝试将 设置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 …