LSTM-CRF实战

引言

在实现了CRF之后，本文我们来看一下如何应用它去做一个简单的命名实体识别。

参考了PyTorch官方教程： ADVANCED: MAKING DYNAMIC DECISIONS AND THE BI-LSTM CRF。

导入包

import torch
import torch.nn as nn
import torch.optim as optim

# 设置随机种子
torch.manual_seed(47)

定义CRF模型

把CRF模型的实现拷贝过来，详细的图解在CRF实现。

# CRF Model
class CRF(nn.Module):
    """
    条件随机场的实现。
    使用前向-后向算法计算输入的对数似然。参考论文 Neural Architectures for Named Entity Recognition 。
    基于Python3.10.9和torch-1.13.1

    """

    def __init__(
        self,
        num_tags: int,
        batch_first: bool = False,
    ) -> None:
        """初始化CRF的参数

        Args:
            num_tags (int): 标签数量
            batch_first (bool, optional): 是否batch维度在前，默认为False
        """
        super().__init__()
        self.num_tags = num_tags
        self.batch_first = batch_first

        # 转移分数(A)，表示两个标签之间转移的得分
        # transitions[i,j] 表示由第i个标签转移到第j个标签的得分(可以理解为可能性/置信度)
        self.transitions = nn.Parameter(torch.Tensor(num_tags, num_tags))
        # 新引入了两个状态：start和end
        # 从start状态开始转移的分数
        self.start_transitions = nn.Parameter(torch.Tensor(num_tags))
        # 转移到end状态的分数
        self.end_transitions = nn.Parameter(torch.Tensor(num_tags))

        self.reset_parameters()
        
            
    def __repr__(self):
        return f"CRF(num_tags={self.num_tags})"

    def reset_parameters(self) -> None:
        nn.init.uniform_(self.transitions, -0.1, 0.1)
        nn.init.uniform_(self.start_transitions, -0.1, 0.1)
        nn.init.uniform_(self.end_transitions, -0.1, 0.1)

    def forward(
        self,
        emissions: torch.Tensor,
        tags: torch.LongTensor,
        mask: torch.ByteTensor | None = None,
        reduction: str = 'sum'
    ) -> torch.Tensor:
        """计算给定的标签序列tags的负对数似然

        Args:
            emissions (torch.Tensor):  发射分数P 形状 (seq_len, batch_size, num_tags), 代表序列中每个单词产生每个标签的得分
            tags (torch.LongTensor): 标签序列 如果batch_first=False 形状 (seq_len, batch_size) ，否则 形状为 (batch_size, seq_len)
            mask (torch.ByteTensor | None, optional): 表明哪些元素为填充符，和tags的形状一致。  如果batch_first=False  形状 (seq_len, batch_size) ，否则 形状为 (batch_size, seq_len)
                默认为None，表示没有填充符。
            reduction (str): 汇聚函数： none|sum|mean|token_mean 。 none：不应用汇聚函数； 

        Returns:
            torch.Tensor: 输入tags的负对数似然
        """

        if mask is None:
            # mask 取值为0或1，这里1表示有效标签，默认都为有效标签
            mask = torch.ones_like(tags)

        if self.batch_first:
            # 转换为seq维度在前的形式
            emissions = emissions.permute(0, 1)
            tags = tags.permute(0, 1)
            mask = mask.permute(0, 1)

        # 计算标签序列tags的得分
        score = self._compute_score(emissions, tags, mask)
        # 计算划分函数 partition Z(x)
        partition = self._compute_partition(emissions, mask)
        # negative log likelihood 
        nllh = partition - score
        
        if reduction == 'none':
            return nllh
        if reduction == 'sum':
            return nllh.sum()
        if reduction == 'mean':
            return nllh.mean()
        # 否则为 'token_mean'
        return nllh.sum() / mask.type_as(emissions).sum()
    

    def _compute_score(
        self,
        emissions: torch.Tensor,
        tags: torch.LongTensor,
        mask: torch.ByteTensor,
    ) -> torch.Tensor:
        """计算标签序列tags的得分


        Args:
            emissions (torch.Tensor): 发射分数P 形状 (seq_len, batch_size, num_tags)
            tags (torch.LongTensor): 标签序列 形状 (seq_len, batch_size)
            mask (torch.ByteTensor): 表明哪些元素为填充符 形状 (seq_len, batch_size)

        Returns:
            torch.Tensor: 批次内标签tags的得分， 形状(batch_size,)
        """

        seq_len, batch_size = tags.shape
        # first_tags (batch_size,)
        first_tags = tags[0]

        # 由start标签转移到批次内所有标签序列第一个标签的得分
        score = self.start_transitions[first_tags]
        # 加上 批次内第一个(index=0)发射得分，即批次内第0个输入产生批次内对应第0个标签的得分

        score += emissions[0, torch.arange(batch_size), first_tags]

        mask = mask.type_as(emissions)  # 类型保持一致
        # 这里的index从1开始，也就是第二个位置开始
        for i in range(1, seq_len):
            # 第i-1个标签转移到第i个标签的得分 + 第i个单词产生第i个标签的得分
            # 乘以mask[i]不需要计算填充单词的得分
            # score 形状(batch_size,)
            score += (
                self.transitions[tags[i - 1], tags[i]]
                + emissions[i, torch.arange(batch_size), tags[i]]
            ) * mask[i]

        # last_tags = tags[-1] × 这是错误的！，因为可能包含填充单词
        valid_last_idx = mask.long().sum(dim=0) - 1  # 有效的最后一个索引
        last_tags = tags[valid_last_idx, torch.arange(batch_size)]

        # 最后加上最后一个标签转移到end标签的转移得分
        score += self.end_transitions[last_tags]
        return score

    def _compute_partition(
        self, emissions: torch.Tensor, mask: torch.ByteTensor
    ) -> torch.Tensor:
        """利用CRF的前向算法计算partition的分数

        Args:
            emissions (torch.Tensor): 发射分数P 形状 (seq_len, batch_size, num_tags)
            mask (torch.ByteTensor): 表明哪些元素为填充符  (seq_len, batch_size)

        Returns:
            torch.Tensor: 批次内的partition分数 形状(batch_size,)
        """

        seq_len = emissions.shape[0]
        # score (batch_size, num_tags) 对于每个批次来说，第i个元素保存到目前为止以i结尾的所有可能序列的得分
        score = self.start_transitions.unsqueeze(0) + emissions[0]

        for i in range(1, seq_len):
            # broadcast_score: (batch_size, num_tags, 1) = (batch_size, pre_tag, current_tag)
            # 所有可能的当前标签current_tag广播
            broadcast_score = score.unsqueeze(2)
            # 广播成 (batch_size, 1, num_tags)
            # shape: (batch_size, 1, num_tags)
            broadcast_emissions = emissions[i].unsqueeze(1)
            # (batch_size, num_tags, num_tags) = (batch_size, num_tags, 1) + (num_tags, num_tags) + (batch_size, 1, num_tags)
            current_score = broadcast_score + self.transitions + broadcast_emissions
            # 在前一时间步标签上求和  -> (batch_size, num_tags)
            # 对于每个批次来说，第i个元素保存到目前为止以i结尾的所有可能标签序列的得分
            current_score = torch.logsumexp(current_score, dim=1)
            # mask[i].unsqueeze(1) -> (batch_size, 1)
            # 只有mask[i]是有效标签的current_score才将值设置到score中，否则保持原来的score
            score = torch.where(mask[i].bool().unsqueeze(1), current_score, score)

        # 加上到end标签的转移得分 end_transitions本身保存的是所有的标签到end标签的得分
        # score (batch_size, num_tags)
        score += self.end_transitions
        # 在所有的标签上求(logsumexp)和
        # return (batch_size,)
        return torch.logsumexp(score, dim=1)

    def decode(
        self, emissions: torch.Tensor, mask: torch.ByteTensor = None
    ) -> list[list[int]]:
        """使用维特比算法找到最有可能的序列

        Args:
            emissions (torch.Tensor):  发射分数P 形状 (seq_len, batch_size, num_tags), 代表序列中每个单词产生每个标签的得分
            mask (torch.ByteTensor | None, optional): 表明哪些元素为填充符。  如果batch_first=False  形状 (seq_len, batch_size) ，否则 形状为 (batch_size, seq_len)
                默认为None，表示没有填充符。

        Returns:
            list[list[int]]: 批次内的最佳标签序列
        """
        if mask is None:
            mask = torch.ones(emissions.shape[:2], dtype=torch.uint8)

        if self.batch_first:
            # 转换为seq维度在前的形式
            emissions = emissions.permute(0, 1)
            mask = mask.permute(0, 1)

        return self._viterbi(emissions, mask)

    def _viterbi(
        self, emissions: torch.Tensor, mask: torch.ByteTensor = None
    ) -> list[list[int]]:
        """维特比算法的实现

        Args:
            emissions (torch.Tensor): 发射分数P 形状 (seq_len, batch_size, num_tags)
            mask (torch.ByteTensor): 表明哪些元素为填充符 形状 (seq_len, batch_size)

        Returns:
            list[list[int]]: 批次内的最佳标签序列
        """
        seq_len, batch_size = mask.shape
        # 由start到当前时间步所有标签的转移得分 + 批次内所有当前时间步产生所有标签的发射得分
        # (num_tags,) +  (batch_size, num_tags) -> (batch_size, num_tags)
        # score 形状 (batch_size, num_tags) 保存了当前位置，到每个tag的最佳累计得分： 前一累计得分+转移得分+发射得分
        score = self.start_transitions + emissions[0]
        # 保存了目前为止到当前时间步所有标签的最佳候选路径 最终有seq_len-1个(batch_size, num_tags)的Tensor
        history: list[torch.Tensor] = []

        for i in range(1, seq_len):
            # broadcast_score: (batch_size, num_tags, 1) = (batch_size, pre_tag, current_tag)
            # 所有可能的当前标签current_tag广播
            broadcast_score = score.unsqueeze(2)
            # 广播成 (batch_size, 1, num_tags)
            # shape: (batch_size, 1, num_tags)
            broadcast_emissions = emissions[i].unsqueeze(1)
            # (batch_size, num_tags, num_tags) = (batch_size, num_tags, 1) + (num_tags, num_tags) + (batch_size, 1, num_tags)
            current_score = broadcast_score + self.transitions + broadcast_emissions
            # 计算前一时间步到当前时间步的某个标签的最佳得分
            # best_score 形状 (batch_size, num_tags) indices 形状 (batch_size, num_tags)
            # indices 相当于 torch.argmax(current_score, dim=1) ，得到产生最大值对应的索引
            best_score, indices = torch.max(current_score, dim=1)

            # mask[i].unsqueeze(1) -> (batch_size, 1)
            # 只有mask[i]是有效标签的best_score才将值设置到score中，否则保持原来的score
            score = torch.where(mask[i].bool().unsqueeze(1), best_score, score)
            # 记录得到最佳得分的前一个索引
            history.append(indices)

        # 加上到end标签的转移得分 end_transitions本身保存的是所有的标签到end标签的得分
        # score (batch_size, num_tags)
        score += self.end_transitions
        # 计算出最后一个时间步到end标签的最大得分 以及对应的索引(tag)
        # best_score 形状(batch_size,)  indices 形状(batch_size,)
        best_score, indices = torch.max(score, dim=1)
        # 序列最后有效标签的个数
        seq_end_tags = mask.long().sum(dim=0) - 1
        # 保存需要返回的结果
        best_paths: list[list[int]] = []

        # 因为批次内每个样本的最后一个有效标签可能不同，因此需要写成for循环
        for i in range(batch_size):
            best_last_tag = indices[i]
            # 通过item()变成普通int
            this_path = [best_last_tag.item()]
            # history 有 seq_len-1个(batch_size, num_tags)， 但是是顺序添加的，history[: seq_end_tags[i]]取有效的history，再逆序
            for hist in reversed(history[: seq_end_tags[i]]):
                # 先取批次内第i个样本的路径，再取到this_path[-1]，即best_last_tag的最佳标签
                best_last_tag = hist[i][this_path[-1]]
                # 转换为int加入到最佳路径中
                this_path.append(best_last_tag.item())
            # 这里是通过回溯的方法添加的最佳路径，因此最后还需要reversed逆序，变回顺序的，存入best_tags列表
            this_path.reverse()
            best_paths.append(this_path)

        return best_paths

定义BiLSTM-CRF模型

这里定义的比较简单，一个BiLSTM网络加上上面实现的CRF层。词向量是随机初始化的。

# Bi-LSTM CRF
class BiLSTMCRF(nn.Module):
    def __init__(self, vocab_size, tag_to_idx, embedding_dim, hidden_dim):
        super().__init__()
        
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.tag_to_idx = tag_to_idx
        # 标签集大小
        self.tagset_size = len(tag_to_idx) 
        
        self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
        # 双向LSTM 拼接双向的结果后就变成了hidden_dim
        self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2, num_layers=1, bidirectional=True)
        # 映射LSTM的输出到标签空间
        self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)
        
        self.crf = CRF(num_tags=self.tagset_size)
        
    def _get_lstm_features(self, sentences):
        # sentence （seq_len， batch_size）
        # embeds   (seq_len， batch_size, embed_dim)
        embeds = self.word_embeds(sentences)
        # lstm_out (seq_len, batch_size, hidden_dim )
        lstm_out, hidden = self.lstm(embeds)

        # 转换成标签大小
        # lstm_feats (seq_len, batch_size, tagset_size)
        lstm_feats = self.hidden2tag(lstm_out)
        return lstm_feats
    
    def forward(self, sentence, targets, mask):
        """
        前向传播，计算损失
        
        """
        # lstm_feats (seq_len, batch_size, tagset_size)
        # sentence （seq_len， batch_size）
        lstm_feats = self._get_lstm_features(sentence)
        # (batch_size,)
        return self.crf.forward(lstm_feats, targets, mask) 
    
    def inference(self, sentence, mask=None):
        """
        利用维特比算法进行推理
        """
        lstm_feats = self._get_lstm_features(sentence)
        tag_seq = self.crf.decode(lstm_feats, mask)
        return tag_seq

        
        

数据预处理

这里用中文进行测试，最小单位为“字”。

Python中可以通过下面方法快速把中文字符串拆分成字的列表：

list("中国人")

['中', '国', '人']

下面定义一个简单的数据集，并构建相应的词典：

PAD_TOKEN = ""


# 构造一些训练数据
training_data = [(
    list("小明在北京大学的燕园看了中国男篮的一场比赛"),
    "B-PER I-PER O B-ORG I-ORG I-ORG I-ORG O B-LOC I-LOC O O B-ORG I-ORG I-ORG I-ORG O O O O O".split()
), (
    list("小红参加了学校组织的校园歌手大赛"),
    "B-PER I-PER O O O B-ORG I-ORG O O O O O O O O O".split()
)]

word_to_ix = {PAD_TOKEN: 0}
tag_to_ix = {}

for sentence, tags in training_data:
    for word, tag in zip(sentence, tags):
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)
        if tag not in tag_to_ix:
            tag_to_ix[tag] = len(tag_to_ix)
    

可以看一下这两个词典是怎么样的：

word_to_ix

{'': 0,
 '小': 1,
 '明': 2,
 '在': 3,
 '北': 4,
 '京': 5,
 '大': 6,
 '学': 7,
 '的': 8,
 '燕': 9,
 '园': 10,
 '看': 11,
 '了': 12,
 '中': 13,
 '国': 14,
 '男': 15,
 '篮': 16,
 '一': 17,
 '场': 18,
 '比': 19,
 '赛': 20,
 '红': 21,
 '参': 22,
 '加': 23,
 '校': 24,
 '组': 25,
 '织': 26,
 '歌': 27,
 '手': 28}

tag_to_ix

{'B-PER': 0,
 'I-PER': 1,
 'O': 2,
 'B-ORG': 3,
 'I-ORG': 4,
 'B-LOC': 5,
 'I-LOC': 6}

接下来定义两个辅助函数：

# 辅助函数
def prepare_sequence(seq, to_ix, max_len=None):
    if not max_len:
        max_len = len(seq)
    # 填充
    return torch.tensor([to_ix[w] for w in seq] + [0] * (max_len - len(seq)), dtype=torch.long)




def preprocess(data, word_to_ix, tag_to_ix):
    max_len = max(len(s) for s in data[0])
    seq_out, tag_out = [], []
    
    for seq, tags in data:      
        seq_out.append(prepare_sequence(seq, word_to_ix, max_len))
        tag_out.append(prepare_sequence(tags, tag_to_ix, max_len))
    
    seq_out = torch.stack(seq_out)
    tag_out = torch.stack(tag_out)
    # 构建mask
    mask = (seq_out != 0).to(torch.int8)

    return seq_out, tag_out, mask

可以测试一下：

prepare_sequence(list("小明在北京大学的燕园看了中国男篮的一场比赛"), word_to_ix)

tensor([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,  8, 17,
        18, 19, 20])

下面我们对整个玩具数据集应用preprocess函数，构造一个批大小为2的数据：

sequences, tags, mask = preprocess(training_data, word_to_ix, tag_to_ix)
sequences.shape

torch.Size([2, 21])

这里我们需要转换成batch_first=False的形式：

sequences = torch.transpose(sequences, 0, 1)
tags = torch.transpose(tags, 0, 1)
mask = torch.transpose(mask, 0, 1)

验证一下sequences：

sequences

tensor([[ 1,  1],
        [ 2, 21],
        [ 3, 22],
        [ 4, 23],
        [ 5, 12],
        [ 6,  7],
        [ 7, 24],
        [ 8, 25],
        [ 9, 26],
        [10,  8],
        [11, 24],
        [12, 10],
        [13, 27],
        [14, 28],
        [15,  6],
        [16, 20],
        [ 8,  0],
        [17,  0],
        [18,  0],
        [19,  0],
        [20,  0]])

变成了seq_len维度在前，上面的0表示填充符。

也可以看下对应的mask:

mask

tensor([[1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 1],
        [1, 0],
        [1, 0],
        [1, 0],
        [1, 0],
        [1, 0]], dtype=torch.int8)

训练

定义模型

# 训练
EMBEDDING_DIM = 5
HIDDEN_DIM = 4

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

BiLSTMCRF(
  (word_embeds): Embedding(29, 5)
  (lstm): LSTM(5, 2, bidirectional=True)
  (hidden2tag): Linear(in_features=4, out_features=7, bias=True)
  (crf): CRF(num_tags=7)
)

上面打印出了模型的架构。下面看模型在未训练之前的预测是怎样的：

with torch.no_grad():
    print(model.inference(sequences, mask))

[[3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2], [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]]

推理

这里的推理就是预测这个玩具数据集的结果，看它预测的是否准确：

with torch.no_grad():
    print(model.inference(sequences, mask))

[[0, 1, 2, 3, 4, 4, 4, 2, 5, 6, 2, 2, 3, 4, 4, 4, 2, 2, 2, 2, 2], 
 [0, 1, 2, 2, 2, 3, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2]]

可以对比下真实标签：

tags.transpose(0, 1).tolist()

[[0, 1, 2, 3, 4, 4, 4, 2, 5, 6, 2, 2, 3, 4, 4, 4, 2, 2, 2, 2, 2],
 [0, 1, 2, 2, 2, 3, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 0, 0, 0, 0, 0]]

可以发现，去掉填充符之后结果是一样的，说明确实学到了一些东西。

参考

1. ADVANCED: MAKING DYNAMIC DECISIONS AND THE BI-LSTM CRF
2. 动手实现CRF