【深度学习2】DataLoader的collate_fn参数

参考文章：https://blog.csdn.net/weixin_42028364/article/details/81675021

collate_fn是torch.utils.data.DataLoaderd的参数之一，用来定义如何取样，可以不赋值，也可以自己定义函数来实现自己想要的功能。

比如：

def collate_fn(data):
    images,captions, sentence_num, max_word_num, image_id = zip(*data)
    images = torch.stack(images, 0)
    max_sentence_num = max(sentence_num) 
    max_word_num = max(max_word_num)
    prob = np.zeros((len(captions), max_sentence_num + 1))
    for i, caption in enumerate(captions):
        for j, sentence in enumerate(caption):
            prob[i][j] = len(sentence) > 0
    return images, targets, prob, image_id



data_loader = torch.utils.data.DataLoader(dataset=dataset,
                                          batch_size=batch_size,
                                          shuffle=shuffle,
                                          collate_fn=collate_fn, 
                                          num_workers=num_workers)