自然语言处理-BERT-CRF深度解析-05

BERT-CRF：深度解析与完整实现

系列文章第 4 篇 | 难度：⭐⭐⭐⭐⭐

1. 为什么需要 BERT？

BiLSTM-CRF 的两个瓶颈：

1. 从零训练：词向量从随机初始化，需要大量标注数据
2. 上下文窗口有限：LSTM 处理长距离依赖仍有衰减

BERT（Bidirectional Encoder Representations from Transformers, Devlin et al., 2018）解决了这两个问题：

• 在海量无标注文本上预训练，学习丰富的语言表示
• Transformer 的 Self-Attention 直接建模任意距离的依赖
• 微调（Fine-tuning）只需少量标注数据即可达到 SOTA

2. BERT 架构深度解析

2.1 输入表示

BERT 的输入是三种嵌入的逐元素求和：

$$\text{Input}=\text{TokenEmb}(x)+\text{SegmentEmb}(s)+\text{PositionEmb}(t)$$

对于序列标注任务（单句），分段嵌入全为 0，位置嵌入为可学习参数。

特殊 Token：

[CLS] 张 三 在 北 京 工 作 [SEP]

• [CLS]：句子级表示（分类任务用，NER 一般忽略）
• [SEP]：句子结束标记

2.2 Transformer Encoder 层

BERT-base 由 12 层 Transformer Encoder 堆叠，每层包含：

Multi-Head Self-Attention：

$${\text{head}}_i=\text{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$$

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

$$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W^O$$

其中 $Q=K=V=H^{l-1}$（自注意力），$d_k=d_{\text{model}}/h$。

Feed-Forward Network（FFN）：

$$\text{FFN}(x)=\text{GELU}(x{W}_1+b_1)W_2+b_2$$

BERT 使用 GELU 而非 ReLU，更平滑。

残差连接 + Layer Norm：

$$H^l=\text{LayerNorm}(H^{l-1}+\text{MultiHead}(H^{l-1}))$$

$$H^l=\text{LayerNorm}(H^l+\text{FFN}(H^l))$$

2.3 预训练任务

MLM（Masked Language Modeling）：

随机遮盖 15% 的 token，预测被遮盖的词：

• 80% 替换为 [MASK]
• 10% 替换为随机词
• 10% 保持不变

NSP（Next Sentence Prediction）：

50% 真实相邻句对，50% 随机句对，二分类。

（注：后续研究 RoBERTa 发现 NSP 用处不大，中文 BERT 也通常不依赖 NSP。）

2.4 BERT 关键参数

3. BERT-CRF 微调架构

3.1 完整数据流

输入: ["张", "三", "在", "北", "京"]     ← 字符级（中文）
  ↓  WordPiece Tokenizer
[CLS] 张 三 在 北 京 [SEP]              ← 加特殊 token
  ↓  BERT 12层 Transformer
H ∈ R^{(T+2) × 768}                     ← 上下文表示
  ↓  取非特殊 token 位置（去掉[CLS][SEP]）
H' ∈ R^{T × 768}
  ↓  Dropout(0.1) + Linear(768 → |tags|)
Emissions ∈ R^{T × |tags|}              ← 发射分数
  ↓  CRF 层（与 BiLSTM-CRF 中相同）
Y* = Viterbi(Emissions, Transitions)    ← 最优标签序列

3.2 WordPiece 对齐问题（重点！）

中文 BERT 以字为单位，英文 BERT 以 WordPiece 子词为单位。

英文示例：

原始词:    ["playing",  "NLP",  "is",  "fun"]
子词:      ["playing",  "NL",  "##P", "is",  "fun"]
word_ids:  [0,          1,      1,     2,     3   ]   ← 位置映射
标签对齐:  [B-X,        B-Y,   -100,  O,     O   ]   ← -100 忽略

只取每个词的首个子词对应的 BERT 输出用于标签预测，其余子词的损失用 -100 屏蔽。

3.3 损失函数

与 BiLSTM-CRF 完全相同：

$$\mathcal{L}=-\log P(Y^{\ast }|X)=\log Z(X)-s(X,Y^{\ast })$$

$$=\underset{\text{前向算法}}{\underbrace{\text{logsumexp over all paths}}}-\underset{\text{直接计算}}{\underbrace{\text{真实路径得分}}}$$

梯度通过 BERT、Linear 和 CRF 转移矩阵共同反向传播。

4. 完整可运行代码（HuggingFace + torchcrf）

"""
BERT-CRF - 中文 NER 完整实现
依赖：transformers>=4.0, torch>=1.9, torchcrf, seqeval
安装：pip install transformers torch torchcrf seqeval
中文预训练模型：hfl/chinese-macbert-base（推荐）或 bert-base-chinese
"""

import os
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from transformers import (
    BertTokenizerFast,
    BertModel,
    get_linear_schedule_with_warmup,
)
from torchcrf import CRF
from seqeval.metrics import f1_score, classification_report
import numpy as np


# ─────────────────────────────────────────────
# 0. 配置
# ─────────────────────────────────────────────

class Config:
    # 模型
    bert_model   = 'bert-base-chinese'   # 或 'hfl/chinese-macbert-base'
    max_len      = 128
    dropout      = 0.1

    # 训练
    batch_size   = 16
    num_epochs   = 10
    bert_lr      = 2e-5      # BERT 小学习率
    head_lr      = 1e-3      # 分类头大学习率
    weight_decay = 0.01
    warmup_ratio = 0.1
    clip_grad    = 1.0

    # 标签
    label_pad_id = -100      # 忽略的标签 id（PAD/特殊token位置）

    device = 'cuda' if torch.cuda.is_available() else 'cpu'


# ─────────────────────────────────────────────
# 1. 标签映射
# ─────────────────────────────────────────────

class LabelVocab:
    def __init__(self, tag_list):
        # 确保 O 在最前，便于调试
        tags = ['O'] + [t for t in sorted(set(tag_list)) if t != 'O']
        self.tag2id = {t: i for i, t in enumerate(tags)}
        self.id2tag = {i: t for t, i in self.tag2id.items()}
        self.num_tags = len(tags)

    def encode(self, tags):
        return [self.tag2id[t] for t in tags]

    def decode(self, ids):
        return [self.id2tag.get(i, 'O') for i in ids]


# ─────────────────────────────────────────────
# 2. 数据集
# ─────────────────────────────────────────────

class NERDataset(Dataset):
    """
    支持：
      - 中文字符级（bert-base-chinese）
      - 英文 WordPiece 级（需要对齐）
    """

    def __init__(self, data, tokenizer, label_vocab, config):
        self.config      = config
        self.label_vocab = label_vocab
        self.samples     = []

        for words, tags in data:
            encoding = tokenizer(
                words,
                is_split_into_words=True,
                max_length=config.max_len,
                truncation=True,
                padding='max_length',
                return_tensors='pt',
            )
            # 标签对齐
            aligned_labels = self._align_labels(
                encoding.word_ids(), tags, label_vocab
            )

            self.samples.append({
                'input_ids':      encoding['input_ids'].squeeze(0),
                'attention_mask': encoding['attention_mask'].squeeze(0),
                'token_type_ids': encoding.get('token_type_ids',
                    torch.zeros_like(encoding['input_ids'])).squeeze(0),
                'labels':         torch.tensor(aligned_labels, dtype=torch.long),
                'orig_words':     words,
                'orig_tags':      tags,
            })

    def _align_labels(self, word_ids, tags, label_vocab):
        """
        将词级标签对齐到 token 级
        - [CLS]/[SEP]/[PAD] 位置：-100（忽略）
        - WordPiece 首个子词：对应词的标签
        - WordPiece 后续子词：-100（忽略）
        """
        aligned = []
        prev_word_id = None
        for word_id in word_ids:
            if word_id is None:
                aligned.append(Config.label_pad_id)
            elif word_id != prev_word_id:
                # 首个子词：取真实标签
                label = tags[word_id] if word_id < len(tags) else 'O'
                aligned.append(label_vocab.tag2id[label])
            else:
                # 后续子词：忽略
                aligned.append(Config.label_pad_id)
            prev_word_id = word_id

        # 补 padding
        while len(aligned) < Config.max_len:
            aligned.append(Config.label_pad_id)
        return aligned[:Config.max_len]

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx]


def collate_fn(batch):
    return {
        'input_ids':      torch.stack([b['input_ids']      for b in batch]),
        'attention_mask': torch.stack([b['attention_mask'] for b in batch]),
        'token_type_ids': torch.stack([b['token_type_ids'] for b in batch]),
        'labels':         torch.stack([b['labels']         for b in batch]),
        'orig_words':     [b['orig_words'] for b in batch],
        'orig_tags':      [b['orig_tags']  for b in batch],
    }


# ─────────────────────────────────────────────
# 3. BERT-CRF 模型
# ─────────────────────────────────────────────

class BertCRF(nn.Module):
    def __init__(self, bert_model_name, num_tags, dropout=0.1):
        super().__init__()
        self.bert    = BertModel.from_pretrained(bert_model_name)
        self.dropout = nn.Dropout(dropout)
        self.linear  = nn.Linear(self.bert.config.hidden_size, num_tags)
        self.crf     = CRF(num_tags, batch_first=True)

        # 初始化线性层
        nn.init.xavier_uniform_(self.linear.weight)
        nn.init.zeros_(self.linear.bias)

    def forward(self, input_ids, attention_mask, token_type_ids=None,
                labels=None):
        # BERT 编码
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        seq_output = outputs.last_hidden_state   # (B, L, 768)
        emissions  = self.linear(self.dropout(seq_output))   # (B, L, T)

        if labels is not None:
            # 训练：CRF 损失
            # 构建 mask（有效位置：attention_mask=1 且 label != -100）
            mask = (attention_mask.bool() &
                    (labels != Config.label_pad_id))
            # 将 -100 替换为 0（不影响结果，CRF 不看被 mask 的位置）
            labels_crf = labels.clone()
            labels_crf[labels_crf == Config.label_pad_id] = 0
            loss = -self.crf(emissions, labels_crf,
                             mask=mask, reduction='mean')
            return loss
        else:
            # 推断：Viterbi 解码
            mask = attention_mask.bool()
            return self.crf.decode(emissions, mask=mask)

    def freeze_bert_layers(self, num_layers=6):
        """冻结 BERT 前 N 层（低资源场景）"""
        for i, layer in enumerate(self.bert.encoder.layer):
            if i < num_layers:
                for p in layer.parameters():
                    p.requires_grad = False
        print(f"已冻结 BERT 前 {num_layers} 层")

    def unfreeze_all(self):
        for p in self.bert.parameters():
            p.requires_grad = True


# ─────────────────────────────────────────────
# 4. 训练器
# ─────────────────────────────────────────────

class BertCRFTrainer:
    def __init__(self, model, config):
        self.model  = model.to(config.device)
        self.config = config

        # 差异化学习率：BERT 用小 lr，分类头用大 lr
        bert_params  = list(model.bert.named_parameters())
        head_params  = list(model.linear.named_parameters()) + \
                       list(model.crf.named_parameters())

        # 权重衰减：不对 bias 和 LayerNorm 做衰减
        no_decay = ['bias', 'LayerNorm.weight', 'LayerNorm.bias']
        optimizer_grouped = [
            {'params': [p for n, p in bert_params if not any(nd in n for nd in no_decay)],
             'lr': config.bert_lr, 'weight_decay': config.weight_decay},
            {'params': [p for n, p in bert_params if     any(nd in n for nd in no_decay)],
             'lr': config.bert_lr, 'weight_decay': 0.0},
            {'params': [p for n, p in head_params if not any(nd in n for nd in no_decay)],
             'lr': config.head_lr, 'weight_decay': config.weight_decay},
            {'params': [p for n, p in head_params if     any(nd in n for nd in no_decay)],
             'lr': config.head_lr, 'weight_decay': 0.0},
        ]
        self.optimizer = AdamW(optimizer_grouped)
        self.best_f1   = 0.0

    def build_scheduler(self, total_steps):
        warmup = int(total_steps * self.config.warmup_ratio)
        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=warmup,
            num_training_steps=total_steps,
        )

    def train_epoch(self, loader):
        self.model.train()
        total_loss = 0.0
        for batch in loader:
            self.optimizer.zero_grad()
            loss = self.model(
                input_ids      = batch['input_ids'].to(self.config.device),
                attention_mask = batch['attention_mask'].to(self.config.device),
                token_type_ids = batch['token_type_ids'].to(self.config.device),
                labels         = batch['labels'].to(self.config.device),
            )
            loss.backward()
            nn.utils.clip_grad_norm_(self.model.parameters(),
                                     self.config.clip_grad)
            self.optimizer.step()
            self.scheduler.step()
            total_loss += loss.item()
        return total_loss / len(loader)

    @torch.no_grad()
    def evaluate(self, loader, label_vocab):
        self.model.eval()
        all_pred, all_true = [], []

        for batch in loader:
            pred_ids_list = self.model(
                input_ids      = batch['input_ids'].to(self.config.device),
                attention_mask = batch['attention_mask'].to(self.config.device),
                token_type_ids = batch['token_type_ids'].to(self.config.device),
            )
            labels = batch['labels']   # (B, L)

            for b in range(len(pred_ids_list)):
                # 找有效位置（label != -100 且 attention_mask=1）
                valid_mask = (labels[b] != Config.label_pad_id)
                n_valid    = valid_mask.sum().item()

                true_ids = labels[b][valid_mask].tolist()
                pred_ids = pred_ids_list[b][:n_valid]

                all_true.append(label_vocab.decode(true_ids))
                all_pred.append(label_vocab.decode(pred_ids))

        f1 = f1_score(all_true, all_pred)
        return f1, all_true, all_pred

    def fit(self, train_loader, val_loader, label_vocab,
            save_path='bert_crf_best.pt'):
        total_steps = len(train_loader) * self.config.num_epochs
        self.build_scheduler(total_steps)

        print(f"\n{'Epoch':>6} {'Loss':>8} {'Val F1':>8} {'LR(BERT)':>12}")
        print("-" * 45)

        for epoch in range(1, self.config.num_epochs + 1):
            loss = self.train_epoch(train_loader)
            f1, _, _ = self.evaluate(val_loader, label_vocab)
            cur_lr = self.optimizer.param_groups[0]['lr']

            flag = ' ← best' if f1 > self.best_f1 else ''
            print(f"{epoch:>6} {loss:>8.4f} {f1:>8.4f} {cur_lr:>12.2e}{flag}")

            if f1 > self.best_f1:
                self.best_f1 = f1
                torch.save(self.model.state_dict(), save_path)

        print(f"\n最优 Val F1: {self.best_f1:.4f}")


# ─────────────────────────────────────────────
# 5. 推理工具
# ─────────────────────────────────────────────

class NERPredictor:
    def __init__(self, model, tokenizer, label_vocab, config):
        self.model       = model.to(config.device)
        self.tokenizer   = tokenizer
        self.label_vocab = label_vocab
        self.config      = config
        self.model.eval()

    @torch.no_grad()
    def predict(self, words_list):
        """
        words_list: list of list of str
        返回: list of list of str（标签序列）
        """
        encoding = self.tokenizer(
            words_list,
            is_split_into_words=True,
            max_length=self.config.max_len,
            truncation=True,
            padding=True,
            return_tensors='pt',
        )
        pred_ids_list = self.model(
            input_ids      = encoding['input_ids'].to(self.config.device),
            attention_mask = encoding['attention_mask'].to(self.config.device),
        )

        results = []
        for b, words in enumerate(words_list):
            # word_ids 对齐到原始词
            word_ids = encoding.word_ids(batch_index=b)
            pred_tags = []
            prev_wid  = None
            tag_idx   = 0
            for wid in word_ids:
                if wid is None:
                    continue
                if wid != prev_wid and tag_idx < len(pred_ids_list[b]):
                    pred_tags.append(
                        self.label_vocab.id2tag[pred_ids_list[b][tag_idx]]
                    )
                    tag_idx += 1
                prev_wid = wid
            results.append(pred_tags[:len(words)])
        return results


# ─────────────────────────────────────────────
# 6. 主程序（使用示例数据；替换为真实数据即可生产使用）
# ─────────────────────────────────────────────

def make_demo_data():
    train_data = [
        (['张三', '在', '北京', '工作'],       ['B-PER', 'O', 'B-LOC', 'O']),
        (['李四', '来自', '上海'],              ['B-PER', 'O', 'B-LOC']),
        (['王五', '和', '赵六', '去', '深圳'],  ['B-PER', 'O', 'B-PER', 'O', 'B-LOC']),
        (['北京', '是', '中国', '首都'],        ['B-LOC', 'O', 'B-LOC', 'O']),
        (['张三', '认识', '李四'],              ['B-PER', 'O', 'B-PER']),
        (['华为', '总部', '在', '深圳'],        ['B-ORG', 'O', 'O', 'B-LOC']),
        (['腾讯', '是', '中国', '科技公司'],    ['B-ORG', 'O', 'B-LOC', 'O']),
        (['阿里巴巴', '总部', '在', '杭州'],    ['B-ORG', 'O', 'O', 'B-LOC']),
        (['小米', '成立', '于', '北京'],        ['B-ORG', 'O', 'O', 'B-LOC']),
        (['赵六', '在', '广州', '生活'],        ['B-PER', 'O', 'B-LOC', 'O']),
        (['百度', '总部', '在', '北京'],        ['B-ORG', 'O', 'O', 'B-LOC']),
        (['上海', '是', '金融', '中心'],        ['B-LOC', 'O', 'O', 'O']),
    ]
    test_data = [
        (['张三', '去', '上海', '出差'],  ['B-PER', 'O', 'B-LOC', 'O']),
        (['赵六', '在', '北京', '工作'],  ['B-PER', 'O', 'B-LOC', 'O']),
        (['字节', '总部', '在', '北京'],  ['B-ORG', 'O', 'O', 'B-LOC']),
    ]
    return train_data, test_data


def main():
    print("=" * 60)
    print("BERT-CRF - 中文 NER 完整演示")
    print("=" * 60)

    config = Config()
    print(f"使用设备: {config.device}")
    print(f"BERT 模型: {config.bert_model}")

    train_data, test_data = make_demo_data()

    # 标签映射
    all_tags = [t for _, tags in train_data + test_data for t in tags]
    label_vocab = LabelVocab(all_tags)
    print(f"标签集合({label_vocab.num_tags}): {list(label_vocab.tag2id.keys())}")

    # Tokenizer
    print(f"\n加载 tokenizer: {config.bert_model}")
    tokenizer = BertTokenizerFast.from_pretrained(config.bert_model)

    # 数据集
    train_dataset = NERDataset(
        [d for d in train_data], tokenizer, label_vocab, config
    )
    test_dataset  = NERDataset(
        [d for d in test_data],  tokenizer, label_vocab, config
    )
    train_loader = DataLoader(train_dataset, batch_size=config.batch_size,
                              shuffle=True,  collate_fn=collate_fn)
    test_loader  = DataLoader(test_dataset,  batch_size=config.batch_size,
                              shuffle=False, collate_fn=collate_fn)

    # 模型
    print(f"\n加载 BERT: {config.bert_model}")
    model = BertCRF(config.bert_model, label_vocab.num_tags, config.dropout)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"总参数量: {total_params:,}")

    # 训练
    trainer = BertCRFTrainer(model, config)
    trainer.fit(train_loader, test_loader, label_vocab)

    # 最终评估
    model.load_state_dict(
        torch.load('bert_crf_best.pt', map_location=config.device)
    )
    f1, y_true, y_pred = trainer.evaluate(test_loader, label_vocab)
    print(f"\n最终测试集 F1: {f1:.4f}")
    print(classification_report(y_true, y_pred))

    # 单句推理
    print("\n[推理示例]")
    predictor = NERPredictor(model, tokenizer, label_vocab, config)
    test_sents = [
        ['字节跳动', '总部', '在', '北京'],
        ['马云', '创立', '了', '阿里巴巴'],
    ]
    preds = predictor.predict(test_sents)
    for sent, pred in zip(test_sents, preds):
        print(f"  {sent}")
        print(f"  {pred}")
        print()


if __name__ == '__main__':
    main()

5. 工程部署优化

5.1 ONNX 导出（加速推理）

import torch.onnx

dummy_input = (
    torch.zeros(1, 128, dtype=torch.long),   # input_ids
    torch.ones(1, 128, dtype=torch.long),    # attention_mask
)

# 注意：CRF 层不支持直接 ONNX 导出，需单独处理
# 方案：导出 BERT+Linear 部分，CRF Viterbi 用 ONNX Runtime 外的 CPU 实现
torch.onnx.export(
    model.bert,
    dummy_input,
    'bert_encoder.onnx',
    opset_version=13,
    input_names=['input_ids', 'attention_mask'],
    output_names=['last_hidden_state'],
    dynamic_axes={
        'input_ids':      {0: 'batch', 1: 'seq'},
        'attention_mask': {0: 'batch', 1: 'seq'},
    }
)

5.2 模型量化（INT8）

from transformers import pipeline

# 使用 bitsandbytes 量化（推理显存减半）
model_int8 = BertCRF.from_pretrained(
    'bert-base-chinese',
    load_in_8bit=True,   # 需要 bitsandbytes
    device_map='auto',
)

5.3 知识蒸馏到 BiLSTM-CRF

# BERT-CRF 作为 Teacher，BiLSTM-CRF 作为 Student
# 软标签蒸馏损失（KL 散度）
def distillation_loss(student_logits, teacher_logits, temperature=4.0):
    teacher_probs = torch.softmax(teacher_logits / temperature, dim=-1)
    student_log   = torch.log_softmax(student_logits / temperature, dim=-1)
    return nn.KLDivLoss(reduction='batchmean')(student_log, teacher_probs) * (temperature ** 2)

# 总损失 = α * 硬标签CRF损失 + (1-α) * 软标签蒸馏损失
alpha = 0.5
total_loss = alpha * crf_loss + (1 - alpha) * distillation_loss(
    student_emissions, teacher_emissions.detach()
)

6. 中文 NER 推荐预训练模型

推荐：中文通用 NER 首选 hfl/chinese-roberta-wwm-ext，垂直领域先找对应的领域预训练模型。

7. 与 BiLSTM-CRF 的核心差异

8. 文档索引