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

最新推荐文章于 2024-06-19 09:12:04 发布

iviyandyssg

最新推荐文章于 2024-06-19 09:12:04 发布

阅读量718

点赞数

分类专栏：记录文章标签：深度学习 pytorch

本文链接：https://blog.csdn.net/iviyandyssg/article/details/120039908

版权

记录专栏收录该内容

5 篇文章 0 订阅

订阅专栏

这篇文章详细介绍了如何在PyTorch中使用collate_fn函数来自定义数据加载器的行为，包括如何组织输入样本和处理caption等复杂结构。通过实例展示了如何根据句子长度和ID对数据进行预处理，使得数据适配于模型训练。

摘要由CSDN通过智能技术生成

参考文章：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)