[fairseq]translation task model 以及transformer的实现

translation.py

• 在train的时候如果不指定sourse language 和target language，那么，language pair必须是这样的形式.-.(...).idx,split是train valid test
• 也就是说language pair要在第一个.后面，以及用-隔开两个语言对

parts = filename.split('.')
if len(parts) >= 3 and len(parts[1].split('-')) == 2:
    return parts[1].split('-')

• 加载的词典的命名必须是dict.{lang1}.txt

src_dict = cls.load_dictionary(os.path.join(args.data[0], 'dict.{}.txt'.format(args.source_lang)))
tgt_dict = cls.load_dictionary(os.path.join(args.data[0], 'dict.{}.txt'.format(args.target_lang)))

• 其实发现translaion task 其实没有什么东西，全是一些如何加载预训练模型，以及如何加载数据，如何将数据处理成翻译需要的形式，因为主要是继承fairseq_task的功能，而fairseq_task本身就是一个seq2seq，因此只用translation.py实现加载数据什么的就行了。

fairseq_task.py

• 首先可以看到一个比较有意思的点，是创建词典的时候说词典的大小必须为8的倍数这样nvidia能够加速。

def build_dictionary(cls, filenames, workers=1, threshold=-1, nwords=-1, padding_factor=8):
    """Build the dictionary

    Args:
        filenames (list): list of filenames
        workers (int): number of concurrent workers
        threshold (int): defines the minimum word count
        nwords (int): defines the total number of words in the final dictionary,
            including special symbols
        padding_factor (int): can be used to pad the dictionary size to be a
            multiple of 8, which is important on some hardware (e.g., Nvidia
            Tensor Cores).
    """
    d = Dictionary()
    for filename in filenames:
        Dictionary.add_file_to_dictionary(filename, d, tokenizer.tokenize_line, workers)
    d.finalize(threshold=threshold, nwords=nwords, padding_factor=padding_factor)
    return d

• 里面的train_step，可以看到task还承担了调用train valid test,以及inference(decode)的过程，里面获得loss等等是通过criterion来获得的。
• sample是一个minibatch，就是fairseq的translation task类里面实现的的读取数据的操作。

def train_step(self, sample, model, criterion, optimizer, ignore_grad=False):
    """
    Do forward and backward, and return the loss as computed by *criterion*
    for the given *model* and *sample*.

    Args:
        sample (dict): the mini-batch. The format is defined by the
            :class:`~fairseq.data.FairseqDataset`.
        model (~fairseq.models.BaseFairseqModel): the model
        criterion (~fairseq.criterions.FairseqCriterion): the criterion
        optimizer (~fairseq.optim.FairseqOptimizer): the optimizer
        ignore_grad (bool): multiply loss by 0 if this is set to True

    Returns:
        tuple:
            - the loss
            - the sample size, which is used as the denominator for the
                gradient
            - logging outputs to display while training
    """
    model.train()
    loss, sample_size, logging_output = criterion(model, sample)
    if ignore_grad:
        loss *= 0
    optimizer.backward(loss)
    return loss, sample_size, logging_output

• inference_step的时候是通过generator.generate来产生数据的。

def inference_step(self, generator, models, sample, prefix_tokens=None):
    with torch.no_grad():
        return generator.generate(models, sample, prefix_tokens=prefix_tokens)

• criterion是在FairseqCriterion中定义的，但是这个task对应的criterion是在哪里呢？
• train_step，test_step，valid_step是在task中定义的，然后通过train.py中创建trainer来调用的。
• inference_step调用的是generatre，这个位于sequence_generator.py里面。

sequence_generator.py

• 可以选择使用sampling或者是beam search或者是sample topk，temperature控制采样的概率>1更趋近于均匀采样，<1更趋近于尖锐的采样。

sampling (bool, optional): sample outputs instead of beam search
         (default: False)
sampling_topk (int, optional): only sample among the top-k choices
         at each step (default: -1)
sampling_temperature (float, optional): temperature for sampling,
        where values >1.0 produces more uniform sampling and values
        <1.0 produces sharper sampling (default: 1.0)

• generate的时候可以传入一个models然后获得EnsembleModel。
• 要搞清楚encoder_input是什么，首先要搞清楚sample是什么。

encoder_input = {
    k: v for k, v in sample['net_input'].items()
    if k != 'prev_output_tokens'
}

• 没有在循环里面说明是一次得到所有的encoder输出。

encoder_outs = model.forward_encoder(encoder_input)

• 还有这个encoder_outs = model.reorder_encoder_out(encoder_outs, new_order)是干啥的？是根据new_order顺序对encdoer_outs重新排序的
• 最为重要的函数是finalize_hypos，有一个参数是step，是当前的时间步，具体功能是？
• model.forward_decoder是最为关键的一步，根据之前所生成的tokens和encoder的输出得到当前步每个单词的概率，以及atten_scores，这里应该写了teacher forcing。
• 在这里可以看到forward_decoder是一个step一个step进行的，而encoder的时候是全部的encoder出来。
• 所以说控制是否生成unk只用decoder的时候输出概率的时候将lprobs[:,self.unk] -= float('inf')即可。

reorder_state = None
batch_idxs = None
for step in range(max_len + 1):  # one extra step for EOS marker
    # reorder decoder internal states based on the prev choice of beams
    if reorder_state is not None:
        if batch_idxs is not None:
            # update beam indices to take into account removed sentences
            corr = batch_idxs - torch.arange(batch_idxs.numel()).type_as(batch_idxs)
            reorder_state.view(-1, beam_size).add_(corr.unsqueeze(-1) * beam_size)
        model.reorder_incremental_state(reorder_state)
        model.reorder_encoder_out(encoder_outs, reorder_state)

    lprobs, avg_attn_scores = model.forward_decoder(tokens[:, :step + 1], encoder_outs)

    lprobs[:, self.pad] = -math.inf  # never select pad
    lprobs[:, self.unk] -= self.unk_penalty  # apply unk penalty

• 在search_step这一行得到具体的分数，indices和beams，search_step就是根据本次decoder出现的概率在inference的时候使用beam search得到的词以及得到的累计的分数score。
  
• 最后generator函数返回的是根据句子的score由大到小排序的一个列表，列表里面是一个字典，字典中有一项是score是这个生成出来句子的得分，就是beam search得到每个位置词概率相加的分数。

# bsz是batch size
# list of completed sentences
finalized = [[] for i in range(bsz)]
finished = [False for i in range(bsz)]
worst_finalized = [{'idx': None, 'score': -math.inf} for i in range(bsz)]
...
# the return of generator in sequence_generator.py
# sort by score descending
for sent in range(len(finalized)):
    finalized[sent] = sorted(finalized[sent], key=lambda r: r['score'], reverse=True)

return finalized

• 找到了生成的概率，所以要看一下search.step生成的是个什么东西了。
  
• score是每个source word截止当前步总的分数。
  
• token起初的时候填充的都是pad，第0个位置是eos。
  
• generate还有一个比较有意思的参数是prefix_tokens用来强制开始时候产生开始的word。
• 往token里面填充新的单词是通过swap buffers来实现的，不知道这么实现是为了什么？？
  
• 想知道scores是干啥的，首先要看search.step里面
  
• 从这里我们可以看到lprobs是bsz ，beam_size，每个句子产生的self.vocab_size的对应于每个次的概率。 score是根据cand_scores而来的，而cand_scores是search.step得到的分数。
  

search.step

首先厘清trainer.py，task和model和train的关系

task

• translation task里面定义了加载数据load_dataset，加载模型的方法load_pretrained_model，setup_task初始化task（就是加载词典）。
• fairseq_task里面定义了build_generator，train_step，valid_step，inference_step，get_batch_iterator。

train.py

• 在train.py里面，使用epoch_itr = task.get_batch_iterator得到训练数据，model = task.build_model(args)得到模型，criterion = task.build_criterion(args)得到训练器。
• 然后将每次更新交给函数train来做，valid和lr调度的时候由具体的函数来做。原来还可以通过args.max_epoch设置训练最多的epoch。

max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf

• 原来subset是这么用的用,隔开测试多个文件，valid_subsets = args.valid_subset.split(',')
• 在train函数里面首先获得训练的数据，train函数是一次训练一个epoch，里面的每个batch的训练是使用trainer来调用的。
  
• 很好奇梯度累计是哪里实现的，trainer.train_step是一次实现一个batch。
• valide的结构和train差不多

trainer.py

• 在这初始化时候将model.half以及model.cuda
  

• trainer里面是通过调用task.train_step来进行前向和后向传递的。
  

• valid也是调用task.valid_step。
  

• 所以说trainer就是task.train_step的一个皮，而task.train_step里面主要是由criterion和optimizer来实现的，梯度累计是在optimizer里面实现的
  

• criterion里面的实现也是先将数据通过model forward以后得到net_output这个应该是logit，然后得到log probs

• 以及得到的net_output = model(**sample['net_input'])是否里面有teacher forcing，有的，得到的net_output直接是decoder每个位置数据的单词以及所有单词的概率，以及一些其他信息比如decoder句子的得分。
  

• 所以说这个target是干啥的target = model.get_targets(sample, net_output).view(-1)

• 因此所有的问题现在回到model上来了。

model

fairseq_model

BaseFairseqModel

• 从BaseFairseqModel首先看一下fairseq_model有什么功能。
• 从FairseqModel，可以看到基础的架构是一个encoder一个decoder。
  
• 所以现在的prev_output_tokens是啥，根据criterion里面的net_output = model(**sample['net_input'])可以看到看就是teacher forcing需要的正确答案，或者是推断的时候之前步的decoder输出。
  

FairseqMultiModel

• 就是多个encoder_decoder结合起来，这里有share-embeddings。
  
• forward也是得到了多个decoder的output。
  
• 可以看到LanguageModel只有一个decoder。
  
• src_tokens是训练的数据[[1,3,5,3],[123,2,2,3],[12,2,3,1]]，得到的是一个句子序列，只不过每个单词的位置是所有词表的概率。
  

TransformerModel

• 可以通过参数args.decoder_embed_path和args.encoder_embed_path加载预训练好的embedding具体是通过args.decoder_embed_path。
  
• seq2seq
  

transformerEncoder

• forward中搭建多层的layer，只用nn.ModuleList然后加入layer的子组件。
  
• 原始的embedding加上位置的embedding然后做dropout。
  
• transformerEncoderLayer，layer_norm加的也太频繁了吧，基本上每做一个操作都要加一个layer_norm。
  
• 根据新的index对encoder_out进行重新排列
  

transformerEncoderLayer

• 关于先residual后layernorm，原始的论文和t2t的不一样，fairseq中的实现是原始论文，但是可以通过参数args.encoder_normalize_before=True来切换到t2t的模式。
  
• 每一层有两个layernorm，layernorm和batchnorm一样是有参数的。
  
  
  
• 原来画了两个真是两层fead_forward，fead_forward后面必跟一个dropout，而且两层fead_forward只有一个有relu，self.relu_dropout和self.dropout数值上有什么不同么？
• 所以说一层有两个子层，第一个子层主要是multi-head attention，第二个子层是两层feed_forward，encoder返回的是出来好的x，就是embedding矩阵。
  
• mult-head attention为了避免乘性attention出现极端值因此进行了scale
  
• 实现qkv的一点细节
  

• 都是来自一个大的矩阵，然后取出自己想要的部分

  
  
• 原来torch里面的parameter可以向操作numpy一样。

self.in_proj_weight = Parameter(torch.Tensor(3 * embed_dim, embed_dim))

• multi-head attention也不是每个头乘以不同的矩阵，而是乘以一个大的矩阵，然后reshape成最后一维度是head_dim,然后乘以同样reshape后的k，
  
• 通过masked_fill将有些位置的attention填充成-inf这样就概率就为0
  
• torch.bmm以后reshape成原始attention的大小，然后过一个linear
  
• 就是一系列的reshape操作，先加大bs维度，得到bs维度很大，但是最后一维是head_dim的矩阵，然后reshape成原始的最后一维是embed_dim矩阵。
  

TransformerDecoderLayer

• transformerDecoder的输入是pre_output_tokens+ position embedding
  
• 而训练中的 pre_output_tokens是真实的数据，采用了teacher forcing的技术。
  
• transformerDecoderLayer的输入，x是已经embedding化的
• decoder 第一部分是一个self attention layer norm和residual
  
• 第二个attention是以自己为query，以encoder的output作为k和v进行attention
  
• 剩下的操作，可以看到和encoder一样也是一个linear加上了relu一个没有加relu。