CRF的简洁实现

提供基于Tensorflow2.x的CRF简洁实现，并提供NER、POS、CWS例子。

平时做序列标注任务都是用开源的CRF实现，最近看了若干开源的CRF实现（这里不具体说了^-^）其实都不够简洁，这里提供一种十分清晰简洁的实现，见地址：

https://github.com/allenwind/tensorflow-crf

里面提供序列标注常见任务的例子。

这里的实现还是要依赖tensorflow_addons.text中的相关函数，主要是把这些函数封装成Layer并处理mask、loss、train_step等等。CRF是深度学习流行时代解决序列标注和信息抽取任务中最流行的概率图模型，其相关的理论在过去的文章概率图模型系列（5）：CRF中分析过。接下来简述其实现思路。

实现思路

CRF模型（线性链）可以简单表示为，

$$P_c(y_1,\dots,y_n|x_1,\dots,x_n) = \frac{1}{Z(\boldsymbol{x})}\exp\left( \sum_{i=1}^{n}\Big[t(y_{i-1},y_{i})+h(y_{i};\boldsymbol{x})\Big]\right)$$

其中，$t(y_{i-1},y_{i})$是CRF自带的参数矩阵，用于约束标签的转移；而，$h(y_{i};\boldsymbol{x})$​是上游模型学习的，如CNN、LSTM、BERT，相当于自动完成特征工程，是CRF Layer的输入。因此CRF在Tensorflow（Keras）下实现有两点：

• $h(y_{i};\boldsymbol{x})$，CRF层的输入，可以来自CNN、BERT的特征编码
• $t(y_{i-1},y_{i})$，CRF的参数矩阵，用于对标签$y_i$的约束

封装状态转移矩阵和viterbi解码

首先是层的封装，CRF可以看做是包括状态转移矩阵trans和viterbi解码的层，

import tensorflow as tf
import tensorflow_addons as tfa

class CRF(tf.keras.layers.Layer):

    def __init__(self, lr_multiplier=1, trans_initializer="glorot_uniform", **kwargs):
        super(CRF, self).__init__(**kwargs)
        self.supports_masking = True
        self.lr_multiplier = lr_multiplier
        self.trans_initializer = tf.keras.initializers.get(trans_initializer)

    def build(self, input_shape):
        assert len(input_shape) == 3
        units = input_shape[-1]
        self.trans = self.add_weight(
            name="trans",
            shape=(units, units),
            initializer=self.trans_initializer
        )

    def call(self, inputs, mask=None):
        assert mask is not None
        lengths = tf.reduce_sum(tf.cast(mask, tf.int32), axis=-1)
        viterbi_tags, _ = tfa.text.crf_decode(inputs, self.trans, lengths)
        # (bs, seq_len), (bs, seq_len, units), (bs,), (units, units)
        return viterbi_tags, inputs, lengths, self.trans

当然，viterbi解码也可以写到CRF层外，仅仅包括trans矩阵。不过这样在推断阶段解码在TF图外，性能可能不高。其次，Loss也写到Layer上，不过扩张性不好。

CRF的call方法必须传入相应的mask，根据Tensorflow的mask机制，传入方法包括：

• 手动计算mask并显式传入，见例子
• 设置Masking层
• Embedding层设置mask_zero=True，见例子

CRF的学习率调整

假设CRF的状态转移矩阵参数为$\boldsymbol{\theta}$，那么SGD的优化可以表示为，

$$\boldsymbol{\theta} \leftarrow \boldsymbol{\theta}-\alpha \frac{\partial L(\boldsymbol{\theta})}{\partial \boldsymbol{\theta}}$$

现在把状态转移矩阵作变换，缩小到原来的$\frac{1}{\lambda}$，即$\boldsymbol{\psi} = \frac{\boldsymbol{\theta}}{\lambda}$，对$\psi$作优化有，

$$\begin{align} \boldsymbol{\psi} &\leftarrow \boldsymbol{\psi}-\alpha \frac{\partial L(\boldsymbol{\theta})}{\partial \boldsymbol{\psi}} \newline &\leftarrow \boldsymbol{\psi}-\alpha \frac{\partial L(\boldsymbol{\theta})}{\partial \boldsymbol{\theta}} \frac{\partial \boldsymbol{\theta}}{\partial \boldsymbol{\psi}} \newline &\leftarrow \boldsymbol{\psi} - \alpha\lambda \frac{\partial L(\boldsymbol{\theta})}{\partial \boldsymbol{\theta}} \end{align}$$

其中$L(\boldsymbol{\theta})$的参数是$\boldsymbol{\theta}$而不是$\boldsymbol{\psi}$，因为我们在实现的适合只让优化器看到前者。两边乘以$\lambda$有，

$$\boldsymbol{\theta} \leftarrow \boldsymbol{\theta}-\alpha \lambda^{2} \frac{\partial L(\boldsymbol{\theta})}{\partial \boldsymbol{\theta}}$$

这就意味着，把参数缩小到原来的$\frac{1}{\lambda}$，相当于学习率放大到原来的$\lambda^2$倍。不过，等等，其他优化器也是这样吗？SGD的改进有两个方向，分别是Momentum和自适应学习率。Momentum并不调整学习率本身，但是自适应学习率会涉及到原来学习率$\alpha$的调整，以AdaGrad为例，

$$\begin{aligned} & \boldsymbol{g}_{t}(\theta) = \nabla_{\theta}L_{R_{t-1}}(\theta_{t-1}) \\ & G_{t} = \sum_{\tau = 1}^{t} \boldsymbol{g}_{\tau} \odot \boldsymbol{g}_{\tau} \\ & \theta_{t} = \theta_{t-1} - \gamma_{0} \frac{1}{\sqrt{G_{t} + \varepsilon }} \odot \boldsymbol{g}_{t} \end{aligned}$$

注意到$\frac{1}{\sqrt{G_{t} + \varepsilon }}$带有一个$\lambda$因子，因此，对于自适应学习率的优化器，把参数缩小到原来的$\frac{1}{\lambda}$，相当于学习率放大到原来的$\lambda$倍。日常训练模型最常用的adam优化器，应该使用这种方法来调整层的学习率。

因此，把参数缩小到原来的$\frac{1}{\lambda}$，对于自适应的优化器来说，相当于学习率放大到原来的$\lambda$倍，正好一个倒数关系，而对于SGD来说相当于学习率放大到原来的$\lambda^2$​倍。

对于预训练模型，fine-tune的时候学习率较小，而CRF的参数却是随机初始化，进而导致训练中两者的学习率不匹配，这种情况很必要为CRF层设置更大的学习率。

在实践中，一般默认学习率倍数为1，调优的时候在逐步增大学习率倍数。

CRF训练逻辑

其次，把训练（包含loss）、测试、预测的逻辑分离开来，训练需要考虑loss计算和参数更新，而预测阶段只需要计算scores，并输出最优路径（viterbi path）即可。

class ModelWithCRFLoss(tf.keras.Model):
    """把CRFLoss包装到train_step上，容易做各种扩展如loss、对抗训练"""

    def __init__(self, base, **kwargs):
        super(ModelWithCRFLoss, self).__init__(**kwargs)
        self.base = base
        self.accuracy_fn = tf.keras.metrics.Accuracy(name="accuracy")

    def call(self, inputs):
        return self.base(inputs)

    def train_step(self, data):
        x, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
        with tf.GradientTape() as tape:
            viterbi_tags, lengths, crf_loss = self.compute_loss(
                x, y, sample_weight, training=True
            )
        grads = tape.gradient(crf_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        mask = tf.sequence_mask(lengths, y.shape[1])
        self.accuracy_fn.update_state(y, viterbi_tags, mask)
        results = {"crf_loss": crf_loss, "accuracy": self.accuracy_fn.result()}
        return results

    def test_step(self, data):
        x, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
        viterbi_tags, lengths, crf_loss = self.compute_loss(
            x, y, sample_weight, training=False
        )
        mask = tf.sequence_mask(lengths, y.shape[1])
        self.accuracy_fn.update_state(y, viterbi_tags, mask)
        results = {"crf_loss": crf_loss, "accuracy": self.accuracy_fn.result()}
        return results

    def predict_step(self, data):
        # 指定predict函数在预测时只返回viterbi tags
        x, *_ = tf.keras.utils.unpack_x_y_sample_weight(data)
        viterbi_tags, *_ = self(x, training=False)
        return viterbi_tags

    def compute_loss(self, x, y, sample_weight, training):
        viterbi_tags, potentials, lengths, trans = self(x, training=training)
        crf_loss, _ = tfa.text.crf_log_likelihood(potentials, y, lengths, trans)
        if sample_weight is not None:
            crf_loss = crf_loss * sample_weight
        return viterbi_tags, lengths, tf.reduce_mean(-crf_loss)

注意到，train_step、test_step、predict_step有着不同的逻辑，其中predict_step只返回viterbi解码后的结果。这种写法可以做各种扩展，如在train_step上添加对抗训练，在compute_loss修改具体任务的Loss。

以上是实现的简述，使用起来也十分简单，简单的例子如下，

import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras import *
from crf import CRF, ModelWithCRFLoss

vocab_size = 5000
hdims = 128
inputs = Input(shape=(None,), dtype=tf.int32)
# 设置mask_zero=True获得全局mask
x = Embedding(vocab_size, hdims, mask_zero=True)(inputs)
x = Bidirectional(LSTM(hdims, return_sequences=True))(x)
x = Dense(4)(x)
crf = CRF(trans_initializer="orthogonal")
outputs = crf(x)
base = Model(inputs, outputs)
model = ModelWithCRFLoss(base)
model.summary()
model.compile(optimizer="adam")
X = tf.random.uniform((32*100, 64), minval=0, maxval=vocab_size, dtype=tf.int32)
y = tf.random.uniform((32*100, 64), minval=0, maxval=4, dtype=tf.int32)
model.fit(X, y)
tags = model.predict(X)
print(tags.shape)

是不是很简单。其中设置mask_zero=True获得全局mask，此外还可以手动计算mask以获得更大的灵活性。

开源代码

CRF的实现和例子见tensorflow-crf，包括NER、CWS、POS等例子。

关于mask的两种处理：

• 手动计算mask并传入到CRF：crf_example2.py
• 自动计算mask的例子：crf_example1.py

命名实体识别（NER）：

• task_ner_bilstm_crf.py
• task_ner_cnn_crf.py

中文分词（CWS）：

• task_cws_bilstm_crf.py
• task_cws_cnn_crf.py
• 更多中文分词的实现见chinese-cut-word。

词性标注（POS）：

• task_pos_bilstm_crf.py
• task_pos_cnn_crf.py

深度学习时代，CRF是序列标注中最常用的解码器，而序列标注的思路可以解决很多NLP问题，如信息抽取、机器阅读理解（MRC）文本摘要甚至对对联和对诗句等等，后期有需要和时间再更新~

转载请包括本文地址：https://allenwind.github.io/blog/13527
更多文章请参考：https://allenwind.github.io/blog/archives/