transformers的Trainer中使用CRF
0. 关于CRF
条件随机场（CRF）是序列标注任务中常用的模型，其基本作用是给定一个序列的特征，对序列中每一个节点的状态进行预测，既可以单独用于序列标注任务，也可以在bert等编码器的基础上，将编码特征作为输入，可以有效地提高序列标注模型的准确性。
在bert4keras中，crf层的调用非常方便，只需要一行代码即可实现，而基于pytorch的最常用的NLP模块transformers的样例中，却并没有给出CRF的使用方法。在transformers提供的Trainer中，也没有预留可以加入CRF的地方，使得用户可以方便的操作（huggingface太追求面面俱到了，原本是为了让使用者利用非常简短的代码就可以实现功能，但是在过度封装的情况下，任何一点小的改动都变得比较困难）。
关于transformers做序列标注，可以直接参考官方文档。
经过了一些实践之后，认为其提供的trainer并不适合开发者使用，更适合不需要修改代码的情况下使用，而如果有自己的模型需要搭建，我个人认为使用trainer并不是明智的选择，包括其中的各种闭环和控制操作都需要对整个模块有一定的理解，才能够熟练的使用。即便如此，我还是决定写这篇博客，希望能给有需要的同学提供一些参考。
1. 下载一个pytorch实现的crf模块
pypi上关于crf的实现模块有很多，在这里推荐Allennlp实现的版本，pytorch-crf，版本是0.7.2，其代码比较清晰，注释也很多，读起来比较容易。关于其他版本的crf实现，我没有做尝试。
下载地址：
https://pypi.org/project/pytorch-crf/#files
联网状态下可直接利用pip安装：
pip install pytorch-crf
引用方法：
from torchcrf import CRF
2. torchcrf的基本使用方法
这个模块的使用非常简便，只需要在两个位置添加。
第一步，在构造方法中创建CRF对象。需要传入一个参数num_labels，也就是标签总数（计入BIO的）。
class model(torch.nn.Model):
def __init__(self, your_args_here, num_labels):
# Your code here.
self.crf = CRF(num_tags=config.num_labels, batch_first=True)
第二步，在forward方法中，计算crf损失。
def forward(self, input_ids, attention_mask, token_type_ids, labels):
# Your code here.
loss = self.crf(logits, labels)
loss = -1 * loss
return loss
注意在这个地方loss要乘以-1，起初我在修改这部分的时候发现损失始终得不到优化，检查crf的代码才发现，计算的两项logsumexp是反着的。
3. 对transformers模块进行修改
熟悉了crf的基本使用方法，我们就可以按照自己的需求取修改transformers模块了。在序列标注任务中，通常使用的默认模型是BertForTokenClassification，所以干脆在这里边加入crf。同时我们希望保留该模型原有的功能，那么就加入一个参数use_crf，用来控制是否使用crf模块。
直接定位到源码位置，开始修改。脚本位于transformers下的models/bert下，名为modeling_bert.py。可以参考下面的路劲找到这个脚本。
vim /root/anaconda3/envs/your_env/lib/python3.6/site-packages/transformers/models/bert/modeling_bert.py
打开之后，先要在头部位置加上引用：
from torchcrf import CRF
直接搜BertForTokenClassification，找到这个类。修改**init**方法：
# 增加了一个参数use_crf来控制是否使用crf
# 如果使用，则创建crf对象
def __init__(self, config, use_crf=False):
super().__init__(config)
self.num_labels = config.num_labels
self.bert = BertModel(config, add_pooling_layer=False)
self.use_crf = use_crf
if self.use_crf:
self.crf = CRF(num_tags=config.num_labels, batch_first=True)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
self.init_weights()
接下来修改它的forward方法：
# 从sequence_output = self.dropout(sequence_output)这一行开始，前面的不用改
# 其实就只是根据self.use_crf的值增加了一个分支
# 如果使用了crf则将损失替换为crf的损失，仅此而已。
sequence_output = self.dropout(sequence_output)
logits = self.classifier(sequence_output)
loss = None
loss_fct = CrossEntropyLoss()
if labels is not None:
# print(logits.shape)
# print(labels.shape)
if self.use_crf:
# 如果使用crf，mask放在crf内部计算
# print(logits.shape)
# print(labels.shape)
if attention_mask is not None:
loss = self.crf(logits, labels, mask=attention_mask.byte())
loss = -1 * loss
# print(loss)
else:
loss = self.crf(logits, labels)
loss = -1 * loss
else:
if attention_mask is not None:
active_loss = attention_mask.view(-1) == 1
active_logits = logits.view(-1, self.num_labels)
active_labels = torch.where(
active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
)
loss = loss_fct(active_logits, active_labels)
else:
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
# 后面的return也还是原样
在这一部分代码中，注意到非crf的分支对attention_mask进行了操作的，而crf的分支中却并没有这一步操作，这是因为torchcrf的代码本身有对attention_mask进行判断的操作，所以无需再这里再算一次。
现在transformers中的模型代码已经修改好了，可以通过以下的方法创建模型：
from transformers import BertForTokenClassification
model = BertForTokenClassification.from_pretrained('bert-base-chinese',
num_labels=len(label_list),
use_crf=True)
但是如果直接把这个model传入trainer进行训练的话，会报CUDA的错误，把device设置成cpu就会发现，其实这个错误是’list out of range’。
这是因为transformers的Trainer类会默认采用-99进行padding，这个值传入crf中之后，crf会按照label去进行索引，-99自然索引不到，所以报了list out of range。
4. 对torchcrf模块进行修改
针对3中最后提到的报错问题，我们打开crf的脚本进行简单的修改。
vim /root/anaconda3/envs/your_env/lib/python3.6/site-packages/torchcrf/__init__.py
找到def _compute_score方法，在如下的位置修改：
# 只添加了一行tags = torch.tensor(np.maximum(tags.cpu().numpy(), 0))
# 道理很简单，就是不让-99溢出，把它替换成0即可
# 至于padding的mask问题不用担心，因为我们传了mask进来
# 代码中有其他地方对mask进行了判断的
def _compute_score(
self, emissions: torch.Tensor, tags: torch.LongTensor,
mask: torch.ByteTensor) -> torch.Tensor:
# emissions: (seq_length, batch_size, num_tags)
# tags: (seq_length, batch_size)
# mask: (seq_length, batch_size)
assert emissions.dim() == 3 and tags.dim() == 2
assert emissions.shape[:2] == tags.shape
assert emissions.size(2) == self.num_tags
assert mask.shape == tags.shape
assert mask[0].all()
seq_length, batch_size = tags.shape
mask = mask.float()
# seq_ends = mask.long().sum(dim=0) - 1
# 导致报错的原因是tags里边padding使用-100
tags = torch.tensor(np.maximum(tags.cpu().numpy(), 0))
# 剩下的部分维持原样
经过这些修改，训练过程就可以正常进行了。
5. 关于评估
训练的问题解决了，那么评估呢？固然可以不进行修改，依然使用argmax的方法进行解码，但是这样一来crf就白用了，而且以CRF的损失进行优化，以argmax进行解码的话，导致两个步骤的不一致性，会令模型的效果变差。
在Trainer中，评估用到的方法叫prediction_loop，这个方法写的比较复杂，把所有batch和step中的所有结果全都堆叠在一起了，然后再重新组织（好像是这个样子，记不太清了），要修改它的话需要很多reshape的操作，比较麻烦，所以我也就没有再继续进行。
如果执意想用Trainer的话，不妨思考一下这个地方具体该怎么修改，或者在callback的控制中，控制其不进行评估，只训练，然后另写评估和预测方法。
考虑到修改源码的时间成本，我没有再继续用trainer进行训练，而是自己另写了一套训练和评估的方法。（其实通过阅读大量的论文和项目可以发现，尽管transformers模块被广泛的使用，但几乎没有人是直接采用了它的Trainer进行训练的）
我把另写的评估方法也贴出来，供看到这里的同学参考。
def evaluate(model, valid_dataloader, id2label, device):
"""
对模型进行评估
---------------
ver: 2021-08-30
by: changhongyu
"""
pred_true = 1e-10
pred_all = 1e-10
true_all = 1e-10
for n, (input_ids, token_type_ids, attention_mask, true_labels) in enumerate(tqdm(valid_dataloader)):
input_ids = input_ids.squeeze(0).to(device)
token_type_ids = token_type_ids.squeeze(0).to(device)
attention_mask = attention_mask.squeeze(0).to(device)
true_labels = true_labels.squeeze(0).to(device)
with torch.no_grad():
eval_loss = model(input_ids, token_type_ids, attention_mask, true_labels)
logits = model(input_ids, token_type_ids, attention_mask) # (b, l, num_labels)
# pred_labels = torch.argmax(logits, dim=-1) # (b, l)
pred_labels = model.crf.decode(logits) # 主要看这一行
# pred_labels = pred_labels.detach().cpu().numpy().tolist()[0] # (l)
pred_labels = pred_labels.detach().cpu().numpy().tolist()[0][0] # (l) # 还有这一行
true_labels = true_labels.detach().cpu().numpy().tolist()
# 这个方法是从labels获取entities的方法
# 不同的数据形式对应不一样的结构
gold_ents = get_entities_from_labels(true_labels, id2label)
pred_ents = get_entities_from_labels(pred_labels, id2label)
# if n % 30 == 1:
# print(gold_coarse)
# print(pred_coarse)
# print('---------')
pred_all += len(pred_ents)
true_all += len(gold_ents)
for ent in pred_ents:
if ent in gold_ents:
pred_true += 1
precision = pred_true / pred_all
recall = pred_true / true_all
f1 = 2 * precision * recall / (precision + recall)
return f1, precision, recall
总结一下，CRF是序列标注任务中非常基础，也是非常重要的方法，即便现在的很多模型中，NER任务不再按照序列标注的形式来进行，但CRF的重要性仍然是不可忽视的，阅读一下源码对理解CRF的运作过程，也会有很大的帮助。transformers模块作为一个便捷的工具，采用trainer可以是初学者很快上手，但我认为并不适合开发者在其基础上进行大量的修改，除非已经对transformers的整个流程和逻辑已经习惯了且非常熟练，如果源码都没有从头到尾读一遍的话，还是自己设计自己的流程吧，不然遇到问题了再去debug会很麻烦。
如果这篇文章对你有帮助的话，麻烦点个赞投个币支持一下。我们下期再见。