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 = "<START>"
STOP_TAG = "<STOP>"
EMBEDDING_DIM = 5
HIDDEN_DIM = 4

• 首先，设定超参数：在"B", “I”, “O"三个标签的基础上添加了句首标签”<START>“和句尾标签”<STOP>"；每个词嵌入的编码维度设定为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