Pytorch中的DataLoader处理机制

# 该代码主要是为了讲解介绍Dataloader的工作机制
from torch.utils import data
# 任何Dataset数据类的子类，并重写相关的函数
class NerDataset(data.Dataset):
    # 将需要的参数进行初始化
    def __init__(self, examples, tokenizer, label_map, max_seq_length):
        self.examples=examples
        self.tokenizer=tokenizer
        self.label_map=label_map
        self.max_seq_length=max_seq_length
    # 该函数是为了计算样本的数量，其意义在于保证__getitem__()的参数inx的取值范围在(0, len)之间
    def __len__(self):
        return len(self.examples)
    # 迭代传入__getitem__()索引，并返回相应的内容
    def __getitem__(self, idx):
        feat=example2feature(self.examples[idx], self.tokenizer, self.label_map, max_seq_length)
        return feat.input_ids, feat.input_mask, feat.segment_ids, feat.predict_mask, feat.label_ids

    # classmethod修饰符的函数不需要实例化，不需要 self 参数。第一个参数cls表示自身类，用于调用类的属性、类的方法、实例化对象等(如cls.tokenizer)。
    # 此函数的第二个参数batch，即是__getitem__()迭代batch_size次的返回结果来作为参数。
    @classmethod
    def pad(cls, batch):

        seqlen_list = [len(sample[0]) for sample in batch]
        maxlen = np.array(seqlen_list).max()

        f = lambda x, seqlen: [sample[x] + [0] * (seqlen - len(sample[x])) for sample in batch] # 0: X for padding
        input_ids_list = torch.LongTensor(f(0, maxlen))
        input_mask_list = torch.LongTensor(f(1, maxlen))
        segment_ids_list = torch.LongTensor(f(2, maxlen))
        predict_mask_list = torch.ByteTensor(f(3, maxlen))
        label_ids_list = torch.LongTensor(f(4, maxlen))

        return input_ids_list, input_mask_list, segment_ids_list, predict_mask_list, label_ids_list

train_dataset = NerDataset(train_examples,tokenizer,label_map,max_seq_length)

train_dataloader = data.DataLoader(dataset=train_dataset,
                                batch_size=batch_size,
                                shuffle=True,
                                num_workers=4,
                                collate_fn=NerDataset.pad)

我们详细解释一下data.Dataloader()的工作流程：

① 构建Dataset的子类（我们常使用的TensorDataset就是这样的一个类），得到实例化对象并传参给dataset。

② 将DataLoader看作一个迭代器，以sampler的style或者shuffle的style（详见Dataloader的参数）对dataset迭代batch_size次，来获取一个batch_data。每次迭代都调用dataset对象的__getitem__()函数获取一个样本，其中__getitem__()的参数idx取值方法是由style决定的(即随机取值，还是顺序取值)，idx取值范围是（0，self.__len__()）

③ 获取到batch_data之后，batch_data将作为collate_fn()函数的参数进行下一步操作，并最终返回collate_fn()操作后的batch_data结果。

④最终，我们将得到一个③返回的DataLoader数据对象（即collate_fn()结果），该数据对象供torch来进行下一步操作使用。

（注意：我们需要在上述过程中将变量类型处理成torch.tensor，一般选择过程①或③）

接着，你可以用下面的代码示例来解读一下：

或者，我们可以看看该代码的数据获取方式：
# 这里的train_features是一个对象列表，每个对象包含多种属性，每个属性(如f.input_ids)是一个列表。因此，all_input_ids是列表的列表。
all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features], dtype=torch.long)
all_valid_ids = torch.tensor([f.valid_ids for f in train_features], dtype=torch.long)
all_lmask_ids = torch.tensor([f.label_mask for f in train_features], dtype=torch.long)
# TensorDataset的参数可以是多个，每个都是一个列表
train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids,all_valid_ids,all_lmask_ids)
if args.local_rank == -1:
    train_sampler = RandomSampler(train_data)
else:
    train_sampler = DistributedSampler(train_data)
# DataLoader的结果可以得到一个列表train_dataloader,每个元素包含了TensorDataset的参数，每个参数的大小为batch_size
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size)

参考文献：

https://blog.csdn.net/dong_liuqi/article/details/114521240

https://blog.csdn.net/lyb3b3b/article/details/83713820

https://blog.csdn.net/weixin_39670464/article/details/111221971

https://www.runoob.com/python/python-func-classmethod.html

https://www.jianshu.com/p/8ea7fba72673

https://blog.csdn.net/u014380165/article/details/79058479

https://github.com/Louis-udm/NER-BERT-CRF/blob/master/NER_BERT_CRF.py