[code] Transformer For Summarization Source Code Reading [2]

Basic Information

作者：李丕绩（腾讯AI Lab）

模型：Transformer + copy mechanism for abstractive summarization

数据集：CNN/Daily Mail

Debug

1. run (main.py)

1. 指定变量existing_model_name，载入之前训练保存的模型参数

2. 进入run

3. 利用函数init_modules，初始化模型参数modules, constants, options。其中具体的超参数，在config.py中的class DeepmindConfigs中已有定义；将已经处理好的字典载入modules

4. 打印模型信息：函数print_basic_info，主要是modules, constants, options等参数，或者操作（初始化model之后，也可以将model打印一下）

5. 载入数据集

6. 函数datar.batched返回batch list（每个子列表中是batch size个样本的索引），总的样本数，以及batch的总数目

7. 实例化模型：model = Model(modules, consts, options)

8. 定义优化器：optimizer = torch.optim.Adagrad(model.parameters(), lr=consts["lr"], initial_accumulator_value=0.1)

9. 如果需要载入之前训练的模型，使用函数load_model载入model和optimizer的参数

10. 进入epoch循环：每个epoch都会对训练集重新进行打乱，重新生成batch list

11. 进入batch循环：按照当前batch list中子序列中的索引id，索引出对应的原始数据

12. 利用函数datar.get_data得到一个批量的数据——返回一个类class BatchData的实例（包含进行了word2id的输入x，目标y，考虑拓展词表的x_ext，y_ext，以及记录补零位的x_mask，y_mask；此时所有的数据都是numpy array）

13. 拓展词表x_ext_word，拓展词表是一个batch中所有样本共享的，拼接到fixed vocab的后面；y_ex区别于y之处：后者的OOV被设置为<unk>，但是前者给出了OOV在extended vocabulary中的索引；y_input是对y进行了shift，头部添加了标志位<bos>，作为解码输入

14. 根据一个batch的数据进行训练，惯常操作：

    # 梯度置零
    model.zero_grad()
    
    # 前向传播
    y_pred, cost = model()  # 输出y_pred是为了打印预测结果
    
    # 反向传播
    cost.backward()  # 保证反向传播的路径通畅
    
    # 梯度裁剪
    torch.nn.utils.clip_grad_norm_(model.parameters(), consts["norm_clip"])
    
    # 优化参数
    optimizer.step()
    

15. 打印信息

    1. print_size = num_files // print_time；表示每个epoch打印print_time次信息，也就是每处理print_size个数据打印一次信息
    2. 打印信息的同时保存一次训练好的model
    3. 每个epoch打印一次epoch的平均loss，以及所用的时间
    4. 如果上一个epoch的total error大于当前epoch的total error，则继续优化；否则，优化结束。
    5. 这部分是可以优化的：
       1. 多打印一些信息，少保存些model
       2. 每个epoch进行一次validation，并且打印出来生成的摘要
       3. 这个implementation没有用scheduled sampling可以尝试加入

16. 结束训练

17. 进行inference（此时必须载入训练好的model，以及预设的参数）

18. 载入测试数据，利用datar.batched划分batch，函数datar.get_data获取一个batch的数据

19. 编码：用模型的encodermodel.encode

20. 解码（摘要生成）：此处使用了beam search method；还分为了copy模式、以及non-copy模式

21. 进入beam_decoding函数（简述操作，不析代码）：

    1. 输入了考虑OOV的x_ext，x_mask（OOV对mask没有影响，mask的是padding位），word_embed（encoder的编码输出）；padding_mask与x_mask是互为相反（1,0位置）
    2. x_ext_words是一个列表，子列表中是每个样本的OOV词表；max_ext_len指的是当前batch中，最大的OOV词表容量；
    3. 还输入了ground truth y，以及原始的summary，为了计算生成的摘要的ROUGE
    4. 执行beam search decoding：采样；排序；选择TOP N；迭代进行；得到beam search的结果，并写入文件

2. model (model.py)

基本流程：

forward()
│
└─── encode()
│
└─── decode()
│   └───decoding()
│   └───word_prob_layer()
│
└─── label_smoothing_loss() or nll_loss()

2.1 input

传入model的数据（全部来自于类BatchData），该pytorch代码的实现中，batch_first全部设置为False：

torch.LongTensor(batch.x).to(options["device"])
# shape: [max_x_len, batch_size]
# 不考虑OOV的word2id处理后的source article

torch.LongTensor(batch.y_inp).to(options["device"])
# shape: [max_y_len, batch_size]
# 是y的基础上向后shift了一位，作为解码输入

torch.LongTensor(batch.y).to(options["device"])
# shape: [max_y_len, batch_size]
# 不考虑OOV的word2id处理后的target summary

torch.LongTensor(batch.x_ext).to(options["device"])
# shape: 
# 考虑了OOV的word2id处理后的source article

torch.LongTensor(batch.y_ext).to(options["device"])
# shape: 
# 考虑了OOV的word2id处理后的target summary

torch.FloatTensor(batch.x_mask).to(options["device"])
# shape: [max_x_len, batch_size, 1]
# 欲将一个batch中的所有序列对其，需要将每个序列补足到最大长度，mask掉padding位

torch.FloatTensor(batch.y_mask).to(options["device"])
# shape: [max_y_len, batch_size, 1]

batch.max_ext_len
# scalar
# 每个sample对应一个OOV list，一个batch中OOV的最大容量

2.2 encode

Args: 
	x: source article经过word2id的表示，尚未进行embedding

Returns：
	encoding hidden state：经过Transformer encoder编码的结果
	source padding mask：padding位置1

1. 预处理

   1. token embedding + positional embedding
   2. 层标准化（对embedding进行normalization）
   3. dropout（对于embedding）

2. 进入N encoder stack逐层进行编码。每一层的输出是隐含状态，每一层输入是上一层的输出

   xs = []
   for layer_id, layer in enumerate(self.enc_layers):
       x, _ ,_ = layer(x, self_padding_mask=padding_mask)  # 进行N个blocks的编码
       xs.append(x)
       
   

   列表 xs 保存了每一层的隐含状态。

   encoder layer的输入输出：

   1. 输入的shape：[max_x_len, batch_size, embedding_dim]；

   2. 输出的shape：[max_x_len, batch_size, embedding_dim]

   每一层的输入有二：1. 上一层的输出，2. padding mask；每一层encoder是一个Transformer单元，执行的操作有：

   1. 记录residual
   2. self-attention （query = key = value = x + dropout
   3. add residual + attention normalization
   4. 记录residual
   5. 全连接（从embed_dim映射到d_ff，维度变大）
   6. GELU activation function + dropout
   7. 全连接（从d_ff映射回embed_dim）
   8. dropout + add residual + feed forward normalization

   其中，attention normalization和 feed forward normalization使用的使用一个class LayerNorm。所做的操作都是减去均值，除以标准差，并且经过一个全连接层的映射。

3. final encoder layer output 作为 输入编码返回

2.3 decode

训练时用的是Teacher Forcing的decoder，需要输入shifted ground truth作为decoder input。

Args:
    y_inp		：shifted y
    			[max_y_len, batch_size]
        
    mask_y		：非padding的位置置为1
    			[max_y_len, batch_size, 1]
        
    mask_x		：非padding的位置置为1
    			[max_x_len, batch_size, 1]
        
    hs			：source article经过编码的隐含向量
    			[max_x_len, batch_size, embed_dim]
        
    src_padding_mask	：padding的位置置为1
    			[max_x_len, batch_size]
        
    x_ext		：考虑OOV的source article对应的ids（numpy array）
    			[max_x_len, batch_size]
        
    max_ext_len	：最大的OOV容量

Returns:
    y_dec
    
    attn_dist

1. 预处理（与embedding相同）

   1. token embedding + positional embedding
   2. 层标准化（对embedding进行normalization）
   3. dropout（对于embedding）
   4. attention mask，使用了Transformer中定义的类class SelfAttentionMask。此处返回了一个上三角矩阵（非零元素全为1），作为self_attn_mask

2. 进入N层解码。与encoder不同的时，decoder输入了encoder hidden states作为external memory，及其对应的padding（均为encoder的输出），以及self-attention mask。

   Transformer内部的具体操作流程为：

   1. 记录residual
   2. self-attention （query = key = value = x）+ dropout （注意！encoder和decoder的self-attention中的参数不一样，虽然都是相同的attention结构；model中对encoder、decoder定义的时候各自对Transformer分别进行了实例化）
   3. add residual + attention normalization
   4. 外部attention： y_inp的嵌入表示作为query，编码状态作为key和value，进行Multihead Attention的计算
   5. 记录residual
   6. 全连接（从embed_dim映射到d_ff，维度变大）
   7. GELU activation function + dropout
   8. 全连接（从d_ff映射回embed_dim）
   9. dropout + add residual + feed forward normalization

   记录transformer中的形状变化：

   x, residual 			[sequence_length, batch_size, embedding_dim]
   self-attention 			[sequence_length, batch_size, embedding_dim]
   dropout, add & norm 	操作不改变形状
   
   residual 				[sequence_length, batch_size, embedding_dim]
   external-attention 		[sequence_length, batch_size, embedding_dim]
   dropout, add & norm 	操作不改变形状
   
   residual 				[sequence_length, batch_size, embedding_dim]
   fully connected layer	[sequence_length, batch_size, ff_dim]
   gelu, dropout			操作不改变形状
   fully connected layer	[sequence_length, batch_size, embedding_dim]
   dropout, add & norm 	操作不改变形状
   

   最终Transformer的输入输出形状相同

2.4 vocabulary probability distribution (word_rob_layer.py)

利用解码结果，映射到单词表维度上进行vocabulary probability distribution的计算

Args:
    h: 			decoder output (hidden states)
    y_emb: 		decoder input（shifted）
    memory: 	encoder output (hidden states)
    mask_x: 	encoder output (padding_mask)
    xids: 		source article in the id representation, considering OOVs
    max_ext_len: maximum length of OOV lists

Returns:
    pred
    dists		decoder-encoder attention [max_y_len, batch_size, max_x_len]

计算单词表上概率分布的流程：

1. query = h，key = value = memory 进行一次external attention的计算，返回contexts，以及对应的attention weights，记作dists，shape：[num_query, batch_size, num_kv]

2. pred的计算：

   1. 基于decoder states h，decoder input y_emb，以及encoder-decoder context（external attention的输出）

   2. 将三者进行concatenation；shape：[max_y_len, batch_size, embed_dim * 3]

   3. 再将concatenation的hidden state维度映射到vocabulary size（fixed vocabulary）；shape：[max_y_len, batch_size, vocab_size]

   4. 再在单词表的维度上进行softmax表示成概率分布

   5. 给pred拼接上全零Tensor，使得单词表的维度大小变为vocabulary size + max_ext_len（extended vocabulary），shape: [max_y_len, batch_size, vocab_ext_size]

   6. 利用上述的concatenation，经过全连接层，和sigmoid激活函数，得到一个gate，[max_y_len, batch_size, 1]

   7. xids 的shape原本是 [max_x_len, batch_size]， 经过一次转置，在第一个维度上复制max_y_len次，形状变为 [max_y_len, batch_size, max_x_size]

   8. 用下述函数得到最终的概率分布：

      pred = (g * pred).scatter_add(2, xids, (1 - g) * dists)
      
      # selfTensor.scatter_add_(dim, indexTensor, otherTensor)
      # 该函数将otherTensor的所有值加到selfTensor中，加入位置由indexTensor指明
      # 理论上indexTensor和otherTensor的形状应该相同，index_Tensor中的元素取值范围应该小于selfTensor的指定维度dim
      
      # For a 3-D tensor, :attr:`self` is updated as::
      #    self[index[i][j][k]][j][k] += other[i][j][k]  # if dim == 0
      #    self[i][index[i][j][k]][k] += other[i][j][k]  # if dim == 1
      #    self[i][j][index[i][j][k]] += other[i][j][k]  # if dim == 2
      

   
   (g * pred)	  shape: [max_y_len, batch_size, vocab_ext_size]
   
   xids				shape: [max_y_len, batch_size, max_x_len]
   
   (1 - g) * dists shape: [max_y_len, batch_size, max_x_len]
   
   xids 中的元素（id）对应着(g * pred)中的每一个单词，也对应着(1 - g) * dists中的每一个（出现在source article中的）单词的注意力权重；
   
   将**这些单词**对应的（decoder-encoder）**attention weights**，与利用decoder输出计算得到的**pred** 进行加权求和。
   
   9. 返回最终的（拓展）单词表上的概率分布，以及decoder-encoder之间的attention weights
   
   

2.5 loss function (label_smoothing.py)

Args:
	y_inp: 		解码输入（shifted）
	y_tgt: 		目标序列（not shifted）
	mask_y:  	非padding位置1，padding位被置0

Returns: 
	avg_loss:	返回该样本中的平均loss：总loss / 总输出词数（而非 总loss / batch size）

1. 预处理

   y_pred = T.log(y_pred.clamp(min=1e-8))
   

   将y_pred中的概率值小于1e-8的数据，全部设置为1e-8

2. label smoothing

   def forward(self, output, target):
       # 计算real_size，此处相当与之前的max_ext_len
       if output.size(1) > self.size:
           real_size = output.size(1) - self.size
       else:
           real_size = 0
   
       model_prob = self.one_hot.repeat(target.size(0), 1)
       if real_size > 0:
           ext_zeros = torch.full((model_prob.size(0), real_size), self.smoothing_value).to(self.device)
           model_prob = torch.cat((model_prob, ext_zeros), -1)
           
       model_prob.scatter_(1, target, self.confidence)
       model_prob.masked_fill_((target == self.padding_idx), 0.)
       
       return F.kl_div(output, model_prob, reduction='sum')
   

   流程：

   1. 将y_pred（shape: [max_y_len, batch_size, vocab_ext_size]）展开成 [max_y_len * batch_size, vocab_ext_size]

   2. 将y_tgt（shape: [max_y_len, batch_size]）展开成 [max_y_len * batch_size, 1]

   3. model_prob由one_hot向量repeat而来，shape: [max_y_len * batch_size, vocab_size]

   4. ext_zeros的shape [max_y_len * batch_size, max_ext_len]，其中的元素的值均为smoothing value = 2e-6

   5. 将两者进行concatenation，得到model_prob，shape：[max_y_len * batch_size, vocab_ext_size]

   6. 在model_prob中target对应元素的位置，加上confidence = 0.9

      model_prob.scatter_(1, target, self.confidence)
      
      # torch.Tensor.scatter_(dim, index, src) -> Tensor
      # 该函数将src加到self中，加入位置由index指明
      
      # For a 3-D tensor, :attr:`self` is updated as::
      #    self[index[i][j][k]][j][k] += src  # if dim == 0
      #    self[i][index[i][j][k]][k] += src  # if dim == 1
      #    self[i][j][index[i][j][k]] += src  # if dim == 2
      

   7. 对结果进行mask

      model_prob.masked_fill_((target == self.padding_idx), 0.)
      
      # torch.Tensor.masked_fill_(mask, value)
      #     Args:
      #         mask (ByteTensor): the binary mask
      #         value (float): the value to fill in with
      # mask中元素为1的位置，被替换为value
      # selfTensor应该与maskTensor形状一样
      

   8. 返回model_prob与output之间的KL散度作为 label smoothing loss