快速从零构建成DataLoader

1 引言

在上一篇文章中，笔者详细介绍了在文本处理过程中如何通过torch.vocab来快速根据原始语料构建我们所需要的词表。在接下来的这篇文章中，笔者就来详细介绍一下如何在上一步的基础上快速构建Pytorch中的DataLoader，以便于后续网络的训练。

下面，笔者将会详细介绍在两种场景下文本数据集的构建过程。第一种场景是输入为序列输出为标签，也就是类似于文本分类数据集；第二种场景就是输入输出均序列，例如翻译模型，或者文本生成模型等。

2 构建类文本分类数据集

对于构建类似文本分类的数据集来说，总体上可以分为4个步骤：①构建字典；②将文本中的每一个词（字）转换为Token序列；③对不同长度的样本序列按照某个标准进行padding处理；④构建DataLoader类。现在假设我们有如下格式的原始数据：

1
问君能有几多愁，恰似一江春水向东流。  0
2
年年岁岁花相似，岁岁年年人不同。   0
3
去年今日此门中，人面桃花相映红。   2
4
人面不知何处去，桃花依旧笑春风。   1
5
渺沧海之一粟，羡长江之无穷。 3
6
人面不知何处去，桃花依旧笑春风。   1
7
月来客栈。  1

其中文本与标签之间通过一个控制符进行分割，下面我们开始来一步步构建数据集。

2.1 构建字典

由于这部分内容在上一篇文章中已经具体介绍过，所以这里直接贴出代码即可，如下：

​x
1
def tokenizer(s, word=False):
2
    """
3
    word: 是否采用分字模式
4
    """
5
    if word:
6
        r = [w for w in s]
7
    else:
8
        s = jieba.cut(s, cut_all=False)
9
        r = " ".join(s).split()
10
    return r
11
​
12
def build_vocab(tokenizer, filepath, word, min_freq, specials=None):
13
    """
14
    根据给定的tokenizer和对应参数返回一个Vocab类
15
    Args:
16
        tokenizer:  分词器
17
        filepath:  文本的路径
18
        word: 是否采用分字的模式对汉字进行处理
19
        min_freq: 最小词频，去掉小于min_freq的词
20
        specials: 特殊的字符，如<pad>，<unk>等
21
    Returns:
22
    """
23
    if specials is None:
24
        specials = ['<unk>', '<pad>', '<bos>', '<eos>']
25
    counter = Counter()
26
    with open(filepath, encoding='utf8') as f:
27
        for string_ in f:
28
            counter.update(tokenizer(string_.strip(), word))
29
    return Vocab(counter, min_freq=min_freq, specials=specials)

在完成上述过程后，我们将得到一个Vocab类的实例化对象，通过它便可以得到最终生成的字典：

xxxxxxxxxx
1
1
{'<unk>': 0, '<pad>': 1, '岁': 2, '年': 3, '。': 4, '，': 5, '不': 6, '人': 7, '似': 8, '春': 9, '花': 10, '一': 11, '东': 12,...}

此时，我们就需要定义一个类，并在类的初始化过程中根据训练语料完成字典的构建，代码如下：

xxxxxxxxxx
21
1
class LoadSentenceClassificationDataset():
2
    def __init__(self, train_file_path=None, # 训练集路径
3
                 tokenizer=None,
4
                 batch_size=2,
5
                 word=True, # 是否采用分字的模式对汉字进行处理
6
                 min_freq=1, # 最小词频，去掉小于min_freq的词
7
                 max_sen_len='same'): #最大句子长度，默认设置其长度为整个数据集中最长样本的长度
8
        # 根据训练预料建立英语和德语各自的字典
9
        self.tokenizer = tokenizer
10
        self.min_freq = min_freq
11
        self.specials = ['<unk>', '<pad>']
12
        self.word = word
13
        self.vocab = build_vocab(self.tokenizer,
14
                                 filepath=train_file_path,
15
                                 word=self.word,
16
                                 min_freq=self.min_freq,
17
                                 specials=self.specials)
18
        self.PAD_IDX = self.vocab['<pad>']
19
        self.UNK_IDX = self.vocab['<unk>']
20
        self.batch_size = batch_size
21
        self.max_sen_len = max_sen_len

2.2 转换为Token序列

在得到构建的字典后，便可以通过如下函数来将训练集、验证集和测试集转换成Token序列：

xxxxxxxxxx
18
1
    def data_process(self, filepath):
2
        """
3
        将每一句话中的每一个词根据字典转换成索引的形式，同时返回所有样本中最长样本的长度
4
        :param filepath: 数据集路径
5
        :return:
6
        """
7
        raw_iter = iter(open(filepath, encoding="utf8"))
8
        data = []
9
        max_len = 0
10
        for raw in raw_iter:
11
            line = raw.rstrip("\n")
12
            s, l = line.split('\t')
13
            tensor_ = torch.tensor([self.vocab[token] for token in
14
                                    self.tokenizer(s, self.word)], dtype=torch.long)
15
            l = torch.tensor(int(l), dtype=torch.long)
16
            max_len = max(max_len, tensor_.size(0))
17
            data.append((tensor_, l))
18
        return data, max_len

在上述代码中，其中第9行用来保存当前数据中最长样本的长度，在后续padding时会用到；第13-14行是先将原始文本序列tokenize，然后再转换成每个词（字）对应的token；第18行将返回包含所有样本的一个列表，以及当前语料中最长样本的长度。

例如如下两行样本

xxxxxxxxxx
2
1
问君能有几多愁，恰似一江春水向东流。  0
2
年年岁岁花相似，岁岁年年人不同。    0

在经过该函数处理后得到的结果为

xxxxxxxxxx
2
1
[(tensor([61, 36, 58, 45, 33, 37, 40,  4, 39, 18, 16, 23, 12, 49, 35, 30, 51,  3]), tensor(0)),
2
 (tensor([ 5,  5,  7,  7,  8, 24, 18,  4,  7,  7,  5,  5,  6, 10, 34,  3]), tensor(0)) ...]

2.3 padding处理

从上面的输出结果也可以看到，对于不同的样本来说其对应的长度通常来说都是不同的。但是在将数据输入到相应模型时却需要保持同样的长度，因此在这里我们就需要对Token序列化后的样本进行padding处理，具体代码如下：

xxxxxxxxxx
30
1
def pad_sequence(sequences, batch_first=False, max_len=None, padding_value=0):
2
    """
3
    对一个List中的元素进行padding
4
        sequences:
5
        batch_first: 是否把batch_size放到第一个维度
6
        padding_value:
7
        max_len : 最大句子长度，默认为None，即在每个batch中以最长样本的长度对其它样本进行padding；
8
        当指定max_len的长度小于一个batch中某个样本的长度，那么在这个batch中还是会以最长样本的长度对其它样本进行padding
9
        建议指定max_len的值为整个数据集中最长样本的长度
10
    Returns:
11
    """
12
    max_size = sequences[0].size()
13
    trailing_dims = max_size[1:]
14
    length = max_len
15
    max_len = max([s.size(0) for s in sequences])
16
    if length is not None:
17
        max_len = max(length, max_len)
18
    if batch_first:
19
        out_dims = (len(sequences), max_len) + trailing_dims
20
    else:
21
        out_dims = (max_len, len(sequences)) + trailing_dims
22
    out_tensor = sequences[0].data.new(*out_dims).fill_(padding_value)
23
    for i, tensor in enumerate(sequences):
24
        length = tensor.size(0)
25
        # use index notation to prevent duplicate references to the tensor
26
        if batch_first:
27
            out_tensor[i, :length, ...] = tensor
28
        else:
29
            out_tensor[:length, i, ...] = tensor
30
    return out_tensor

上述代码是根据torch.nn.utils.rnn中的pad_sequence函数修改而来，增加了可以指定一个全局最大长度的参数max_len。在经过pad_sequence函数处理后，所有的样本就会保持同样的长度。例如上面的tokenize后的结果在经过padding处理后将变为

xxxxxxxxxx
2
1
[(tensor([61, 36, 58, 45, 33, 37, 40,  4, 39, 18, 16, 23, 12, 49, 35, 30, 51,  3]), tensor(0)),
2
 (tensor([ 5,  5,  7,  7,  8, 24, 18,  4,  7,  7,  5,  5,  6, 10, 34,  3,  1,  1]), tensor(0)) ...]

即第2个样本的末尾padding了两个1。

2.4 构建DataLoader迭代器

在经过前面的一系列处理后，我们便可以通过如下代码来构建DataLoader迭代器：

xxxxxxxxxx
25
1
    def load_train_val_test_data(self, train_file_paths, val_file_paths, test_file_paths):
2
        train_data, max_sen_len = self.data_process(train_file_paths)  # 得到处理好的所有样本
3
        if self.max_sen_len == 'same':
4
            self.max_sen_len = max_sen_len
5
        val_data, _ = self.data_process(val_file_paths)
6
        test_data, _ = self.data_process(test_file_paths)
7
        train_iter = DataLoader(train_data, batch_size=self.batch_size,  # 构造DataLoader
8
                                shuffle=True, collate_fn=self.generate_batch)
9
        valid_iter = DataLoader(val_data, batch_size=self.batch_size,
10
                                shuffle=True, collate_fn=self.generate_batch)
11
        test_iter = DataLoader(test_data, batch_size=self.batch_size,
12
                               shuffle=True, collate_fn=self.generate_batch)
13
        return train_iter, valid_iter, test_iter
14
​
15
    def generate_batch(self, data_batch):
16
        batch_sentence, batch_label = [], []
17
        for (sen, label) in data_batch:  # 开始对一个batch中的每一个样本进行处理。
18
            batch_sentence.append(sen)
19
            batch_label.append(label)
20
        batch_sentence = pad_sequence(batch_sentence,  # [batch_size,max_len]
21
                                      padding_value=self.PAD_IDX,
22
                                      batch_first=True,
23
                                      max_len=self.max_sen_len)
24
        batch_label = torch.tensor(batch_label, dtype=torch.long)
25
        return batch_sentence, batch_label

在上述代码中，第1-13行便是用来构造最后需要返回的DataLoader迭代器；而第15-25行则是自定义一个函数来对每个batch中的样本进行处理，该函数将作为一个参数传入到类DataLoader中。同时，由于在DataLoader中是对每一个batch的数据进行处理，所以，当max_len=None时这就意味着上面的pad_sequence操作最终表现出来的结果就是不同的样本，padding后在同一个batch中长度是一样的，而在不同的batch之间可能是不一样的。因为此时pad_sequence是以一个batch中最长的样本为标准对其它样本进行padding。当max_len = 'same'时，最终表现出来的结果就是，所有样本在padding后的长度都等于训练集中最长样本的长度。

最终，在定义完成类LoadSentenceClassificationDataset后，便可以通过如下方式进行使用：

xxxxxxxxxx
16
1
if __name__ == '__main__':
2
    path = 'data_02.txt'
3
    data_loader = LoadSentenceClassificationDataset(train_file_path=path,
4
                                                    tokenizer=tokenizer,
5
                                                    batch_size=2,
6
                                                    word=True,
7
                                                    max_sen_len='same')
8
    train_iter, valid_iter, test_iter = data_loader.load_train_val_test_data(path, path, path)
9
    for sen, label in train_iter:
10
        print("batch:", sen)
11
        # batch: tensor([[6, 14, 10, 25, 19, 21, 11, 4, 13, 8, 20, 22, 26, 12, 27, 3, 1, 1],
12
        #                [61, 36, 58, 45, 33, 37, 40, 4, 39, 18, 16, 23, 12, 49, 35, 30, 51, 3]])
13
        print("batch size:", sen.shape)
14
        # batch size: torch.Size([2, 18])
15
        print("labels:", label)
16
        # labels: tensor([1, 0])

当然，暂时不想理解代码的朋友，可以直接将原始数据整理成上述一样的格式然后导入类LoadSentenceClassificationDataset使用即可（下载地址见文末[1]）。

3 构建类翻译模型数据集

通常，在NLP中还有一类任务就是模型的输入和输出均为序列的形式，这就需要在训练的过程中同时将这两部分输入到模型中。例如在翻译模型的训练过程中，需要同时将原始序列和目标序列一同输入到网络模型里。对于构建这种类似翻译模型的数据集总体上同样也可以采用上面的4个步骤：①构建字典；②将文本中的每一个词（字）转换为Token序列；③对不同长度的样本序列按照某个标准进行padding处理；④构建DataLoader类。只是在每一步中都需要分别对原始序列和目标序列进行处理。

现在假设我们有如下格式的平行语料数据：

xxxxxxxxxx
6
1
# 原始序列 source sequence
2
Zwei junge weiße Männer sind im, Freien in der Nähe vieler Büsche.
3
Mehrere Männer mit Schutzhelmen bedienen ein Antriebsradsystem.
4
Ein kleines Mädchen klettert in ein Spielhaus aus Holz.
5
Ein Mann in einem blauen Hemd steht auf einer Leiter und putzt ein Fenster.
6
Zwei Männer stehen am Herd und bereiten Essen zu.

xxxxxxxxxx
6
1
# 目标序列 target sequence
2
Two young, White males are outside near many bushes.
3
Several men in hard hats are operating a giant pulley system.
4
A little girl climbing into a wooden playhouse.
5
A man in a blue shirt is standing on a ladder cleaning a window.
6
Two men are at the stove preparing food.

从上面的语料可以看出，这是一个用于训练翻译模型的数据，原始序列为德语，目标序列为英语。下面我们开始来一步步的构建数据集。

3.1 构建字典

在构建字典的过程中，整体上与2.1节内容中的一样，只是在这里需要同时对原始序列和目标序列分别构建一个字典。具体代码如下所示：

xxxxxxxxxx
12
1
def my_tokenizer(s):
2
    s = s.replace(',', " ,").replace(".", " .")
3
    return s.split()
4
​
5
def build_vocab(tokenizer, filepath, specials=None):
6
    if specials is None:
7
        specials = ['<unk>', '<pad>', '<bos>', '<eos>']
8
    counter = Counter()
9
    with open(filepath, encoding='utf8') as f:
10
        for string_ in f:
11
            counter.update(tokenizer(string_))
12
    return Vocab(counter, specials=specials)

在上述代码中，第1-3行为自定义的一个tokenizer；虽然上述两种语料可以直接通过空格来对每个词进行分割，但是还需要做的就是在单词和符号之间加上一个空格，以便把符号分割出来。第5-12行定义的build_vocab函数还是同之前的一样，没有发生改变。

在完成上述过程后，我们将得到两个Vocab类的实例化对象。

一个为原始序列的字典：

xxxxxxxxxx
1
1
{'<unk>': 0, '<pad>': 1, '<bos>': 2, '<eos>': 3, '.': 4, 'Männer': 5, 'ein': 6, 'in': 7, 'Ein': 8, 'Zwei': 9, 'und': 10, ',': 11, ......}

一个为目标序列的字典：

xxxxxxxxxx
1
1
{'<unk>': 0, '<pad>': 1, '<bos>': 2, '<eos>': 3, '.': 4, 'a': 5, 'are': 6, 'A': 7, 'Two': 8, 'in': 9, 'men': 10, ',': 11, 'Several': 12,......}

此时，我们就需要定义一个类，并在类的初始化过程中根据训练语料完成字典的构建，代码如下：

xxxxxxxxxx
11
1
class LoadEnglishGermanDataset():
2
    def __init__(self, train_file_paths=None, tokenizer=None, batch_size=2):
3
        # 根据训练预料建立英语和德语各自的字典
4
        self.tokenizer = tokenizer
5
        self.de_vocab = build_vocab(self.tokenizer, filepath=train_file_paths[0])
6
        self.en_vocab = build_vocab(self.tokenizer, filepath=train_file_paths[1])
7
        self.specials = ['<unk>', '<pad>', '<bos>', '<eos>']
8
        self.PAD_IDX = self.de_vocab['<pad>']
9
        self.BOS_IDX = self.de_vocab['<bos>']
10
        self.EOS_IDX = self.de_vocab['<eos>']
11
        self.batch_size = batch_size

3.2 转换为Token序列

在得到构建的字典后，便可以通过如下函数来将训练集、验证集和测试集转换成Token序列：

xxxxxxxxxx
16
1
    def data_process(self, filepaths):
2
        """
3
        将每一句话中的每一个词根据字典转换成索引的形式
4
        :param filepaths:
5
        :return:
6
        """
7
        raw_de_iter = iter(open(filepaths[0], encoding="utf8"))
8
        raw_en_iter = iter(open(filepaths[1], encoding="utf8"))
9
        data = []
10
        for (raw_de, raw_en) in zip(raw_de_iter, raw_en_iter):
11
            de_tensor_ = torch.tensor([self.de_vocab[token] for token in
12
                                       self.tokenizer(raw_de.rstrip("\n"))], dtype=torch.long)
13
            en_tensor_ = torch.tensor([self.en_vocab[token] for token in
14
                                       self.tokenizer(raw_en.rstrip("\n"))], dtype=torch.long)
15
            data.append((de_tensor_, en_tensor_))
16
        return data

在上述代码中，第11-4行分别用来将原始序列和目标序列转换为对应词表中的Token形式。在处理完成后，就会得到类似如下的结果：

xxxxxxxxxx
4
1
 [(tensor([9,37, 46, 5, 42, 36, 11, 16,7, 33, 24, 45, 13,4]),tensor([8,45, 11, 13,28, 6, 34,31, 30,16,  4])),
2
  (tensor([22, 5, 40, 25, 30,  6, 12,  4]), tensor([12, 10,  9, 22, 23,  6, 33,  5, 20, 37, 41,  4])),
3
  (tensor([8, 38, 23, 39,  7,  6, 26, 29, 19,  4]), tensor([ 7, 27, 21, 18, 24,  5, 44, 35,  4])),
4
  (tensor([ 9,  5, 43, 27, 18, 10, 31, 14, 47,  4]), tensor([ 8, 10,  6, 14, 42, 40, 36, 19,  4]))  ]

其中左边的一列就是原始序列的Token形式，右边一列就是目标序列的Token形式，每一行构成一个样本。

3.3 padding处理

同样，从上面的输出结果可以看到，无论是对于原始序列来说还是目标序列来说，在不同的样本中其对应长度都不尽相同。但是在将数据输入到相应模型时却需要保持同样的长度，因此在这里我们就需要对Token序列化后的样本进行padding处理。同时需要注意的是，一般在这种生成模型中，模型在训练过程中只需要保证同一个batch中所有的原始序列等长，所有的目标序列等长，也就是说不需要在整个数据集中所有样本都保证等长。

因此，在实际处理过程中无论是原始序列还是目标序列都会以每个batch中最长的样本为标准对其它样本进行padding，具体代码如下：

xxxxxxxxxx
11
1
    def generate_batch(self, data_batch):
2
        de_batch, en_batch = [], []
3
        for (de_item, en_item) in data_batch:  # 开始对一个batch中的每一个样本进行处理。
4
             de_batch.append(de_item)  # 编码器输入序列不需要加起止符
5
             # 在每个idx序列的首位加上 起始token 和 结束 token
6
             en = torch.cat([torch.tensor([self.BOS_IDX]), en_item, torch.tensor([self.EOS_IDX])], dim=0)
7
             en_batch.append(en)
8
         # 以最长的序列为标准进行填充
9
         de_batch = pad_sequence(de_batch, padding_value=self.PAD_IDX)  # [de_len,batch_size]
10
         en_batch = pad_sequence(en_batch, padding_value=self.PAD_IDX)  # [en_len,batch_size]
11
         return de_batch, en_batch

在上述代码中，第6-7行用来在目标序列的首尾加上特定的起止符；第9-10行则是分别对一个batch中的原始序列和目标序列以各自当中最长的样本为标准进行padding（这里的pad_sequence导入自torch.nn.utils.rnn）。

3.4 构建DataLoader迭代器

在经过前面的一系列处理后，我们便可以通过如下代码来构建DataLoader迭代器：

xxxxxxxxxx
11
1
    def load_train_val_test_data(self, train_file_paths, val_file_paths, test_file_paths):
2
        train_data = self.data_process(train_file_paths)
3
        val_data = self.data_process(val_file_paths)
4
        test_data = self.data_process(test_file_paths)
5
        train_iter = DataLoader(train_data, batch_size=self.batch_size,
6
                                shuffle=True, collate_fn=self.generate_batch)
7
        valid_iter = DataLoader(val_data, batch_size=self.batch_size,
8
                                shuffle=True, collate_fn=self.generate_batch)
9
        test_iter = DataLoader(test_data, batch_size=self.batch_size,
10
                               shuffle=True, collate_fn=self.generate_batch)
11
        return train_iter, valid_iter, test_iter

最终，在定义完成类LoadEnglishGermanDataset后，便可以通过如下方式进行使用：

xxxxxxxxxx
11
1
if __name__ == '__main__':
2
    train_filepath = ['train_.de',
3
                      'train_.en']
4
    data_loader = LoadEnglishGermanCorpus(train_filepath, tokenizer=my_tokenizer, batch_size=2)
5
    train_iter, valid_iter, test_iter = data_loader.load_train_val_test_data(train_filepath,
6
                                                                             train_filepath,
7
                                                                             train_filepath)
8
    for src, tgt in train_iter:
9
        print("src shape：", src.shape)  # [de_tensor_len,batch_size]
10
        print("tgt shape:", tgt.shape)  # [de_tensor_len,batch_size]
11
        print("===================")

运行结果为：

xxxxxxxxxx
9
1
src shape： torch.Size([12, 2])
2
tgt shape: torch.Size([14, 2])
3
===================
4
src shape： torch.Size([16, 2])
5
tgt shape: torch.Size([13, 2])
6
===================
7
src shape： torch.Size([17, 1])
8
tgt shape: torch.Size([17, 1])
9
===================

从上述结果可以看出，对于同一个batch来说，原始序列的长度都相同，目标序列的长度也都相同。最后，如果暂时不想理解代码的朋友，可以直接将原始数据整理成上述一样的格式，然后导入类LoadEnglishGermanDataset使用即可。

4 总结

在这篇文章中，笔者首先介绍了如何快速的构建类似文本分类模型中的DataLoader迭代器，包括词表字典的构建、序列转换为Token、对不同长度的样本进行padding等；然后介绍了如何构建类似翻译模型中的数据迭代器，并且也对其用法进行了示例。

本次内容就到此结束，感谢您的阅读！如果你觉得上述内容对你有所帮助，欢迎分享至一位你的朋友！若有任何疑问与建议，请添加笔者微信nulls8或加群进行交流。青山不改，绿水长流，我们月来客栈见！

引用

[1] 完整代码：https://github.com/moon-hotel/DeepLearningWithMe