torch.utils.data.DataLoader中collate_fn

在PyTorch中，torch.utils.data.DataLoader中的collate_fn参数用于指定如何将一个batch的数据样本整合成一个batch的张量。默认情况下，collate_fn使用torch.stack函数将数据样本堆叠在一起。如果数据样本具有不同的大小，则需要自定义collate_fn函数来处理。

例如，如果数据样本是一个元组，其中第一个元素是图像张量，第二个元素是标签张量，则可以使用以下自定义collate_fn函数：

def custom_collate_fn(batch):
    images = []
    labels = []
    for image, label in batch:
        images.append(image)
        labels.append(label)
    images = torch.stack(images, dim=0)
    labels = torch.tensor(labels)
    return images, labels