This post is all you need（⑦基于Transformer的对联模型）

1 引言

各位朋友大家好，欢迎来到月来客栈。经过前面6篇文章的介绍，对于Transformer相信大家都应该理解得差不多了。不过要想做到灵活运用Transformer结构，那就还得多看看其它场景下的运用。在接下来的这篇文章中，笔者将会以一个含有70余万条的对联数据集为例，来搭建一个基于Transformer结构的对联生成模型。同时，这也是介绍Transformer结构系列文章的最后一篇。

如图1所示便是一个基于Transformer结构的对联生成模型。可以看出，其实它与前面介绍的基于Transformer结构的翻译模型没有特别的变换，唯一不同的可能就是在对联生成模型中解码器和编码器共用了同一个词表（因为两者都是中文）。

2 数据预处理

2.1 语料介绍

按老规矩，在正式介绍模型搭建之前我们还是先来看看数据集都长什么样。本次所使用到的数据集是一个网上公开的对联数据集，在github中搜索"couplet-dataset”就能找到，其一共包含有770491条训练样本，4000条测试样本。和翻译数据集类似，对联数据集也包含上下两句：

1
# 上联 in.txt
2
风 弦 未 拨 心 先 乱 
3
花 梦 粘 于 春 袖 口 
4
晋 世 文 章 昌 二 陆 

xxxxxxxxxx
4
1
# 下联 out.txt
2
夜 幕 已 沉 梦 更 闲 
3
莺 声 溅 落 柳 枝 头 
4
魏 家 词 赋 重 三 曹 

如上所示便是3条样本，分别存放在in.txt和out.txt这两个文件中。可以看出，原始数据已经做了分字这步操作，所有后续我们只需要进行简单的split操作即可。

2.2 数据集构建

总体上来说对联生成模型的数据集构建过程和翻译模型的数据集构建过程基本上没有太大差别，主要步骤同样也是：①构建字典；②将文本中的每一个词（字）转换为Token序列；③对不同长度的样本序列按照某个标准进行padding处理；④构建DataLoader类。

第1步：定义tokenize

由于原始数据每个字已经被空格隔开了，所以这里tokenizer的定义只需要进行split操作即可，代码如下：

xxxxxxxxxx
2
1
def my_tokenizer(s):
2
    return s.split()

可以看到，其实也非常简单。例如对于如下文本来说

xxxxxxxxxx
1
1
腾 飞 上 铁 ， 锐 意 改 革 谋 发 展 ， 勇 当 千 里 马 

其tokenize后的结果为：

xxxxxxxxxx
1
1
['腾', '飞', '上', '铁', '，', '锐', '意', '改', '革', '谋', '发', '展', '，', '勇', '当', '千', '里', '马']

第2步：建立词表

在介绍完tokenize的实现方法后，我们就可以正式通过torchtext.vocab中的Vocab方法来构建词典了，代码如下：

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

在上述代码中，第3行代码用来指定特殊的字符；第5-10行分别用来遍历in.txt文件和out.txt文件中的每一个样本（每行一个）并进行tokenize和计数，其中对于counter.update进行介绍可以参考[1]；第8行则是返回最后得到词典。值得注意的是，由于在对联生成这一场景中编码器和解码器共用的是一个词表，所以这里同时对in.txt和out.txt文件进行了遍历。

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

xxxxxxxxxx
1
1
{'<unk>': 0, '<pad>': 1, '<bos>': 2, '<eos>': 3, '，': 4, '风': 5, '春': 6, '一': 7, '人': 8, '月': 9, '山': 10, '心': 11, '花': 12, '天': 13, ...}

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

xxxxxxxxxx
11
1
class LoadCoupletDataset():
2
    def __init__(self, train_file_paths=None, tokenizer=None,
3
                 batch_size=2, min_freq=1):
4
        # 根据训练预料建立字典，由于都是中文，所以共用一个即可
5
        self.tokenizer = tokenizer
6
        self.vocab = build_vocab(self.tokenizer, filepath=train_file_paths, min_freq=min_freq)
7
        self.specials = ['<unk>', '<pad>', '<bos>', '<eos>']
8
        self.PAD_IDX = self.vocab['<pad>']
9
        self.BOS_IDX = self.vocab['<bos>']
10
        self.EOS_IDX = self.vocab['<eos>']
11
        self.batch_size = batch_size

第3步：转换为Token序列

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

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

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

xxxxxxxxxx
1
1
[(tensor([   5,  549,  250, 1758,   11,  228,  651]), tensor([ 154, 1420,  310,  598,   29,  206,  164])), (tensor([  12,   29, 3218,  262,    6,  628,  419]), tensor([ 441,   62, 2049,   93,   66,  304,  111])), (tensor([1137,   40,   47,  286,  819,  364, 1383]), tensor([1803,   49,  586,  556,  126,   25, 1830])), (tensor([  7, 291, 138, 115, 216, 151,   9]), tensor([ 15, 311, 107,  57,  80,   5,  21]))]

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

第4步：padding处理

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

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

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

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

第5步：构造mask向量

在处理完成前面几个步骤后，进一步需要根据src_input和tgt_input来构造相关的mask向量，具体代码如下：

​x
1
    def generate_square_subsequent_mask(self, sz, device):
2
        mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
3
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
4
        return mask
5
​
6
    def create_mask(self, src, tgt, device='cpu'):
7
        src_seq_len = src.shape[0]
8
        tgt_seq_len = tgt.shape[0]
9
        tgt_mask = self.generate_square_subsequent_mask(tgt_seq_len, device)  # [tgt_len,tgt_len]
10
        # Decoder的注意力Mask输入，用于掩盖当前position之后的position，所以这里是一个对称矩阵
11
        src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)
12
        # Encoder的注意力Mask输入，这部分其实对于Encoder来说是没有用的，所以这里全是0
13
        src_padding_mask = (src == self.PAD_IDX).transpose(0, 1)
14
        # 用于mask掉Encoder的Token序列中的padding部分,[batch_size, src_len]
15
        tgt_padding_mask = (tgt == self.PAD_IDX).transpose(0, 1)
16
        # 用于mask掉Decoder的Token序列中的padding部分,batch_size, tgt_len
17
        return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

在上述代码中，第1-4行是用来生成一个形状为[sz,sz]的注意力掩码矩阵，用于在解码过程中掩盖当前position之后的position；第6-17行用来返回Transformer中各种情况下的mask矩阵，其中src_mask在这里并没有作用。

第6步：构造DataLoade与使用示例

经过前面5步的操作，整个数据集的构建就算是已经基本完成了，只需要再构造一个DataLoader迭代器即可，代码如下：

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

在上述代码中，第2-3行便是分别用来将训练集和测试集转换为Token序列；第4-7行则是分别构造2个DataLoader，其中generate_batch将作为一个参数传入来对每个batch的样本进行处理。在完成类LoadCoupletDataset所有的编码过程后，便可以通过如下形式进行使用：

xxxxxxxxxx
23
1
if __name__ == '__main__':
2
    config = Config()
3
    data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
4
                                     batch_size=config.batch_size,
5
                                     tokenizer=my_tokenizer,
6
                                     min_freq=config.min_freq)
7
    train_iter, test_iter = data_loader.load_train_val_test_data(config.test_corpus_file_paths,
8
                                                                 config.test_corpus_file_paths)
9
    print(data_loader.PAD_IDX)
10
    for src, tgt in train_iter:
11
        tgt_input = tgt[:-1, :]
12
        tgt_out = tgt[1:, :]
13
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = data_loader.create_mask(src, tgt_input)
14
        print(src)
15
        print(tgt)
16
        print("src shape：", src.shape)  # [in_tensor_len,batch_size]
17
        print("tgt shape:", tgt.shape)  # [in_tensor_len,batch_size]
18
        print("src input shape:", src.shape)
19
        print("src_padding_mask shape (batch_size, src_len): ", src_padding_mask.shape)
20
        print("tgt input shape:", tgt_input.shape)
21
        print("tgt_padding_mask shape: (batch_size, tgt_len) ", tgt_padding_mask.shape)
22
        print("tgt output shape:", tgt_out.shape)
23
        print("tgt_mask shape (tgt_len,tgt_len): ", tgt_mask.shape)

在介绍完数据集构建的整个过程后，下面就开始正式进入到翻译模型的构建中。如果对于这部分不是特别理解的话，建议先看这篇文章[2]中的数据处理流程图进行理解。

3 基于Transformer的对联生成模型

3.1 网络结构

总体来说，基于Transformer的对联生成模型的网络结构其实就是图1所展示的所有部分，只是在前面介绍Transformer网络结构[3]时笔者并没有把Embedding部分的实现给加进去。这是因为对于不同的文本生成模型，其Embedding部分会不一样（例如在本场景中编码器和解码器共用一个TokenEmbedding即可，而在翻译模型中就需要两个），所以将两者进行了拆分。同时，待模型训练完成后，在inference过程中Encoder只需要执行一次，所以在此过程中也需要单独使用Transformer中的Encoder和Decoder。

首先，我们需要定义一个名为CoupletModel的类，其前向传播过程代码如下所示：

xxxxxxxxxx
39
1
class CoupletModel(nn.Module):
2
    def __init__(self, vocab_size,
3
                 d_model=512, nhead=8, num_encoder_layers=6,
4
                 num_decoder_layers=6, dim_feedforward=2048,
5
                 dropout=0.1):
6
        super(CoupletModel, self).__init__()
7
        self.my_transformer = MyTransformer(d_model=d_model,
8
                                            nhead=nhead,
9
                                            num_encoder_layers=num_encoder_layers,
10
                                            num_decoder_layers=num_decoder_layers,
11
                                            dim_feedforward=dim_feedforward,
12
                                            dropout=dropout)
13
        self.pos_embedding = PositionalEncoding(d_model=d_model, dropout=dropout)
14
        self.token_embedding = TokenEmbedding(vocab_size, d_model)
15
        self.classification = nn.Linear(d_model, vocab_size)
16
​
17
    def forward(self, src=None, tgt=None, src_mask=None,
18
                tgt_mask=None, memory_mask=None, src_key_padding_mask=None,
19
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
20
        """
21
        :param src: Encoder的输入 [src_len,batch_size]
22
        :param tgt: Decoder的输入 [tgt_len,batch_size]
23
        :param src_key_padding_mask: 用来Mask掉Encoder中不同序列的padding部分,[batch_size, src_len]
24
        :param tgt_key_padding_mask: 用来Mask掉Decoder中不同序列的padding部分 [batch_size, tgt_len]
25
         memory_key_padding_mask: 用来Mask掉Encoder输出的memory中不同序列的padding部分 [batch_size, src_len]
26
        :return:
27
        """
28
        src_embed = self.token_embedding(src)  # [src_len, batch_size, embed_dim]
29
        src_embed = self.pos_embedding(src_embed)  # [src_len, batch_size, embed_dim]
30
        tgt_embed = self.token_embedding(tgt)  # [tgt_len, batch_size, embed_dim]
31
        tgt_embed = self.pos_embedding(tgt_embed)  # [tgt_len, batch_size, embed_dim]
32
        outs = self.my_transformer(src=src_embed, tgt=tgt_embed, src_mask=src_mask,
33
                                   tgt_mask=tgt_mask, memory_mask=memory_mask,
34
                                   src_key_padding_mask=src_key_padding_mask,
35
                                   tgt_key_padding_mask=tgt_key_padding_mask,
36
                                   memory_key_padding_mask=memory_key_padding_mask)
37
        # [tgt_len,batch_size,embed_dim]
38
        logits = self.classification(outs)  # [tgt_len,batch_size,tgt_vocab_size]
39
        return logits

在上述代码中，第7-12行便是用来定义一个Transformer结构；第13-15分别用来定义Positional Embedding、Token Embedding和最后的分类器（需要注意的是这里是共用同一个Token Embedding）；第28-38行便是用来执行整个前向传播过程，其中Transformer的整个前向传播过程在前一篇[3]文章中已经介绍过，在这里就不再赘述。

在定义完logits的前向传播过后，便可以通过如下形式进行使用：

xxxxxxxxxx
24
1
if __name__ == '__main__':
2
    src_len = 7
3
    batch_size = 2
4
    dmodel = 32
5
    tgt_len = 8
6
    num_head = 4
7
    src = torch.tensor([[4, 3, 2, 6, 0, 0, 0],
8
                        [5, 7, 8, 2, 4, 0, 0]]).transpose(0, 1)  # 转换成 [src_len, batch_size]
9
    src_key_padding_mask = torch.tensor([[True, True, True, True, False, False, False],
10
                                         [True, True, True, True, True, False, False]])
11
    tgt = torch.tensor([[1, 3, 3, 5, 4, 3, 0, 0],
12
                        [1, 6, 8, 2, 9, 1, 0, 0]]).transpose(0, 1)
13
    tgt_key_padding_mask = torch.tensor([[True, True, True, True, True, True, False, False],
14
                                         [True, True, True, True, True, True, False, False]])
15
    trans_model = CoupletModel(vocab_size=10, d_model=dmodel,
16
                               nhead=num_head, num_encoder_layers=6,
17
                               num_decoder_layers=6, dim_feedforward=30,
18
                               dropout=0.1)
19
    tgt_mask = trans_model.my_transformer.generate_square_subsequent_mask(tgt_len)
20
    logits = trans_model(src, tgt=tgt, tgt_mask=tgt_mask,
21
                         src_key_padding_mask=src_key_padding_mask,
22
                         tgt_key_padding_mask=tgt_key_padding_mask,
23
                         memory_key_padding_mask=src_key_padding_mask)
24
    print(logits.shape)

接着，我们需要再定义一个Encoder和Decoder在inference中进行使用，代码如下：

xxxxxxxxxx
12
1
    def encoder(self, src):
2
        src_embed = self.token_embedding(src)  # [src_len, batch_size, embed_dim]
3
        src_embed = self.pos_embedding(src_embed)  # [src_len, batch_size, embed_dim]
4
        memory = self.my_transformer.encoder(src_embed)
5
        return memory
6
​
7
    def decoder(self, tgt, memory, tgt_mask):
8
        tgt_embed = self.tgt_token_embedding(tgt)  # [tgt_len, batch_size, embed_dim]
9
        tgt_embed = self.pos_embedding(tgt_embed)  # [tgt_len, batch_size, embed_dim]
10
        outs = self.my_transformer.decoder(tgt_embed, memory=memory,
11
                                           tgt_mask=tgt_mask)  # [tgt_len,batch_size,embed_dim]
12
        return outs

在上述代码中，第1-5行用于在inference时对输入序列进行编码并得到memory（只需要执行一次）；第7-11行用于根据memory和当前解码时刻的输入对输出进行预测，需要循环执行多次，这部分内容详见模型预测部分。

3.2 模型训练

在定义完成整个对联生成模型的网络结构后下面就可以开始训练模型了。由于这部分代码较长，所以下面笔者依旧以分块的形式进行介绍：

第1步：载入数据集

xxxxxxxxxx
8
1
def train_model(config):
2
    data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
3
                                     batch_size=config.batch_size,
4
                                     tokenizer=my_tokenizer,
5
                                     min_freq=config.min_freq)
6
    train_iter, test_iter = \
7
        data_loader.load_train_val_test_data(config.train_corpus_file_paths,
8
                                             config.test_corpus_file_paths)

首先我们可以根据前面的介绍，通过类LoadCoupletDataset来载入数据集，其中config中定义了模型所涉及到的所有配置参数。

第2步：定义模型并初始化权重

xxxxxxxxxx
10
1
    couplet_model = CoupletModel(vocab_size=len(data_loader.vocab),
2
                                 d_model=config.d_model,
3
                                 nhead=config.num_head,
4
                                 num_encoder_layers=config.num_encoder_layers,
5
                                 num_decoder_layers=config.num_decoder_layers,
6
                                 dim_feedforward=config.dim_feedforward,
7
                                 dropout=config.dropout)
8
    for p in couplet_model.parameters():
9
        if p.dim() > 1:
10
            nn.init.xavier_uniform_(p)

在载入数据后，便可以定义模型CoupletModel，并根据相关参数对其进行实例化；同时，可以对整个模型中的所有参数进行一个初始化操作。

第3步：定义损失学习率与优化器

xxxxxxxxxx
5
1
    loss_fn = torch.nn.CrossEntropyLoss(ignore_index=data_loader.PAD_IDX)
2
    learning_rate = CustomSchedule(config.d_model)
3
    optimizer = torch.optim.Adam(couplet_model.parameters(),
4
                                 lr=0.,
5
                                 betas=(config.beta1, config.beta2), eps=config.epsilon)

在上述代码中，第1行是定义交叉熵损失函数，并同时指定需要忽略的索引ignore_index。因为根据tgt_output可知，有些位置上的标签值其实是Padding后的结果，因此在计算损失的时候需要将这些位置给忽略掉。第2行代码则是论文[4]中所提出来的动态学习率计算过程，其计算公式为：

具体实现代码为：

xxxxxxxxxx
12
1
class CustomSchedule(nn.Module):
2
    def __init__(self, d_model, warmup_steps=4000):
3
        super(CustomSchedule, self).__init__()
4
        self.d_model = torch.tensor(d_model, dtype=torch.float32)
5
        self.warmup_steps = warmup_steps
6
        self.step = 1.
7
​
8
    def __call__(self):
9
        arg1 = self.step ** -0.5
10
        arg2 = self.step * (self.warmup_steps ** -1.5)
11
        self.step += 1.
12
        return (self.d_model ** -0.5) * min(arg1, arg2)

通过CustomSchedule，就能够在训练过程中动态的调整学习率。学习率随step增加而变换的结果如图2所示：

从图2可以看出，在前warm_up个step中，学习率是线性增长的，在这之后便是非线性下降，直至收敛与0.0004。

第4步：开始训练

xxxxxxxxxx
30
1
    for epoch in range(config.epochs):
2
        losses = 0
3
        start_time = time.time()
4
        for idx, (src, tgt) in enumerate(train_iter):
5
            src = src.to(config.device)  # [src_len, batch_size]
6
            tgt = tgt.to(config.device)
7
            tgt_input = tgt[:-1, :]  # 解码部分的输入, [tgt_len,batch_size]
8
            src_mask, tgt_mask, src_padding_mask, tgt_padding_mask \
9
                = data_loader.create_mask(src, tgt_input, config.device)
10
            logits = couplet_model(
11
                src=src,  # Encoder的token序列输入，[src_len,batch_size]
12
                tgt=tgt_input,  # Decoder的token序列输入,[tgt_len,batch_size]
13
                src_mask=src_mask,  # Encoder的注意力Mask输入，这部分其实对于Encoder来说是没有用的
14
                tgt_mask=tgt_mask,
15
                # Decoder的注意力Mask输入，用于掩盖当前position之后的position [tgt_len,tgt_len]
16
                src_key_padding_mask=src_padding_mask,  # 用于mask掉Encoder的Token序列中的padding部分
17
                tgt_key_padding_mask=tgt_padding_mask,  # 用于mask掉Decoder的Token序列中的padding部分
18
                memory_key_padding_mask=src_padding_mask)  # 用于mask掉Encoder的Token序列中的padding部分
19
            # logits 输出shape为[tgt_len,batch_size,tgt_vocab_size]
20
            optimizer.zero_grad()
21
            tgt_out = tgt[1:, :]  # 解码部分的真实值  shape: [tgt_len,batch_size]
22
            loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
23
            # [tgt_len*batch_size, tgt_vocab_size] with [tgt_len*batch_size, ]
24
            loss.backward()
25
            lr = learning_rate()
26
            for p in optimizer.param_groups:
27
                p['lr'] = lr
28
            optimizer.step()
29
            losses += loss.item()
30
            acc, _, _ = accuracy(logits, tgt_out, data_loader.PAD_IDX)

在上述代码中，第5-9行是用来得到模型各个部分的输入；第10-18行是计算模型整个前向传播的过程；第20-24行则是执行损失计算与反向传播；第26-28则是将每个step更新后的学习率送入到模型中并进行参数更新；第30行是用来计算模型预测的准确率，具体过程将在后续文章中进行介绍。以下便是模型训练过程中的输出：

xxxxxxxxxx
5
1
-- INFO: Epoch: 0, Batch[29/3010], Train loss : 8.965, Train acc: 0.094
2
-- INFO: Epoch: 0, Batch[59/3010], Train loss : 8.618, Train acc: 0.098
3
-- INFO: Epoch: 0, Batch[89/3010], Train loss : 8.366, Train acc: 0.099
4
-- INFO: Epoch: 0, Batch[119/3010], Train loss : 8.137, Train acc: 0.109
5
...

3.3 模型预测

在介绍完模型的训练过程后接下来就来看模型的预测部分。生成模型的预测部分不像普通的分类任务只需要将网络最后的输出做argmax操作即可，生成模型在预测过程中往往需要按时刻一步步进行来进行。因此，下面我们这里定义一个couplet函数来执行这一过程，具体代码如下：

xxxxxxxxxx
11
1
def couplet(model, src, data_loader, config):
2
    vocab = data_loader.vocab
3
    tokenizer = data_loader.tokenizer
4
    model.eval()
5
    tokens = [vocab.stoi[tok] for tok in tokenizer(src)]  # 构造一个样本
6
    num_tokens = len(tokens)
7
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1))  # 将src_len 作为第一个维度
8
    tgt_tokens = greedy_decode(model, src, max_len=num_tokens + 5,
9
                               start_symbol=data_loader.BOS_IDX, config=config,
10
                               data_loader=data_loader).flatten()  # 解码的预测结果
11
    return "".join([vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

在上述代码中，第5行是将待翻译的源序列进行序列化操作；第7-10行则是通过函数greedy_decode函数来对输入进行解码；第11行则是将最后解码后的结果由Token序列在转换成实际的目标语言。同时，greedy_decode函数的实现如下：

xxxxxxxxxx
19
1
def greedy_decode(model, src, max_len, start_symbol, config, data_loader):
2
    src = src.to(config.device)
3
    memory = model.encoder(src)  # 对输入的Token序列进行解码翻译
4
    ys = torch.ones(1, 1).fill_(start_symbol). \
5
        type(torch.long).to(config.device)  # 解码的第一个输入，起始符号
6
    for i in range(max_len - 1):
7
        memory = memory.to(config.device)
8
        tgt_mask = (model.my_transformer.generate_square_subsequent_mask(ys.size(0))
9
                    .type(torch.bool)).to(config.device)  # 根据tgt_len产生一个注意力mask矩阵（对称的）
10
        out = model.decoder(ys, memory, tgt_mask)  # [tgt_len,tgt_vocab_size]
11
        out = out.transpose(0, 1)  # [tgt_vocab_size, tgt_len]
12
        prob = model.classification(out[:, -1])  # 只对对预测的下一个词进行分类
13
        _, next_word = torch.max(prob, dim=1)  # 选择概率最大者
14
        next_word = next_word.item()
15
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
16
        # 将当前时刻解码的预测输出结果，同之前所有的结果堆叠作为输入再去预测下一个词。
17
        if next_word == data_loader.EOS_IDX:  # 如果当前时刻的预测输出为结束标志，则跳出循环结束预测。
18
            break
19
    return ys

在上述代码中，第3行是将源序列输入到Transformer的编码器中进行编码并得到Memory；第4-5行是初始化解码阶段输入的第1个时刻的，在这里也就是' '；第6-18行则是整个循环解码过程，在下一个时刻为EOS_IDX或者达到最大长度后停止；第8-9行是根据当前解码器输入的长度生成注意力掩码矩阵tgt_mask；第10行是根据memory以及当前时刻的输入对当前时刻的输出进行解码；第12-14行则是分类得到当前时刻的解码输出结果；第15行则是将当前时刻的解码输出结果头当前时刻之前所有的输入进行拼接，以此再对下一个时刻的输出进行预测。

最后，我们只需要调用如下函数便可以完成对原始输入上联的下联生成任务：

xxxxxxxxxx
33
1
def do_couplet(src, config):
2
    data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
3
                                     batch_size=config.batch_size,
4
                                     tokenizer=my_tokenizer,
5
                                     min_freq=config.min_freq)
6
    couplet_model = CoupletModel(vocab_size=len(data_loader.vocab),
7
                                 d_model=config.d_model,
8
                                 nhead=config.num_head,
9
                                 num_encoder_layers=config.num_encoder_layers,
10
                                 num_decoder_layers=config.num_decoder_layers,
11
                                 dim_feedforward=config.dim_feedforward,
12
                                 dropout=config.dropout)
13
    couplet_model = couplet_model.to(config.device)
14
    loaded_paras = torch.load(config.model_save_dir + '/model.pkl')
15
    couplet_model.load_state_dict(loaded_paras)
16
    r = couplet(couplet_model, src, data_loader, config)
17
    return r
18
if __name__ == '__main__':
19
    srcs = ["晚风摇树树还挺",
20
            "忽忽几晨昏，离别间之，疾病间之，不及终年同静好",
21
            "风声、雨声、读书声，声声入耳",
22
            "上海自来水来自海上"]
23
    tgts = ["晨露润花花更红",
24
            "茕茕小儿女，孱羸若此，娇憨若此，更烦二老费精神",
25
            "家事、国事、天下事，事事关心",
26
            ""]
27
    config = Config()
28
    for i, src in enumerate(srcs):
29
        r = do_couplet(" ".join(src), config)
30
        print(f"上联：{src}")
31
        print(f" AI：{r}")
32
        print(f"下联：{tgts[i]}")
33
        print("=======")

在上述代码中，第6-15行是定义网络结构，以及恢复本地保存的网络权重；第16行则是开始执行下联生成任务；第19-26行为生成示例，其输出结果为：

xxxxxxxxxx
14
1
上联：晚风摇树树还挺
2
 AI: 朝露沾花花更红
3
下联：晨露润花花更红
4
​
5
上联：忽忽几晨昏，离别间之，疾病间之，不及终年同静好
6
 AI：茕茕小儿女，孱羸若此，娇憨若此，更烦二老费精神
7
下联：茕茕小儿女，孱羸若此，娇憨若此，更烦二老费精神
8
​
9
上联：风声、雨声、读书声，声声入耳
10
 AI：山色、水色、烟霞色，色色宜人
11
下联：家事、国事、天下事，事事关心
12
​
13
上联：上海自来水来自海上
14
 AI：中山落叶松叶落山中

以上完整代码可参见[5]。

4 总结

在这篇文章中，笔者首先介绍了对联生成模型的整个数据预处理过程；接着笔者介绍了基于Transformer结构的对联模型的整体构成，然后循序渐进地带着各位读者来实现了整个模型，包括基础结构的搭建、模型训练的详细实现、动态学习率的调整实现等；最后介绍了如何来实现模型在实际预测过程中的处理流程等，包括源输入序列的构建、解码时刻输入序列的构建等。

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

引用

[1] 你还在手动构造词表？试试torch.vocab

[2] This post is all you need（⑤基于Transformer的翻译模型）

[3] This post is all you need（④Transformer的实现过程）

[4] Attention is all you need

[5] https://github.com/moon-hotel/TransformerCouplet