模型训练检测方案

模型训练监测方案—适用于生成类任务

导读

本文提出一个模型训练辅助工具，适用于文本类任务（NLU、NMT、NLG、NLP等任务线）。

引言

最近在做开发时脑海里突然产生了一个疑惑：神经网络模型的训练过程是不可见的，神经网络的内部结构由全连接或部分连接的神经元组成，或者搭配如Attention、残差、Norm等模块，这些模块工作时涉及了大量的浮点数计算。人类在训练或调试时对这些是不可见的，更不能根据数值变化推断模型的状态。也就是说模型的迭代是自主的、人类不可知也不能干预的，此即算法的“黑盒”特性。对于如何研究神经网络的可解释性，以驱动NN模型能够更好地适配于任务，一直以来都是研究领域的重点课题。那么对于算法开发人员，在神经网络内部不可知的情况下，如何才能确保自己训练的模型是正确的、逐渐变好的呢？本文向大家分享一个训练监测方案，随时随地掌握模型的动态，确保模型正确地收敛（代码已附在最后）。

训练检测方案—参数统计工具

在开始之前，我们先观摩下 Huggingface上的fill-mask测试工具 的显示效果，如下图。图中使用distillroberta进行句子掩码预测任务，当句子’‘Paris is the of France.’'输入模型后，对mask部分给出的top-5预测结果也展示了出来，右边的浮点数对应了每个词的logits概率。能够从中清晰看出模型对这句话的理解，输出中也包含了如“heart”、“Capital”等诸多与“首都”含义相关的词语，由此也能反映模型对语言知识的掌握程度。

这项功能如果运用于日常模型开发中，将会从一个全新的视角深入了解模型的变化过程。因此我这边自主开发了一个简易的训练监测系统，用于对训练过程的Encoder输入、logits输出及统计、labels值、模型预测与labels的相关度等项目进行可视化展示。下面逐一介绍实现过程（以Huggingface的 Tokenizer 框架为例）。

Encoder输入

用于展示当前模型的input_ids所代表的原始内容，即Encoder层的输入内容。我们希望输出结果能够按词切分，即使用space将每个id对应的词隔开，方便对分词进行检查。方法如下：使用tokenizer.batch_decode()算法对input_ids解码，注意batch_decode默认将输入的第一维当作batch维，因此解码结果也按第一维罗列。因此巧妙利用这一特点，将单个句子的ids输入tokenizer分词，得到的结果自然就是以词为单位的列表，效果如下：

sent_ids = tokenizer('在小小的花园里面挖呀挖呀挖。',return_tensors="pt")
# 'input_ids': tensor([[106，41484，38516，10301，18190，3133，18190，3133，18190，10，0]])
# 'attention_mask': tensor([[1，1，1，1，1，1，1，1，1，1，1，1]])
sent_sep_tk = tokenizer.batch_decode(sent_ids['input_ids'][0]), skip_special_tokens=False)
# sent_sep_tk:['在'，'小小的'，'花园'，'里面'，'挖'，'呀'，'挖'，'呀'，'挖'，'</s>']

将上述实例得到的sent_sep_tk使用分隔符拼接起来，得到Encoder输入的分词展示。注意将skip_special_tokens设为False，以便更直观地展现特殊字符的分布。

logits结果

对于模型的logits，可以使用不同方法展示多个方面的内容，本节讨论将logits转换为模型预测结果并输出。我们都知道模型的logits维度为(B×L×V)，其中V是词表大小，其值反映了当前词预测的置信度。因此使用teacher-forcing的词预测作为当前结果（teacher-forcing即每次生成词时使用之前的正确译文引导，是NLG的一种训练策略，当前获得词预测的方案也是类似teacher-forcing），代码实现如下：

logits = model(**features).logits
_, pred_indices = torch.max(logits, -1)  #在维度V上取最大值，变成（B×L）
predict_texts = [tokenizer.batch_decode(x, skip_special_tokens=False) for x in predict_indices]

首先使用max获得teacher-forcing的最大似然子词，在进行解码获得序列。max方法类似于一般分类任务的后处理，batch_decode思路和2.1一样，都是seq_len作第一维度，实现按词存储。

labels

按照2.1的方式解码，需要先还原labels中的ignore_index再解码。

labels = torch.where(labels==-100, 1, labels)
label_texts = [tokenizer.batch_decode(x, skip_special_tokens=False) for x in labels]

拟合差异

拟合差异是本文重点介绍的方案，它可以反映模型输出内容同label之间的差距究竟多大，使用一系列特征来衡量并展示。本节从两个方面展示：①logits中排名top-k的词分布；②labels词在logits中的置信度排名。其中①能够很好地反映当前状态下，模型对于给定样本能做出的假设是什么；②能够反映模型对当前样本的拟合程度，也间接表明从当前开始，使模型能够收敛到正确预测当前标签的难度。分别来介绍下上述实现：

labels_rank = []    
seq_len = logits.shape[1]
for idx in range(logits.shape[0]):
    line = logits[idx] # 由batch级到seq级，维度转变为：（B×L×V）→（L×V）
    line_rank = []
    for pos in range(seq_len):
        idx_top = torch.topk(line[pos], 10)[1]  # 由seq级到token级，topk取当前位置logits排名top-10的索引（分词id）
        pred_top = self.tokenizer.batch_decode(idx_top, skip_special_tokens=False)  
        # 对top-10进行batch_decode，获得前十位词预测分布。
        line_rank.append(((torch.nonzero(line[pos].sort(descending=True)[1]==labels[idx][pos])[0][0] \
        				+ 1).tolist(), pred_top))  
        # line[pos].sort(descending=True)将维度为V的logits按降序排列，获得索引映射后查找label_id，
        # 得到label_id在排序后的logits中的位置
    labels_rank.append(line_rank)

下面来介绍下代码都实现了哪些操作，首先line=logits[idx]映射到单独序列，维度也变为2维。里层的for循环遍历序列中的每一个位置的logits，topk操作选取logits中最大的k个（代码中是10），再对得到的结果按之前的方式解码，即可获取了当前模型状态下teacher-forcing的前十个预测值；真实label在logits中的排名计算则较为麻烦，对应代码中的line_rank. 首先sort方法将logits降序排序（descending=True），并通过索引[1]获得排序后的indices，indices中存放了当前排序后的每个元素在原始张量的位置索引，因此使用indices==label的判断条件搜寻label的位置（nonzero），从而获得排名。由于label也是索引编号，因此使用indices巧妙地一行代码实现排名查询。

再展示模型收敛到一定状态后的效果：

图一是模型训练早期（500步）时，对模型的实时转态检查结果，能够看出早期由于模型欠拟合，并不能有效地预测出真实译文，<模型输出>的内容与参考译文之间相差也很大；<困惑差异>即拟合差异，每一行括号里的第一个数字显示当前参考词在logits中的排名，logits的长度为词表大小，后面显示模型对当前位置的预测置信度前10位的词列表。排名数据显示预测的总体排名都靠近中后位，表明参数的随机性较大，不具备符合语义特征的分布；而top-10词列表中也包含了大多数的简单词如“的”、“是”等，也表明了现阶段模型也只具备了这些高频或区分性不大的词的认知能力，有读者也能看出这些也都符合停用词特征。

到了模型训练的中期，即15k步以上时，由图中可以看出模型输出基本能够符合参考译文的分布。而拟合差异中的排名数据也都位于前列，多个位置的词都排在第一位，表明该词模型已经能正确预测；而后面top-10词列表也表现了很多相关性，例如‘能’后面跟了“能够”、“可以”等含义相近的词，“到”后面跟了“至”、“去”等，它们的含义与目标词也非常接近。上述分析表明了当前时期模型已经具备了一定的语言能力。

潜在能力 上面说到，模型训练后期时，top-k检查结果中不仅能几乎包含目标词，而且其它词也都是和目标词同义或近义的，这也是由于神经网络模型训练时具有的Word2Vec特性，即相似上下文的词表示向量也在空间接近。我在《机器翻译专用词开发实践》中探讨词约束训练时也分析了这一特性，当时提出了 “词类” 说法，即模型对每个位置词的预测是符合词类特性的。如果某个词成为目标词，则其含义相近的词也具有较大的logits， 现在这一特性能够清晰地从top-k分布中验证。这方面的潜在能力是有助于挖掘近义词等任务，尤其是对于医学等专业领域的任务，近义词库稀缺的情况下可以采用已有模型的top查询方式检索潜在词；同时也能做数据增强，因为挖掘到的近义词都对应了句子中的位置，可以直接替换目标词，从而产生句法、语法尽可能合理的增广句子。

其他说明

在什么情况下使用

目前总结出的这项功能的应用背景有：

• 需要了解数据处理的是否正确时（数据是否正确、输入与标签是否对齐、正则化等预处理是否有效、分词是否符合预期）；
• 需要了解模型训练是否有效时（预测结果与真实标签的差异）；
• 需要掌握目前模型收敛到什么程度时，top-k反映当前模型对输入做出的假设，label排名反映模型与预期能力相差多大（如前文所说，训练初期top-k有很多无关词，label排名靠后；后期则top-k明显符合目标词及其近义词分布，label排名也多在个位数以内）。

如何使用

train_dataloader = DataLoader(...)
model = Model(...)
n = 0
for features in train_dataloader:
	if (n+1) % 500 == 0:   # 设置每500步检查一次
		model(**features).logits   # 获取logits值
		input_texts, predict_texts, label_texts, labels_rank \ 
							= monitor_run(features, logits)
		...  # 整理结果、输出等
	...  # 训练、验证等代码

将实现脚本加在训练代码的迭代部分，设置判断条件实现每隔m步做一次状态检查，代码中设置为500步，使用(n+1)是为了不需要训练一开始就检查。

与模型真实值之间是否存在偏差？

偏差肯定是存在的，原因有两方面：①检查时模型依然是train模式，即dropout或BN等还处于训练状态，因此结果会与模型真正预测时不一致；②如果是生成类模型，则这种teacher-forcing的检查方式与测试时的beam-search等方法也不一致，同样会带来偏差。上述偏差因素中①的检查效果与测试效果好坏无法对比，②的检查效果会略好于测试效果，因为beam-search等算法会存在exposure bias，一旦前期发生偏差会传播到后续推理步骤。 因此总体而言检查得到的模型能力会高于实际的能力， 但是本方案提供的信息还是能反映模型整体水平，通过这种检查还是能大致掌握模型的训练是否有效。

注意事项

1、batch_size不能太大

检查模式下batch_size一定做好把控，建议在10以内，因为模型需要返回logits，其包含了数以亿计的浮点数，在训练模式（带grad）下返回logits容易显存溢出。

2、Special_token问题

Special_token的处理对于huggingface开源分词器而言可以设置skip_special_tokens=False，而对于其它分词器则需要具体配置。这里以Tokenizer的BPE词表而言，需要读取词表中的特殊字符，并建立映射。代码如下：

    def load_special(self):
        with open(self.params.tgt_bpe, 'r') as f:        
        	zh_vocab = json.load(f)
        self.mapper = dict([(x['id'], x['content']) for x in zh_vocab['added_tokens']])  
        				# 'added_tokens'为词表中的键名，保存了所有special_tokens，
        				# 它们在被Tokenzier直接解码时会跳过不显示

修改decode脚本如下：

    def decode_thumt_node(self, seq):
        seq = seq.tolist()
        ret = []
        for i in seq:
            if i == 0:
                break
            if i < len(self.mapper):  # 这里假设Special_tokens都是位于词表前列，
            						  #因此判断id是否在前n项内，如果是则直接获取对应的Token.
                cont = self.mapper[i]
            else:
                cont = self.tokenizer.decode([i])
            if cont.startswith('##'):  # 去除BPE解码时的前缀，以简洁地显示
                cont = cont[2:]
            ret.append(cont)
        return ret  

最终完整代码如下：

from transformers import AutoTokenizer, AutoConfig
import torch
from utils.segmenter import ThumtSegmenter
from nltk.translate.bleu_score import sentence_bleu

class Monitor_Run(torch.nn.Module):
    def __init__(self, params) -> None:
        super().__init__()
        self.params = params
        self.tokenizer = AutoTokenizer.from_pretrained(params.model_dir)
        self.eos_tk = self.tokenizer.convert_ids_to_tokens(AutoConfig.from_pretrained(params.model_dir).eos_token_id)
        self.segmenter = ThumtSegmenter()
    
    def forward(self, src_enc, att_mask, logits, labels):
        labels = torch.where(labels == -100, 1, labels)
        input_texts = [' '.join(self.tokenizer.batch_decode(x, skip_special_tokens=False)[:torch.sum(att_mask[id])]) for id, x in enumerate(src_enc)]
            
        _, predict_indices = torch.max(logits, -1)
        predict_texts = [self.tokenizer.batch_decode(x, skip_special_tokens=False) for x in predict_indices]
        
        labels_rank = []    
        seq_len = logits.shape[1]
        for idx in range(logits.shape[0]):
            line = logits[idx]
            line_rank = []
            for pos in range(seq_len):
                idx_top = torch.topk(line[pos], 10)[1]
                pred_top = self.tokenizer.batch_decode(idx_top, skip_special_tokens=False)
                line_rank.append(((torch.nonzero(line[pos].sort(descending=True)[1] == labels[idx][pos])[0][0] + 1).tolist(), pred_top))
            labels_rank.append(line_rank)
            
        label_texts = [self.tokenizer.batch_decode(x, skip_special_tokens=False) for x in labels]
        for i in range(len(label_texts)):
            label_texts[i] = label_texts[i][:label_texts[i].index(self.eos_tk)+1]
        for i in range(len(labels_rank)):
            label_len = len(label_texts[i])
            labels_rank[i] = list(zip(label_texts[i], predict_texts[i][:label_len], labels_rank[i][:label_len]))
            predict_texts[i] = ' '.join(predict_texts[i][:label_len])
            label_texts[i] = ' '.join(label_texts[i])
        
        bleus = []
        for pred, label in zip(predict_texts, label_texts):
            bleus.append(str(round(self.compute_sent_bleu(label, pred), 3)))
        
        return input_texts, predict_texts, label_texts, labels_rank, bleus
    
    def compute_sent_bleu(self, ref, res):
        seg_str = self.segmenter.seg_line_str(ref)
        seg_list = seg_str.split()
        seg_res_str = self.segmenter.seg_line_str(res)
        seg_list_res = seg_res_str.split()
        return sentence_bleu([seg_list], seg_list_res)

写在最后

总体而言，本文提出的监测方案在实现上并不复杂，代码开发量也很少。但是这套功能确实在模型训练中扮演了重要角色，帮助我排查了很多异常问题，因此欢迎有兴趣的小伙伴们尝试使用。