知识蒸馏（KD）详解三：基于BERT的知识蒸馏代码实战

前言

数据集：数据集我们使用huggingface中的"banking77"
任务：拿BERT构建和训练一个意图识别模型
关于BERT模型和知识蒸馏理论知识在前两节已经详细介绍了，可以自己去查看

---

1. 教师模型训练

1.1 训练

教师模型采用BERT骨干网络 + 自定义分类头来训练，完整代码如下：

import os  
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  
import evaluate  
from datasets import load_dataset, DatasetDict  # huggingface  
from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer, set_seed, DataCollatorWithPadding, AutoConfig  # huggingface  
import torch  
import torch.nn as nn  
import torch.nn.functional as F  
from torchinfo import summary  
  
  
# ======================== 配置 ==================================
DATASET = "banking77" # 可选: "banking77" | "clinc150" | "massive_en" | "massive_zh"  
BACKBONE = "bert-base-uncased"  
MAX_LEN = 256  
SEED = 42  
LR = 2e-5  
EPOCHS = 10 
BSZ_TRAIN = 32  
BSZ_EVAL = 32  
  
OUT_DIR= f"runs/{DATASET}_teacher_custom"  
set_seed(SEED)  
  
# 1 ==================== 数据集加载 ====================
def load_intent_dataset(dataset_name):  
  
    if dataset_name == "banking77":  
  
        ds = load_dataset("banking77")  
        text_col, label_col = "text", "label"  
  
        num_labels = len(ds["train"].unique(label_col))  
    else:  
        raise ValueError("Unknown dataset")  
  
    if "validation" not in ds:  
        if "test" in ds:  
            ds = DatasetDict({"train": ds["train"], "validation": ds["test"]})  
        else:  
            split = ds.train_test_split(test_size=0.1, SEED=SEED)  
            ds = DatasetDict({"train": split["train"], "validation": split["test"]})  
  
    return ds, text_col, label_col, num_labels  
  
# 2.======================= 数据预处理 ====================
def tok_fn(batch, tok, max_len, TEXT):  
    return tok(batch[TEXT], truncation=True, max_length=max_len)  
  
  
# 4. ======================= 构建Teacher ====================  
class MeanPooler(torch.nn.Module):  
    """  
    hidden: 通常是 last_hidden_state，形状 [B, S, H]（批大小/序列长/隐藏维）  
    mask: 通常是 attention_mask，形状 [B, S]，有效 token 为 1，padding 为 0。  
    """    
    def forward(self, hidden, mask):  
        mask = mask.unsqueeze(-1) # [b,s,1]  
        summded = (hidden * mask).sum(dim=1)  
        denom = mask.sum(dim=1).clamp(min=1e-6)  
        return summded / denom  
  
class TeacherBackboneHead(torch.nn.Module):  
    def __init__(self, backbone_name, num_labels, config=None):  
        super().__init__()  
        self.backbone = AutoModel.from_pretrained(backbone_name)  
        h = self.backbone.config.hidden_size  
        self.pool = MeanPooler()  
        self.head = nn.Sequential(  
            nn.Dropout(0.2),  
            nn.Linear(h, h),  
            nn.GELU(),  
            nn.Dropout(0.2),  
            nn.Linear(h, num_labels)  
        )  
  
        self.config = config  
  
    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None):  
  
        assert input_ids is not None and attention_mask is not None and token_type_ids is not None, "input_ids, attention_mask, token_ids must be not None"  
  
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, return_dict=True)  
        hidden = out.last_hidden_state  
        sent = self.pool(hidden, attention_mask) # (B, H)  
        logits = self.head(sent)  #(B,C)  
        loss = F.cross_entropy(logits, labels) if labels is not None else None  
  
        return {"loss":loss, "logits": logits}  
  
  
if __name__ == "__main__":  
  
    cfg = AutoConfig.from_pretrained(BACKBONE)  
  
    ds, TEXT, LABEL, N_LABELS = load_intent_dataset(DATASET)  
    tok = AutoTokenizer.from_pretrained(BACKBONE, use_fast=True)  #  
  
    # ds: {'train': ['text', 'label'], 'validation': ['text', 'label']}    
    # ds_enc: {'train': ['label', 'input_ids', 'token_type_ids', 'attention_mask'], 'validation': ['label', 'input_ids', 'token_type_ids', 'attention_mask']}    
    ds_enc = ds.map(tok_fn, batched=True, remove_columns=[TEXT], fn_kwargs={"tok": tok, "max_len": MAX_LEN, "TEXT": TEXT})  
  
    """  
    ds_enc["train"][0] =>    
    {      
    "input_ids": [101, 2026, 3899, ... , 102],      
    "attention_mask": [1, 1, 1, ..., 1],      
    "token_type_ids": [0, 0, 0, ..., 0],      
    "labels": 17    
    }    
    """    
    labels = ds["train"].features["label"].names  
    id2label = {i: name for i, name in enumerate(labels)}  
    cfg.id2label = id2label  
  
    # 每一个mini batch padding  不在tok_fn中做padding，因为Trainer会自动padding  
    collator = DataCollatorWithPadding(tok)  
  
  
    # 3. ======================= 定义指标 ====================    
    acc = evaluate.load("accuracy")  # 预测正确的率  
    # 所有类别求宏平均F1，反应整体类别预测的准确性  
    f1 = evaluate.load("f1")  # F1 = 2 * (precision * recall) / (precision + recall), precision=TP / (TP + FP), recall=TP / (TP + FN)  
    def compute_metrics(eval_pred):  
        logits, labels = eval_pred  
        preds = logits.argmax(-1)  
        acc_result = acc.compute(predictions=preds, references=labels)  
        f1_result = f1.compute(predictions=preds, references=labels, average="macro")  
        return {"accuracy": acc_result["accuracy"], "f1": f1_result["f1"]}  
  
  
    # 5. ======================= 训练Teacher ====================  
    teacher = TeacherBackboneHead(BACKBONE, N_LABELS, config=cfg)  
  
    # 输出模型结构  
    examples = [ds_enc["train"][i] for i in range(2)]  
    batch_ex = collator(examples)  
    inputs_ex = {key: batch_ex[key] for key in ["input_ids", "attention_mask", "token_type_ids"]}  
    class Wrapper(nn.Module):  
        def __init__(self, core):  
            super().__init__()  
            self.core = core  
        def forward(self, input_ids, attention_mask=None, token_type_ids=None):  
            out = self.core(input_ids=input_ids,  
                            attention_mask=attention_mask,  
                            token_type_ids=token_type_ids)  
            return out["logits"]  # [B, num_labels]  
  
    wrapped = Wrapper(teacher)  
    summary(  
        wrapped,  
        input_data=inputs_ex,  
        depth=3,  
        col_names=("input_size","output_size","num_params","trainable"),  
        verbose=2  
    )  
  
    # args  
    args = TrainingArguments(  
        output_dir=OUT_DIR,  
        learning_rate=LR,  
        per_device_train_batch_size=BSZ_TRAIN,  
        per_device_eval_batch_size=BSZ_EVAL,  
        num_train_epochs=EPOCHS,  
        evaluation_strategy="epoch",  
        save_strategy="epoch",  
        logging_steps=50,  
        load_best_model_at_end=True,  
        metric_for_best_model="f1",  
        greater_is_better=True,  
        seed=SEED,  
        fp16=torch.cuda.is_available(),  
        report_to="none"  
    )  
  
    trainer = Trainer(  
        model=teacher,  
        args=args,  
        train_dataset=ds_enc["train"],  
        eval_dataset=ds_enc["validation"],  
        tokenizer=tok,  
        data_collator=collator, # 动态padding  
        compute_metrics=compute_metrics  
    )  
  
    trainer.train()  
    metrics = trainer.evaluate()  
    print(f"Final metrics: {metrics}")  
  
    # 保存（包含自定义头）  
    os.makedirs(f"{OUT_DIR}/best", exist_ok=True)  
    trainer.save_model(f"{OUT_DIR}/best")  
    tok.save_pretrained(f"{OUT_DIR}/best")  # trainer 中会自动保存tokenizer, 这一步可以不要  
    cfg =cfg.save_pretrained(f"{OUT_DIR}/best")  
    print(f"Saved teacher to {OUT_DIR}/best")

---

运行之后，首先我们可以看到模型结构：以下部分截图

我们可以看到我们这个教师模型的参数大概1亿1千万

---

模型训练完成后，如下：

模型最后的精度在0.92，f1也在0.92

---

1.2 测试

我们找几条文本来测试一下效果：
测试代码：

from pyarrow.lib import Mapping  
from transformers import AutoTokenizer, AutoConfig  
import os, torch, numpy as np  
from bert_kd import TeacherBackboneHead, MAX_LEN  
from safetensors.torch import load_file  
  
@torch.no_grad()  
def predict(texts, topk3, model, tok, cfg):  
    batch = tok(texts, return_tensors="pt", padding=True, truncation=True, max_length=MAX_LEN)  
    inputs = {k: v.to(device) for k, v in batch.items()}  
    out = model(**inputs)  
    probs = torch.softmax(out["logits"], dim=-1)  
    topv,topi = torch.topk(probs, topk3, dim=-1)  
    id2label = cfg.id2label  
    for i, text in enumerate(texts):  
  
        top_l = [id2label[id] for id in topi[i].tolist()]  
        top_v = topv[i].tolist()  
  
        print("\nText:", text)  
        print("Top-3:", (top_l, top_v))  
  
if __name__ == "__main__":  
  
  
    TEACHER_DIR = r"runs/banking77_teacher_custom/best"  
    cfg_t = AutoConfig.from_pretrained(TEACHER_DIR)  
    tok_t = AutoTokenizer.from_pretrained(TEACHER_DIR)  
    model_t = TeacherBackboneHead(TEACHER_DIR, num_labels=77, config=cfg_t)  
  
    device = torch.device("cuda:0")  
    state_dict_t = load_file(os.path.join(TEACHER_DIR, "model.safetensors"))  
    model_t.load_state_dict(state_dict_t)  
    model_t = model_t.to(device)  
    model_t.eval()  
  
    samples = [  
        "I lost my card yesterday. Can you block it?",                      # 丢卡  
        "Why was I charged twice for the same transaction?",               # 重复扣款  
        "My contactless payment isn't working at the store.",              # 闪付失效  
        "The transfer hasn't arrived to the recipient yet.",               # 转账未到账  
        "How can I change my PIN code?",                                   # 修改PIN  
        "Is Apple Pay supported for my card?",                             # Apple Pay 支持  
        "The ATM swallowed my card. What should I do?",                    # ATM 吞卡  
        "Balance didn’t update after my bank transfer.",                   # 余额未更新  
        "I want to terminate my account.",                                 # 注销账户  
        "What are the fees for international card payments?",              # 汇率/手续费  
        "Can I get a spare physical card for my account?",                 # 备用卡  
        "Why do I need to verify my identity again?",                      # 身份验证  
    ]  
  
    predict(samples, 3, model=model_t, tok=tok_t, cfg=cfg_t)

---

测试结果：

从测试结果来看，测试的意图识别结果都是正确的。

---

2. 学生模型训练

2.1 代码

以下代码不仅训练了一个带有KD的学生模型，还训练了一个不带KD的学生模型，对比蒸馏与非蒸馏的效果。
教师模型采用1中训练的模型。

import os  
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  
import evaluate  
from datasets import load_dataset, DatasetDict  # huggingface  
from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer, set_seed, DataCollatorWithPadding, AutoConfig, AutoModelForSequenceClassification  # huggingface  
import torch  
import torch.nn as nn  
import torch.nn.functional as F  
from torchinfo import summary  
from bert_kd import TeacherBackboneHead  
from safetensors.torch import load_file  
from dataclasses import dataclass  
from typing import Optional, Union, Callable, Dict, Any, List, Tuple  
  
  
# ======================== 配置 ==================================
DATASET = "banking77"  
TEACHER_DIR = f"runs/{DATASET}_teacher_custom" # 训练好的teacher模型路径  
MAX_LEN = 256  
SEED = 42  
LR = 3e-4  
EPOCHS = 20  
BSZ_TRAIN = 32  
BSZ_EVAL = 32  
  
STUDENT_CKPT = "huawei-noah/TinyBERT_General_4L_312D"  
  
OUT_DIR= f"runs/{DATASET}_tinybert_kd"  
set_seed(SEED)  
  
def load_intent_dataset(dataset_name):  
  
    if dataset_name == "banking77":  
  
        ds = load_dataset("banking77")  
        text_col, label_col = "text", "label"  
  
        num_labels = len(ds["train"].unique(label_col))  
    else:  
        raise ValueError("Unknown dataset")  
  
    if "validation" not in ds:  
        if "test" in ds:  
            ds = DatasetDict({"train": ds["train"], "validation": ds["test"]})  
        else:  
            split = ds.train_test_split(test_size=0.1, SEED=SEED)  
            ds = DatasetDict({"train": split["train"], "validation": split["test"]})  
  
    return ds, text_col, label_col, num_labels  
  
def tok_fn(batch, tok, max_len, TEXT):  
    return tok(batch[TEXT], truncation=True, max_length=max_len)  
  
def dual_tok_fn(batch, tok_s, tok_t,  max_len, TEXT):  
    enc_s = tok_s(batch[TEXT], truncation=True, max_length=max_len)  
    enc_t = tok_t(batch[TEXT], truncation=True, max_length=max_len)  
  
    out = {  
        "input_ids": enc_s["input_ids"],  
        "attention_mask": enc_s["attention_mask"],  
        "token_type_ids": enc_s["token_type_ids"],  
        "t_input_ids": enc_t["input_ids"],  
        "t_attention_mask": enc_t["attention_mask"],  
        "t_token_type_ids": enc_t["token_type_ids"],  
    }  
  
    return out  
  
  
@dataclass  
class KDConfig:  
    T: float = 4.0  
    alpha: float = 0.7  
    mid_weight: float = 0.0   # 如需中间层蒸馏，设 >0 并看下方可选块  
  
# 构建知识蒸馏训练trainer  
class DistillerTrainer(Trainer):  
  
    def __init__(self, *ar, teacher=None, kd=None, **kw):  
        super().__init__(*ar, **kw)  
  
        self.teacher = teacher  
        self.kd = kd  
  
        self.T = kd.T  
        self.alpha = kd.alpha  
  
    def compute_loss(self, model, inputs, return_outputs=False):  
  
        labels = inputs["labels"]  
  
        # students out  
        out_s = model(  
            input_ids=inputs["input_ids"],  
            attention_mask=inputs["attention_mask"],  
            token_type_ids=inputs["token_type_ids"],  
            return_dict=True  
        )  
        z_s = out_s.logits  
  
        # teacher out  
        with torch.no_grad():  
            out_t = self.teacher.backbone(  
                input_ids=inputs["t_input_ids"],  
                attention_mask=inputs["t_attention_mask"],  
                token_type_ids=inputs["t_token_type_ids"],  
                return_dict=True  
            )  
            hidden = out_t.last_hidden_state  
            sent = self.teacher.pool(hidden, inputs["t_attention_mask"]) # (B, H)  
            z_t = self.teacher.head(sent)  #(B,C)  
  
        loss_kd = F.kl_div(F.log_softmax(z_s/self.T, dim=-1), F.softmax(z_t/self.T, dim=-1), reduction="batchmean") * (self.T ** 2)  
        loss_ce = F.cross_entropy(z_s, labels)  
        loss = self.alpha * loss_kd + (1 - self.alpha) * loss_ce  
  
        return (loss, out_s) if return_outputs else loss  
  
  
class DualDataClllatorWithPadding:  
  
    def __init__(self, tok_s, tok_t):  
        super().__init__()  
  
        self.collator_s = DataCollatorWithPadding(tok_s)  
        self.collator_t = DataCollatorWithPadding(tok_t)  
  
        self.keys_s = ["input_ids","attention_mask","token_type_ids", "label"]  
        self.keys_t = ["t_input_ids","t_attention_mask","t_token_type_ids"]  
  
    def __call__(self, features: List[Dict[str, Any]])-> Dict[str, Any]:  
        # features: list of dict  
        features_s = [{k:v for k,v in f.items() if k in self.keys_s} for f in features]  
        features_t = [{k[2:]:v for k,v in f.items() if k in self.keys_t} for f in features]  # 去掉t_前缀  
        batch_s = self.collator_s(features_s)  
        batch_t = self.collator_t(features_t)  
  
        # batch_s 中加入带t_的batch_t  
        for k,v in batch_t.items():  
            batch_s[f"t_{k}"] = v  
  
        return batch_s  
  
  
if __name__ == "__main__":  
  
    # load teacher model  
    TEACHER_DIR = r"runs/banking77_teacher_custom/best"  
    cfg_t = AutoConfig.from_pretrained(TEACHER_DIR)  
    tok_t = AutoTokenizer.from_pretrained(TEACHER_DIR, use_fast=True)  
    model_t = TeacherBackboneHead(TEACHER_DIR, num_labels=77, config=cfg_t)  
    device = torch.device("cuda:0")  
    state_dict_t = load_file(os.path.join(TEACHER_DIR, "model.safetensors"))  
    model_t.load_state_dict(state_dict_t)  
    model_t = model_t.to(device)  
    model_t.eval()  
  
    # load dataset  
    ds, TEXT, LABEL, N_LABELS = load_intent_dataset(DATASET)  
    tok_s = AutoTokenizer.from_pretrained(STUDENT_CKPT, use_fast=True)  
    ds_enc_s = ds.map(tok_fn, batched=True, remove_columns=[TEXT], fn_kwargs={"tok": tok_s, "max_len": MAX_LEN, "TEXT": TEXT})  
    collator_s = DataCollatorWithPadding(tok_s)  
    collator_t = DataCollatorWithPadding(tok_t)  
    ds_enc = ds.map(dual_tok_fn, batched=True, remove_columns=[TEXT], fn_kwargs={"tok_s": tok_s, "tok_t": tok_t, "max_len": MAX_LEN, "TEXT": TEXT})  
    collator = DualDataClllatorWithPadding(tok_s, tok_t)  
  
  
    # 学生模型  
    student = AutoModelForSequenceClassification.from_pretrained(STUDENT_CKPT, num_labels=N_LABELS)  
  
    # cfg  
    cfg = AutoConfig.from_pretrained(STUDENT_CKPT, num_labels=N_LABELS)  
    labels = ds["train"].features["label"].names  
    id2label = {i: name for i, name in enumerate(labels)}  
    cfg.id2label = id2label  
  
    # acc f1  
    acc = evaluate.load("accuracy")  # 预测正确的率  
    f1 = evaluate.load("f1")  
    def compute_metrics(eval_pred):  
        logits, labels = eval_pred  
        preds = logits.argmax(-1)  
        acc_result = acc.compute(predictions=preds, references=labels)  
        f1_result = f1.compute(predictions=preds, references=labels, average="macro")  
        return {"accuracy": acc_result["accuracy"], "f1": f1_result["f1"]}  
  
  
    # 输出模型结构  
    examples = [ds_enc_s["train"][i] for i in range(2)]  
    batch_ex = collator_s(examples)  
    inputs_ex = {key: batch_ex[key] for key in ["input_ids", "attention_mask", "token_type_ids"]}  
    class Wrapper(nn.Module):  
        def __init__(self, core):  
            super().__init__()  
            self.core = core  
        def forward(self, input_ids, attention_mask=None, token_type_ids=None):  
            out = self.core(input_ids=input_ids,  
                            attention_mask=attention_mask,  
                            token_type_ids=token_type_ids)  
            return out["logits"]  # [B, num_labels]  
    wrapped = Wrapper(student)  
    summary(  
        wrapped,  
        input_data=inputs_ex,  
        depth=3,  
        col_names=("input_size","output_size","num_params","trainable"),  
        verbose=2  
    )  
  
    # kd  
    KD = KDConfig()  
  
    # args  
    args_ = TrainingArguments(  
        output_dir=OUT_DIR,  
        learning_rate=LR,  
        per_device_train_batch_size=BSZ_TRAIN,  
        per_device_eval_batch_size=BSZ_EVAL,  
        num_train_epochs=EPOCHS,  
        evaluation_strategy="epoch",  
        save_strategy="epoch",  
        logging_steps=500,  
        load_best_model_at_end=True,  
        metric_for_best_model="f1",  
        greater_is_better=True,  
        seed=SEED,  
        fp16=torch.cuda.is_available(),  
        report_to="none",  
        remove_unused_columns=False  
    )  
  
    # trainer  
    trainer = DistillerTrainer(  
        model=student,  
        args=args_,  
        train_dataset=ds_enc["train"],  
        eval_dataset=ds_enc["validation"],  
        tokenizer=tok_s,  
        data_collator=collator, # 动态padding  
        compute_metrics=compute_metrics,  
        teacher=model_t,  
        kd=KD  
    )  
  
    # train  
    trainer.train()  
    metrics = trainer.evaluate()  
    print(f"Final metrics for student: {metrics}")  
  
    # 保存（包含自定义头）  
    os.makedirs(f"{OUT_DIR}/best", exist_ok=True)  
    trainer.save_model(f"{OUT_DIR}/best")  
    tok_s.save_pretrained(f"{OUT_DIR}/best")  # trainer 中会自动保存tokenizer, 这一步可以不要  
    cfg.save_pretrained(f"{OUT_DIR}/best")  
    print(f"Saved student to {OUT_DIR}/best")  
  
  
    # 对比直接训练  
    student2 = AutoModelForSequenceClassification.from_pretrained(STUDENT_CKPT, num_labels=N_LABELS)  
    trainer2 = Trainer(  
        model=student2,  
        args=args_,  
        train_dataset=ds_enc_s["train"],  
        eval_dataset=ds_enc_s["validation"],  
        tokenizer=tok_s,  
        data_collator=collator_s, # 动态padding  
        compute_metrics=compute_metrics  
    )  
  
  
    # train  
    trainer2.train()  
    metrics2 = trainer2.evaluate()  
    print(f"Final metrics for student no kd: {metrics2}")  
  
    # 保存（包含自定义头）  
    os.makedirs(f"{OUT_DIR}_no_kd/best", exist_ok=True)  
    trainer2.save_model(f"{OUT_DIR}_no_kd/best")  
    tok_s.save_pretrained(f"{OUT_DIR}_no_kd/best")  # trainer 中会自动保存tokenizer, 这一步可以不要  
    cfg.save_pretrained(f"{OUT_DIR}_no_kd/best")  
    print(f"Saved student to {OUT_DIR}_no_kd/best")

---

2.2 结果

直接运行上面代码，首先可以看到student模型结构和参数。可发现其参数量比教师模型少了很多，大概只有教师模型的10%。

最后acc和f1的结果我们可以看到：
KD Student:

---

No KD Student:

---

结果：带KD的模型比不带KD的模型提升了4%左右。

---

结论

1. 基于BERT 骨干网络的预训练模型，结合自定义的分类头，训练一个意图识别模型，需要很少的epoch就能完成，并且得到的效果还不错。说明BERT的预训练模型确实有着不错的语言特征提取能力。
2. KD的效果确实能够在一个小模型上得到跟教师模型差不多的效果，并且由于不带KD直接训练的小模型。说明KD确实起了作用。
3. 这里我们只用了基本的KD手断，比如中间层计算损失也没用，参数也没怎么去调整。
   4. KD的参数很大程度了影响了结果，这个参数也是试了好几次得到相对比较好的效果，最好是加上隐藏层损失，那样会更好一点，这个留给看这篇文章的人自己去尝试。