PyTorch tutorial - BiLSTM CRF 代码解析

本文主要对PyTorch的tutorial之一，Advanced: Making Dynamic Decisions and the Bi-LSTM CRF，进行详细解读，行文顺序上，首先说明一下前面几个辅助函数的作用，然后主体按照Run training的运行顺序进行。(以下删除了原代码注释，可回tutorial中查看)

def argmax(vec):
    _, idx = torch.max(vec, 1)
    return idx.item()

• 给定输入二维序列，在1维度上取最大值，返回对应ID。

def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)

• 利用to_ix这个word2id字典，将序列seq中的词转化为数字表示，包装为torch.long后返回。

def log_sum_exp(vec):
    max_score = vec[0, argmax(vec)]
    max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])
    return max_score + torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))

• 函数目的相当于$\log \sum e^{x_i}$，首先取序列中最大值，输入序列是一个二维序列(shape[1,tags_size])。下面的计算先将每个值减去最大值，再取log_sum_exp，最后加上最大值。具体过程如下：

$$x_m+\log \sum _{i=1}^n{e}^{x_i-x_m}=x_m+\log (e^{-x_m}\sum _{i=1}^n{e}^{x_i})=x_m-x_m+\log \sum _{i=1}^n{e}^{x_i}=\log \sum _{i=1}^n{e}^{x_i}$$

# Author: Robert Guthrie
import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.optim as optim
torch.manual_seed(1)

• 导入torch相关模块，人工设定随机种子以保证相同的初始化参数，实现模型的可复现性。下面开始模型构建及其训练：

START_TAG = ""
STOP_TAG = ""
EMBEDDING_DIM = 5
HIDDEN_DIM = 4

• 首先，设定超参数：在"B", “I”, “O"三个标签的基础上添加了句首标签”“和句尾标签”"；每个词嵌入的编码维度设定为5，LSTM的隐藏层维度设定为4。

training_data = [(
    "the wall street journal reported today that apple corporation made money".split(),
    "B I I I O O O B I O O".split()
), (
    "georgia tech is a university in georgia".split(),
    "B I O O O O B".split()
)]

• 然后，构造训练数据：tutorial中只是针对命名实体识别(NER)任务模拟了两条样例数据。整个训练数据存储为一个列表，其中每条数据是以元组形式存储的语句序列和对应标签序列，语句和标签分别split为词列表和标签列表。

word_to_ix = {}
for sentence, tags in training_data:
    for word in sentence:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)
            
tag_to_ix = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4}

• 接下来，构建词索引表和标签索引表，即数字化以便于计算机处理：为训练数据中所有出现的词添加索引（tutorial只是示例，实际情况下，一般还要针对测试数据中可能出现的未知词以及特殊字符等进行泛化处理）；标签索引直接给定，包含句首句尾。

model = BiLSTM_CRF(len(word_to_ix), tag_to_ix, EMBEDDING_DIM, HIDDEN_DIM)
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)

• 然后进入正式的BiLSTM_CRF模型搭建，并使用随机梯度下降(SGD)对所有参数进行优化，初始学习率设定为0.01，weight_decay表示正则项系数，防止模型过拟合。

class BiLSTM_CRF(nn.Module):
	def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim):
        super(BiLSTM_CRF, self).__init__()
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.tag_to_ix = tag_to_ix
        self.tagset_size = len(tag_to_ix)
        self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2, num_layers=1, bidirectional=True)
        self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)
        self.transitions = nn.Parameter(torch.randn(self.tagset_size, self.tagset_size))
        self.transitions.data[tag_to_ix[START_TAG], :] = -10000
        self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000
        self.hidden = self.init_hidden()
        

• BiLSTM_CRF类的构造函数参数包括词索引表长、标签索引表、词嵌入维度和隐藏层维度，继承torch.nn.Module。私有变量包括输入的词嵌入维度、隐藏层维度、词索引表大小、标签索引表、标签索引表大小（即标签个数）、词嵌入（相当于一个[词索引表长，词嵌入维度]的矩阵，这里是调用nn的Embedding模块初始化的）、LSTM网络层（直接调用的nn中的LSTM模块，设定为单层双向，隐藏层维度设定为指定维度的一半，以便于后期双向拼接）、处理LSTM输出的全连接层（维度变更）、CRF的转移矩阵（T[i,j]表示从j标签转移到i标签，不可能转移到句首标签，也不可能从句尾标签开始转移，因此都设定为极小值）。

# class BiLSTM_CRF		
	def init_hidden(self):
        return (torch.randn(2, 1, self.hidden_dim // 2),
                torch.randn(2, 1, self.hidden_dim // 2))

• 这里是使用随机正态分布初始化LSTM的$h_0$和$c_0$ ，否则，模型自动初始化为零值（维度为[num_layers*num_directions, batch_size, hidden_dim]）。类的构造及初始化结束，接下来回到主函数。

with torch.no_grad():
    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
    precheck_tags = torch.tensor([tag_to_ix[t] for t in training_data[0][1]], dtype=torch.long)
    print(model(precheck_sent))
    

• 在no_grad模式下进行前向推断的检测（torch.no_grad()作用是暂时不进行导数的计算，目的在于减少计算量和内存消耗），取第一条数据语句序列及其对应的标签序列，分别根据word_to_ix和tag_to_ix进行数字化。最后调用BiLSTM_CRF类的forward函数，将数字化的语句序列送入模型进行前向计算。

# class BiLSTM_CRF	
	def forward(self, sentence):  
        lstm_feats = self._get_lstm_features(sentence)
        score, tag_seq = self._viterbi_decode(lstm_feats)
        return score, tag_seq

# class BiLSTM_CRF	
	def _get_lstm_features(self, sentence):
        self.hidden = self.init_hidden()
        embeds = self.word_embeds(sentence).view(len(sentence), 1, -1)
        lstm_out, self.hidden = self.lstm(embeds, self.hidden)
        lstm_out = lstm_out.view(len(sentence), self.hidden_dim)
        lstm_feats = self.hidden2tag(lstm_out)
        return lstm_feats

• 初始化$h_0$和$c_0$，使用之前构造的词嵌入为语句中每个词（word_id）生成向量表示，并将shape改为[seq_len, 1(batch_size), embedding_dim]；LSTM网络根据输入的词向量和初始状态$h_0$和$c_0$，计算得到输出结果lstm_out和最后状态$h_n$和$c_n$。这里用的双向LSTM，lstm_out的shape为[seq_len, 1, (self.hidden_dim//2)*2]，hidden_dim单独作为一个维度，整形成二维，以便于送入hidden2tag中，转换为词-标签([seq_len, tagset_size])表，可以看作每个词被标注为对应标签的得分情况，即维特比算法中的发射矩阵。接下来将该得分返回forward函数，使用维特比算法进行解码，计算最优序列。

# class BiLSTM_CRF	
	def _viterbi_decode(self, feats):
        backpointers = []
        init_vvars = torch.full((1, self.tagset_size), -10000.)
        init_vvars[0][self.tag_to_ix[START_TAG]] = 0
        forward_var = init_vvars
        for feat in feats:
            bptrs_t = []
            viterbivars_t = []
            for next_tag in range(self.tagset_size):
                next_tag_var = forward_var + self.transitions[next_tag]
                best_tag_id = argmax(next_tag_var)
                bptrs_t.append(best_tag_id)
                viterbivars_t.append(next_tag_var[0][best_tag_id].view(1))
            forward_var = (torch.cat(viterbivars_t) + feat).view(1, -1)
            backpointers.append(bptrs_t)
        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
        best_tag_id = argmax(terminal_var)
        path_score = terminal_var[0][best_tag_id]
        best_path = [best_tag_id]
        for bptrs_t in reversed(backpointers):
            best_tag_id = bptrs_t[best_tag_id]
            best_path.append(best_tag_id)
        start = best_path.pop()
        assert start == self.tag_to_ix[START_TAG]
        best_path.reverse()
        return path_score, best_path

• 维特比算法使用动态规划来求解最优路径，根据前向状态和转移矩阵、发射矩阵计算当前状态。backpointers存储每一步的转移路径，便于回溯最优路径。如图1，初始化前向状态，除句首标签对应状态置为0以外，其他状态都置为极小值-10000，保证第一步从START_TAG开始转移。
• 接下来对每个词计算当前状态，独立地看，每个词可以被打作任意一个标签，因此需要计算当前词被打作每一个标签的得分，然后取其最优情况。整体过程如图1所示，将前向状态与当前词的每一个可能标签对应的转移矩阵相加，并对每个标签计算最高得分，再与发射矩阵相加，即得到当前状态，作为下一个词的前向状态。
• 具体实现过程，tutorial中使用bptrs_t记录当前词对应每个标签的最优转移结点，viterbivars_t与bptrs_t相对应，记录对应的最优值。forward_var表示每个标签的前向状态得分，即上一个词被打作每个标签的对应得分值；transitions[next_tag]表示每个标签转移到next_tag的转移得分（即图中转移矩阵的一行）二者相加即得到当前词被打作next_tag的所有可能得分，取其最优值作为当前词被打作这一个标签的最终转移得分，存入viterbivars_t，并将转移结点存入bptrs_t。对每个标签都进行上述计算，求得当前词被打作每个标签的最优得分，将其与当前词的发射矩阵feat相加，得到当前状态，即下一个词的前向状态。
• 对序列中的每个词按上述过程计算，并在最后加上转移到句尾标签STOP_TAG的转移得分，得到terminal_var，表示基于最后一个词的标签的整个序列标注的所有可能总得分情况，取最优值，即为整个序列最终标注得分path_score。根据过程中存储的转移路径结点，反推最优转移路径。最初的转移结点一定是人为构建的START_TAG，删除，并根据这一点确认路径正确性，最后将路径倒序即得到从头开始的最优转移路径best_path。前向计算的测试结束，回到主函数。

for epoch in range(300): 
    for sentence, tags in training_data:
        model.zero_grad()
        sentence_in = prepare_sequence(sentence, word_to_ix)
        targets = torch.tensor([tag_to_ix[t] for t in tags], dtype=torch.long)
        loss = model.neg_log_likelihood(sentence_in, targets)
        loss.backward()
        optimizer.step()

• 300轮迭代训练开始：这里相当于batch_size=1，每次训练一组数据（seq，tag）。每次开始前将上一轮的参数梯度清零，防止累加影响。seq、tags分别数字化为sentence_in、targets。送入neg_log_likelihood函数计算loss，最后反向传播梯度，更新参数值。

# class BiLSTM_CRF	
	def neg_log_likelihood(self, sentence, tags):
        feats = self._get_lstm_features(sentence)
        forward_score = self._forward_alg(feats)
        gold_score = self._score_sentence(feats, tags)
        return forward_score - gold_score

• 这里需要先说明一下loss的定义。CRF是判别式模型，根据当前输入序列计算最优输出序列，即计算条件概率$P(y|x)=\frac{exp[Score(x,y)]}{{\sum }_{y^{\prime }}exp[Score(x,y^{\prime })]}$，分子基于准确的标注路径计算当前特征得分，为gold_score；分母是所有可能的标注序列总得分。取条件概率$P(y|x)$的负对数作为loss，则最大化条件概率等价于最小化loss。因此，模型得loss为：

$$loss=-\log P(y|x)=\log (\sum _{y^{\prime }}{e}^{Score(x,y^{\prime })})-Score(x,y)$$

​ 其中，前一项即为forward_score，后一项为gold_score，下面分别计算这两项。

# class BiLSTM_CRF	
	def _forward_alg(self, feats):
        init_alphas = torch.full((1, self.tagset_size), -10000.)
        init_alphas[0][self.tag_to_ix[START_TAG]] = 0.
        forward_var = init_alphas
        for feat in feats:
            alphas_t = [] 
            for next_tag in range(self.tagset_size):
                emit_score = feat[next_tag].view(1, -1).expand(1, self.tagset_size)
                trans_score = self.transitions[next_tag].view(1, -1)
                next_tag_var = forward_var + trans_score + emit_score
                alphas_t.append(log_sum_exp(next_tag_var).view(1))
            forward_var = torch.cat(alphas_t).view(1, -1)
        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
        alpha = log_sum_exp(terminal_var)
        return alpha

• neg_log_likelihood的第一步和前向计算相同，使用LSTM及线性全连接层计算feats（发射矩阵）。第二步也类似，主要区别有两点：一是将发射矩阵复制为与前向状态和转移矩阵同样的维度，从而可以直接在循环内部相加，取最大值后不需要再处理发射矩阵；二是相加后的得分要再经过log_sum_exp才能作为当前状态。在最后转移到STOP_TAG之后再进行一步log_sum_exp，从而最终得到$\log ({\sum }_{y^{\prime }}{e}^{Score(x,y^{\prime })})$。
• 在这里再说明一下为何要使用log_sum_exp来计算$\log ({\sum }_{y^{\prime }}{e}^{Score(x,y^{\prime })})$。一般句子包括很多词，每个词还对应有多个可能标签，因此，要将所有的路径列举出来挨个计算很困难。使用log_sum_exp就可以在最后得到所有可能路径得分和的形式。
• 举个例子说明一下：假设序列为"机器学习"，可能标签有"B"、"I"和"O"三种。状态矩阵$x_i^j$（j状态下i标签得分）、转移矩阵$y_i^j$（i标签转移到j标签的得分）和发射矩阵$z_i^j$（词j标注为i的得分）分别如下：

状态\标签	B	I	O	START	STOP
START	$x_0^0$	$x_1^0$	$x_2^0$	$x_3^0$	$x_4^0$
机器	$x_0^1$	$x_1^1$	$x_2^1$	$x_3^1$	$x_4^1$
学习	$x_0^2$	$x_1^2$	$x_2^2$	$x_3^2$	$x_4^2$
STOP	$x_0^3$	$x_1^3$	$x_2^3$	$x_3^3$	$x_4^3$

当前\上一步	B	I	O	START	STOP
B	$y_0^0$	$y_1^0$	$y_2^0$	$y_3^0$	$y_4^0$
I	$y_0^1$	$y_1^1$	$y_2^1$	$y_3^1$	$y_4^1$
O	$y_0^2$	$y_1^2$	$y_2^2$	$y_3^2$	$y_4^2$
START	$y_0^3$	$y_1^3$	$y_2^3$	$y_3^3$	$y_4^3$
STOP	$y_0^4$	$y_1^4$	$y_2^4$	$y_3^4$	$y_4^4$

词\标签	B	I	O	START	STOP
机器	$z_0^0$	$z_1^0$	$z_2^0$	$z_3^0$	$z_4^0$
学习	$z_0^1$	$z_1^1$	$z_2^1$	$z_3^1$	$z_4^1$

• 则计算最终状态过程如下：

  初始状态：$(x_0^0,x_1^0,x_2^0,x_3^0,x_4^0)$

  词"机器"的状态：$x_j^1=\log {\sum }_{i=0}^4{e}^{x_i^0+y_i^j+z_j^0},j=0,1,2,3,4$

  词"学习"的状态：$x_k^2=\log {\sum }_{j=0}^4{e}^{x_j^1+y_j^k+z_k^1},k=0,1,2,3,4$

  将$x_j^1$表达式替换到$x_k^2$中，可得：$x_k^2=\log {\sum }_{j=0}^4({\sum }_{i=0}^4{e}^{x_i^0+y_i^j+z_j^0})e^{y_j^k+z_k^1},k=0,1,2,3,4$

  即：$x_k^2=\log {\sum }_{j=0}^4{\sum }_{i=0}^4{e}^{x_i^0+y_i^j+z_j^0+y_j^k+z_k^1},k=0,1,2,3,4$

  最终状态：$x_l^3=\log {\sum }_{k=0}^4{e}^{x_k^2+y_k^l},l=0,1,2,3,4$

  同理，替换可得：$x_l^3=\log {\sum }_{k=0}^4{\sum }_{j=0}^4{\sum }_{i=0}^4{e}^{x_i^0+y_i^j+z_j^0+y_j^k+z_k^1+y_k^l},l=0,1,2,3,4$

  最后再做一步log_sum_exp，可得alpha：

  $alpha=\log {\sum }_{l=0}^4{\sum }_{k=0}^4{\sum }_{j=0}^4{\sum }_{i=0}^4{e}^{x_i^0+y_i^j+z_j^0+y_j^k+z_k^1+y_k^l}$

  每一组$(i,j,k,l)$都对应一条可能路径，即$alpha=\log ({\sum }_{y^{\prime }}{e}^{Score(x,y^{\prime })})$

• By the way，tutorial中也有说过，这部分实际通过矩阵相乘即可计算（可参考李航老师《统计学习方法》中的实例），这里只是为了显式说明具体过程。好了，重头戏结束，回到neg_log_likelihood，接下来计算gold_score。

# class BiLSTM_CRF	
	def _score_sentence(self, feats, tags):
        score = torch.zeros(1)
        tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags])
        for i, feat in enumerate(feats):
            score = score + self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
        score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]
        return score
        

• 初始分值为0，tags前面补上一个句首标签便于计算转移得分，循环用于计算当前的feats在golden路径上的得分，再加上转移到句尾的得分，便得到了$gold\_score=Score(x,y)$，从而得到了模型的loss，然后反向更新参数，完成一次迭代。300代结束以后，回到主函数，进行最后一步前向计算推断，结束训练。

with torch.no_grad():
    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
    print(model(precheck_sent))

• 这里与迭代前的推断测试相同，可以对比训练前后NER的性能。

这样，全篇结束，其实模型虽然名为BI-LSTM CRF，但神经网络部分都利用的集成package，代码主体是在实现CRF部分。建议结合李航老师的《统计学习方法》，可以更深入地理解CRF。

如有代码理解或文字编排等方面的错误，还请批评指正。

OVER！