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]]
可以发现，去掉填充符之后结果是一样的，说明确实学到了一些东西。
参考
- ADVANCED: MAKING DYNAMIC DECISIONS AND THE BI-LSTM CRF
- 动手实现CRF