模型蒸馏与轻量化

from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 蒸馏好的中文小模型
MODEL_NAME = "hfl/chinese-roberta-wwm-ext-small"  # 6层,参数量减半

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME, num_labels=3
)

# 后续训练流程和第 12 篇完全一样
# 只需要换一个 model name,其他代码不变

import torch
import torch.nn.functional as F
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
)

TEMPERATURE = 4.0
ALPHA = 0.5

teacher_model = AutoModelForSequenceClassification.from_pretrained(
    "./chat-moderation-bert"  # 微调好的 BERT-base
)
teacher_model.eval()

student_model = AutoModelForSequenceClassification.from_pretrained(
    "hfl/chinese-electra-180g-small", num_labels=3
)

class DistillationTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")
        student_outputs = model(**inputs)
        student_logits = student_outputs.logits

        with torch.no_grad():
            teacher_outputs = teacher_model(**inputs)
            teacher_logits = teacher_outputs.logits

        hard_loss = F.cross_entropy(student_logits, labels)

        soft_student = F.log_softmax(student_logits / TEMPERATURE, dim=-1)
        soft_teacher = F.softmax(teacher_logits / TEMPERATURE, dim=-1)
        distill_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean")
        distill_loss = distill_loss * (TEMPERATURE ** 2)

        loss = ALPHA * distill_loss + (1 - ALPHA) * hard_loss

        return (loss, student_outputs) if return_outputs else loss

trainer = DistillationTrainer(
    model=student_model,
    args=TrainingArguments(
        output_dir="./distilled-model",
        learning_rate=5e-5,       # Student 用稍大的学习率
        num_train_epochs=5,       # 蒸馏训练多几轮
        per_device_train_batch_size=32,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
    ),
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["test"],
    compute_metrics=compute_metrics,
)

trainer.train()

import torch

model = AutoModelForSequenceClassification.from_pretrained("./chat-moderation-bert")
model.eval()

quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},  # 量化所有 Linear 层
    dtype=torch.qint8
)

# 对比模型大小
import os

torch.save(model.state_dict(), "original.pt")
torch.save(quantized_model.state_dict(), "quantized.pt")

original_size = os.path.getsize("original.pt") / 1024 / 1024
quantized_size = os.path.getsize("quantized.pt") / 1024 / 1024
print(f"原始: {original_size:.1f}MB → 量化: {quantized_size:.1f}MB")
# → 原始: 418.2MB → 量化: 108.3MB

import time

def benchmark(model, tokenizer, test_texts, num_runs=3):
    """测量推理延迟和吞吐量"""
    model.eval()
    encoded = tokenizer(
        test_texts, padding=True, truncation=True,
        max_length=64, return_tensors="pt"
    )

    latencies = []
    for _ in range(num_runs):
        start = time.perf_counter()
        with torch.no_grad():
            model(**encoded)
        latencies.append(time.perf_counter() - start)

    avg_ms = (sum(latencies) / num_runs) * 1000
    per_sample = avg_ms / len(test_texts)
    return {
        "batch_ms": f"{avg_ms:.1f}",
        "per_sample_ms": f"{per_sample:.2f}",
        "throughput": f"{1000 / per_sample:.0f} samples/s",
    }