自定义dataloader里的batch_sampler，你就不需要再做padding了！

最新推荐文章于 2026-06-19 12:42:05 发布

原创

最新推荐文章于 2026-06-19 12:42:05 发布 · 7k 阅读

标签

#dataloader

收录于

博客围绕自然语言处理中FixedLengthSampler的实现展开。通过统计句子长度，保证每个batch只采样相同长度句子，避免padding。详细介绍了建立length_map、state、order，补全不足batch_size样本及迭代生成的步骤，并给出完整代码，可将其类传给torch.utils.data.Dataloader的batch_sampler参数。

前言

最近再看一份代码，居然已经水了三篇博客…and这是第四篇。
在说batch_sampler之前，这里不要脸的放一下之前关于dataset和dataloader的博客

正文

在自然语言处理中，由于一个batch中的句子有长有短，我们对短的做padding。但今天我遇到的fixedlengthsampler给了我另外一种思路，我们可以统计所有句子的长度，然后保证每一个batch在采样的时候，只采样句子长度相同的句子，这样就不需要padding了。这个fixedlengthsampler具体怎么实现呢?

第一步:建立length_map

在这里插入图片描述
这里的self.data_source是已经定义好的，集成了torch.utils.data.Dataset的dataset。假设我们这里的self.data_source.dataset是形似:[[单词1，单词2，…], [单词1，单词2，…], …]这样的句子列表。
我们遍历这个句子列表，记录长度，并以句子长度为length_map的键，而对应的值是一个列表，里面是句子在句子列表中对应的下标。也就是说:length_map形如下（随便举个例子）:
{9:[10,9,11,1,2,3,4,5,6,7,8], 17:[18,20]}