知识蒸馏（Knowledge Distillation）完全指南：模型轻量化与压缩技术

Teacher Model（大而精确）
      │
      │ 传递软目标（概率分布）
      ▼
Student Model（小而快速）

softmax_T(z_i) = exp(z_i / T) / sum_j(exp(z_j / T))

import torch
import torch.nn.functional as F

def temperature_softmax(logits, temperature=1.0):
    """经温度调整后的 softmax。"""
    return F.softmax(logits / temperature, dim=-1)

# 示例
logits = torch.tensor([2.0, 1.0, 0.1, 0.5])
print("T=1:", temperature_softmax(logits, T=1).numpy().round(3))
# [0.596, 0.219, 0.090, 0.096]
print("T=4:", temperature_softmax(logits, T=4).numpy().round(3))
# [0.345, 0.262, 0.195, 0.216] — 更均匀

L_KD = T^2 * KLDiv(softmax_T(student_logits), softmax_T(teacher_logits))

L_CE = CrossEntropy(student_logits, true_labels)

L = alpha * L_CE + (1 - alpha) * L_KD

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torchvision.datasets as datasets


class KnowledgeDistillationLoss(nn.Module):
    """
    Hinton et al. (2015) 的知识蒸馏损失。
    L = alpha * CE(student, labels) + (1-alpha) * T^2 * KLDiv(student_soft, teacher_soft)
    """
    def __init__(self, temperature=4.0, alpha=0.5):
        super().__init__()
        self.T = temperature
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss()
        self.kl_loss = nn.KLDivLoss(reduction='batchmean')

    def forward(self, student_logits, teacher_logits, labels):
        # 硬目标损失
        loss_ce = self.ce_loss(student_logits, labels)

        # 软目标损失（KL 散度）
        student_soft = F.log_softmax(student_logits / self.T, dim=-1)
        teacher_soft = F.softmax(teacher_logits / self.T, dim=-1)
        loss_kd = self.kl_loss(student_soft, teacher_soft) * (self.T ** 2)

        return self.alpha * loss_ce + (1 - self.alpha) * loss_kd


def train_with_distillation(
    teacher, student, train_loader,
    num_epochs=10, temperature=4.0, alpha=0.5,
    device='cuda'
):
    """Teacher-Student 蒸馏训练循环。"""
    teacher = teacher.to(device).eval()  # Teacher 冻结
    student = student.to(device)

    # 冻结 Teacher 参数
    for param in teacher.parameters():
        param.requires_grad = False

    criterion = KnowledgeDistillationLoss(temperature=temperature, alpha=alpha)
    optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

    for epoch in range(num_epochs):
        student.train()
        total_loss = 0.0
        correct = 0
        total = 0

        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)

            # Teacher 推理（无需梯度）
            with torch.no_grad():
                teacher_logits = teacher(images)

            # Student 推理
            student_logits = student(images)

            # KD 损失计算
            loss = criterion(student_logits, teacher_logits, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * images.size(0)
            correct += (student_logits.argmax(1) == labels).sum().item()
            total += images.size(0)

        scheduler.step()
        print(f"Epoch {epoch+1}/{num_epochs} | "
              f"Loss: {total_loss/total:.4f} | "
              f"Acc: {correct/total:.4f}")

    return student


# 实际使用示例
# Teacher: ResNet50 (25M params), Student: ResNet18 (11M params)
teacher = models.resnet50(weights='DEFAULT')
teacher.fc = nn.Linear(2048, 10)

student = models.resnet18(weights=None)
student.fc = nn.Linear(512, 10)

# 参数量比较
t_params = sum(p.numel() for p in teacher.parameters())
s_params = sum(p.numel() for p in student.parameters())
print(f"Teacher: {t_params:,} params")   # ~25.6M
print(f"Student: {s_params:,} params")   # ~11.2M
print(f"压缩率: {t_params/s_params:.1f}x")

class ResponseBasedDistillation(nn.Module):
    """仅使用最终输出的 Response-based 蒸馏。"""
    def __init__(self, temperature=4.0):
        super().__init__()
        self.T = temperature

    def forward(self, student_logits, teacher_logits):
        student_soft = F.log_softmax(student_logits / self.T, dim=-1)
        teacher_soft = F.softmax(teacher_logits / self.T, dim=-1)
        return F.kl_div(student_soft, teacher_soft, reduction='batchmean') * (self.T ** 2)

class FeatureDistillationLoss(nn.Module):
    """匹配中间层 feature 的蒸馏。"""
    def __init__(self, teacher_channels, student_channels):
        super().__init__()
        # 将 Student feature 投影到 Teacher feature 空间
        self.regressor = nn.Sequential(
            nn.Conv2d(student_channels, teacher_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(teacher_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, student_feat, teacher_feat):
        # 将 Student feature 转换到 Teacher 维度
        projected = self.regressor(student_feat)
        # 用 MSE 损失匹配 feature
        return F.mse_loss(projected, teacher_feat.detach())


class HookBasedDistillation:
    """
    用 forward hook 提取中间层 feature。
    """
    def __init__(self, model, layer_names):
        self.features = {}
        self.hooks = []
        for name, layer in model.named_modules():
            if name in layer_names:
                hook = layer.register_forward_hook(
                    self._make_hook(name)
                )
                self.hooks.append(hook)

    def _make_hook(self, name):
        def hook(module, input, output):
            self.features[name] = output
        return hook

    def remove(self):
        for hook in self.hooks:
            hook.remove()


# 使用示例
teacher = models.resnet50(weights='DEFAULT')
student = models.resnet18(weights=None)

# Teacher: layer3 输出（1024 通道），Student: layer3 输出（256 通道）
teacher_hook = HookBasedDistillation(teacher, ['layer3'])
student_hook = HookBasedDistillation(student, ['layer3'])

feat_distill = FeatureDistillationLoss(
    teacher_channels=1024,
    student_channels=256
)

# 训练时
x = torch.randn(4, 3, 224, 224)
teacher_out = teacher(x)
student_out = student(x)

teacher_feat = teacher_hook.features['layer3']
student_feat = student_hook.features['layer3']

loss_feat = feat_distill(student_feat, teacher_feat)

class RelationalKnowledgeDistillation(nn.Module):
    """
    Relational KD：保留样本对之间的距离关系。
    Student 学习 Teacher embedding 空间的结构。
    """
    def __init__(self, distance_weight=25.0, angle_weight=50.0):
        super().__init__()
        self.dist_w = distance_weight
        self.angle_w = angle_weight

    def pdist(self, e, squared=False, eps=1e-12):
        """成对距离计算。"""
        e_sq = (e ** 2).sum(dim=1)
        prod = e @ e.t()
        res = (e_sq.unsqueeze(1) + e_sq.unsqueeze(0) - 2 * prod).clamp(min=eps)
        if not squared:
            res = res.sqrt()
        return res

    def distance_loss(self, teacher_emb, student_emb):
        """保留距离关系的损失。"""
        t_d = self.pdist(teacher_emb)
        # 归一化
        t_d = t_d / (t_d.mean() + 1e-12)
        s_d = self.pdist(student_emb)
        s_d = s_d / (s_d.mean() + 1e-12)
        return F.smooth_l1_loss(s_d, t_d.detach())

    def angle_loss(self, teacher_emb, student_emb):
        """保留角度关系的损失。"""
        # e_i - e_j 向量之间的角度关系
        td = teacher_emb.unsqueeze(0) - teacher_emb.unsqueeze(1)  # (N, N, D)
        sd = student_emb.unsqueeze(0) - student_emb.unsqueeze(1)

        # cosine similarity
        td_norm = F.normalize(td.view(-1, td.size(-1)), dim=-1)
        sd_norm = F.normalize(sd.view(-1, sd.size(-1)), dim=-1)

        t_angle = (td_norm * td_norm.flip(0)).sum(dim=-1)
        s_angle = (sd_norm * sd_norm.flip(0)).sum(dim=-1)

        return F.smooth_l1_loss(s_angle, t_angle.detach())

    def forward(self, teacher_emb, student_emb):
        loss = self.dist_w * self.distance_loss(teacher_emb, student_emb)
        loss += self.angle_w * self.angle_loss(teacher_emb, student_emb)
        return loss

class AttentionTransfer(nn.Module):
    """Attention Transfer：传递 activation map 的空间模式。"""

    def attention_map(self, feat):
        """
        从 feature map 计算 attention map。
        在每个空间位置上，沿通道方向取平方和的开方。
        """
        return F.normalize(feat.pow(2).mean(1).view(feat.size(0), -1))

    def forward(self, student_feats, teacher_feats):
        """
        student_feats, teacher_feats: feature map 列表
        对每一对计算 AT 损失。
        """
        loss = 0.0
        for s_feat, t_feat in zip(student_feats, teacher_feats):
            s_attn = self.attention_map(s_feat)
            t_attn = self.attention_map(t_feat)
            loss += (s_attn - t_attn.detach()).pow(2).mean()
        return loss

from transformers import (
    DistilBertModel, BertModel,
    DistilBertTokenizer
)
import torch
import torch.nn as nn
import torch.nn.functional as F


class DistilBERTLoss(nn.Module):
    """
    DistilBERT 三重损失：
    1. MLM CE 损失（语言建模）
    2. Soft-target KD 损失（匹配 Teacher logit）
    3. Cosine embedding 损失（hidden state 相似度）
    """
    def __init__(self, temperature=2.0, alpha=0.5, beta=0.1):
        super().__init__()
        self.T = temperature
        self.alpha = alpha  # MLM 权重
        self.beta = beta    # Cosine 损失权重

    def forward(self, student_logits, teacher_logits,
                student_hidden, teacher_hidden, mlm_labels):
        # 1. MLM 损失
        loss_mlm = F.cross_entropy(
            student_logits.view(-1, student_logits.size(-1)),
            mlm_labels.view(-1),
            ignore_index=-100
        )

        # 2. Soft KD 损失
        s_soft = F.log_softmax(student_logits / self.T, dim=-1)
        t_soft = F.softmax(teacher_logits / self.T, dim=-1)
        loss_kd = F.kl_div(s_soft, t_soft, reduction='batchmean') * (self.T ** 2)

        # 3. Cosine embedding 损失（hidden state 相似度）
        loss_cos = 1 - F.cosine_similarity(student_hidden, teacher_hidden, dim=-1).mean()

        return (self.alpha * loss_mlm
                + (1 - self.alpha) * loss_kd
                + self.beta * loss_cos)


# 推理示例
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertModel.from_pretrained('distilbert-base-uncased')

text = "Knowledge distillation is a model compression technique."
inputs = tokenizer(text, return_tensors='pt')

with torch.no_grad():
    outputs = model(**inputs)

print(outputs.last_hidden_state.shape)
# torch.Size([1, 11, 768])

class LLMDistillationTrainer:
    """
    LLM 蒸馏训练器。
    将 Teacher 的 logit 分布传递给 Student。
    """
    def __init__(self, teacher, student, temperature=2.0, alpha=0.5):
        self.teacher = teacher
        self.student = student
        self.T = temperature
        self.alpha = alpha

    def compute_loss(self, input_ids, attention_mask, labels):
        # Teacher 推理（no_grad）
        with torch.no_grad():
            teacher_out = self.teacher(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            teacher_logits = teacher_out.logits  # (B, seq_len, vocab_size)

        # Student 推理
        student_out = self.student(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        student_logits = student_out.logits

        # 1. CE 损失（硬目标）
        shift_logits = student_logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        loss_ce = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=-100
        )

        # 2. KD 损失（软目标）
        # 仅对有效 token 计算
        mask = (shift_labels != -100).float()

        s_log_soft = F.log_softmax(
            shift_logits / self.T, dim=-1
        )
        t_soft = F.softmax(
            teacher_logits[..., :-1, :].contiguous() / self.T, dim=-1
        )

        # 逐 token 的 KL 散度
        kl_per_token = F.kl_div(
            s_log_soft, t_soft,
            reduction='none'
        ).sum(-1)  # (B, seq_len)

        loss_kd = (kl_per_token * mask).sum() / mask.sum()
        loss_kd = loss_kd * (self.T ** 2)

        return self.alpha * loss_ce + (1 - self.alpha) * loss_kd

from transformers import (
    WhisperProcessor, WhisperForConditionalGeneration
)

# Distil-Whisper: Whisper-large-v2 的蒸馏
# large-v2: 1550M params → distil-large-v2: 756M params (2x faster, same WER)
processor = WhisperProcessor.from_pretrained("distil-whisper/distil-large-v2")
model = WhisperForConditionalGeneration.from_pretrained("distil-whisper/distil-large-v2")

# 仅保留 6 个解码器层（原始 32 个 → 2 个）
# 编码器保持不变
print(f"Encoder layers: {len(model.model.encoder.layers)}")
print(f"Decoder layers: {len(model.model.decoder.layers)}")

import torch
import torch.nn as nn
import numpy as np


def get_filter_importance(conv_layer, norm='l1'):
    """计算每个过滤器的重要度（norm）。"""
    weight = conv_layer.weight.data  # (out_ch, in_ch, kH, kW)
    if norm == 'l1':
        return weight.abs().sum(dim=(1, 2, 3))
    elif norm == 'l2':
        return weight.pow(2).sum(dim=(1, 2, 3)).sqrt()
    elif norm == 'gm':
        # 基于 Geometric Median（更 robust）
        return weight.view(weight.size(0), -1).norm(p=2, dim=1)


def prune_conv_layer(conv, keep_ratio=0.5):
    """
    过滤器剪枝：移除重要度较低的过滤器。
    Returns: 新的 Conv2d 层（过滤器数量减少）
    """
    importance = get_filter_importance(conv)
    n_keep = int(conv.out_channels * keep_ratio)
    # 选出重要度排名前 n_keep 的过滤器
    _, indices = importance.topk(n_keep)
    indices, _ = indices.sort()

    # 创建新的 conv 层
    new_conv = nn.Conv2d(
        conv.in_channels,
        n_keep,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        bias=conv.bias is not None
    )
    new_conv.weight.data = conv.weight.data[indices]
    if conv.bias is not None:
        new_conv.bias.data = conv.bias.data[indices]

    return new_conv, indices


class StructuredPruner:
    """ResNet 风格模型的结构化剪枝。"""

    def __init__(self, model, prune_ratio=0.5):
        self.model = model
        self.prune_ratio = prune_ratio

    def compute_layer_importance(self):
        """计算每一层过滤器的重要度。"""
        importance_dict = {}
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Conv2d):
                importance = get_filter_importance(module)
                importance_dict[name] = importance
        return importance_dict

    def global_threshold_prune(self, sparsity=0.5):
        """
        对全部过滤器使用全局阈值进行剪枝。
        移除重要度最低的 sparsity 比例的过滤器。
        """
        all_importances = []
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Conv2d):
                imp = get_filter_importance(module)
                all_importances.append(imp)

        # 将所有过滤器的重要度合并为一个
        all_imp = torch.cat(all_importances)
        threshold = all_imp.kthvalue(int(len(all_imp) * sparsity)).values.item()

        # 移除低于阈值的过滤器
        masks = {}
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Conv2d):
                imp = get_filter_importance(module)
                masks[name] = imp >= threshold

        return masks

class AttentionHeadPruner:
    """
    Transformer 的 Multi-Head Attention 头剪枝。
    Michel et al. (2019): Are Sixteen Heads Really Better than One?
    """

    def compute_head_importance(self, model, dataloader, device='cuda'):
        """计算每一层、每个头的重要度。"""
        head_importance = {}

        model.eval()
        for batch in dataloader:
            inputs = {k: v.to(device) for k, v in batch.items()}

            # 用于保存 Attention 权重的 hook
            attention_weights = {}

            hooks = []
            for name, module in model.named_modules():
                if 'attention' in name.lower() and hasattr(module, 'num_heads'):
                    def hook_fn(m, inp, out, layer_name=name):
                        if hasattr(out, 'attentions') and out.attentions is not None:
                            attention_weights[layer_name] = out.attentions
                    hooks.append(module.register_forward_hook(hook_fn))

            with torch.no_grad():
                outputs = model(**inputs, output_attentions=True)

            # Head importance：attention 权重方差之和
            if hasattr(outputs, 'attentions') and outputs.attentions:
                for layer_idx, attn in enumerate(outputs.attentions):
                    # attn: (batch, heads, seq, seq)
                    head_imp = attn.mean(0).var(-1).mean(-1)  # (heads,)
                    key = f"layer_{layer_idx}"
                    if key not in head_importance:
                        head_importance[key] = head_imp
                    else:
                        head_importance[key] += head_imp

            for hook in hooks:
                hook.remove()

        return head_importance

    def prune_heads(self, model, heads_to_prune):
        """
        移除指定的 head。
        heads_to_prune: {layer_idx: [head_indices]}
        """
        model.prune_heads(heads_to_prune)  # HuggingFace 模型内置功能
        return model


# 使用 HuggingFace transformers
from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained('bert-base-uncased')

# 移除 layer 0 的 head 0、3、5
heads_to_prune = {0: [0, 3, 5], 6: [1, 2]}
model.prune_heads(heads_to_prune)

total_params = sum(p.numel() for p in model.parameters())
print(f"Pruned BERT params: {total_params:,}")

class LayerPruner:
    """Transformer 层剪枝。"""

    def compute_layer_importance_by_gradient(
        self, model, dataloader, device='cuda'
    ):
        """
        基于 gradient 计算每一层的重要度。
        重要度 = |gradient * weight| 的平均值
        """
        model.train()
        layer_importance = {}

        for batch in dataloader:
            inputs = {k: v.to(device) for k, v in batch.items()}
            outputs = model(**inputs)
            loss = outputs.loss
            loss.backward()

            for i, layer in enumerate(model.encoder.layer):
                grad_sum = 0.0
                param_count = 0
                for param in layer.parameters():
                    if param.grad is not None:
                        grad_sum += (param.grad * param.data).abs().sum().item()
                        param_count += param.numel()

                key = f"layer_{i}"
                imp = grad_sum / max(param_count, 1)
                if key not in layer_importance:
                    layer_importance[key] = 0.0
                layer_importance[key] += imp

            model.zero_grad()

        return layer_importance

    def drop_layers(self, model, layers_to_drop):
        """移除指定层，并用剩余层重新组装模型。"""
        import copy
        new_model = copy.deepcopy(model)

        remaining = [
            layer for i, layer in enumerate(new_model.encoder.layer)
            if i not in layers_to_drop
        ]
        new_model.encoder.layer = nn.ModuleList(remaining)
        new_model.config.num_hidden_layers = len(remaining)

        return new_model

import torch.nn.utils.prune as prune

model = models.resnet18(weights='DEFAULT')

# 对特定层应用非结构化 L1 剪枝
conv = model.layer1[0].conv1
prune.l1_unstructured(conv, name='weight', amount=0.3)
# 将 weight 的 30% 置零（使用 mask）

# 确认
print(f"Sparsity: {(conv.weight == 0).float().mean():.2f}")

# 结构化剪枝（过滤器单位）
prune.ln_structured(conv, name='weight', amount=0.3, n=2, dim=0)

# 永久应用（移除 mask，实际将 weight 置零）
prune.remove(conv, 'weight')

# 对多个层重复应用
parameters_to_prune = []
for module in model.modules():
    if isinstance(module, nn.Conv2d):
        parameters_to_prune.append((module, 'weight'))

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.4,  # 剪除全部参数的 40%
)

# 确认整体稀疏度
total_zero = sum(
    (m.weight == 0).sum().item()
    for m in model.modules() if isinstance(m, nn.Conv2d)
)
total_params = sum(
    m.weight.numel()
    for m in model.modules() if isinstance(m, nn.Conv2d)
)
print(f"Global sparsity: {total_zero/total_params:.2%}")

class MagnitudePruner:
    """基于幅值的非结构化剪枝。"""

    def __init__(self, model, sparsity=0.5):
        self.model = model
        self.sparsity = sparsity
        self.masks = {}

    def compute_global_threshold(self):
        """计算全部权重的全局阈值。"""
        all_weights = []
        for name, param in self.model.named_parameters():
            if 'weight' in name and param.dim() > 1:
                all_weights.append(param.data.abs().view(-1))

        all_weights = torch.cat(all_weights)
        threshold = all_weights.kthvalue(
            int(len(all_weights) * self.sparsity)
        ).values.item()
        return threshold

    def apply_pruning(self):
        """用全局阈值生成剪枝 mask。"""
        threshold = self.compute_global_threshold()

        for name, param in self.model.named_parameters():
            if 'weight' in name and param.dim() > 1:
                mask = (param.data.abs() >= threshold).float()
                self.masks[name] = mask
                param.data *= mask  # 将低于阈值的权重置零

        total_zeros = sum((m == 0).sum().item() for m in self.masks.values())
        total_params = sum(m.numel() for m in self.masks.values())
        print(f"实际稀疏度: {total_zeros/total_params:.2%}")

    def apply_masks(self):
        """反向传播后重新应用 mask（防止已置零的权重被梯度复原）。"""
        for name, param in self.model.named_parameters():
            if name in self.masks:
                param.data *= self.masks[name]

class GradualMagnitudePruner:
    """
    在训练过程中逐步提高稀疏度的剪枝。
    Zhu & Gupta (2017): To Prune, or Not to Prune
    """
    def __init__(self, model, initial_sparsity=0.0, final_sparsity=0.8,
                 begin_step=0, end_step=1000, frequency=100):
        self.model = model
        self.initial_sparsity = initial_sparsity
        self.final_sparsity = final_sparsity
        self.begin_step = begin_step
        self.end_step = end_step
        self.frequency = frequency
        self.masks = {}
        self._init_masks()

    def _init_masks(self):
        for name, param in self.model.named_parameters():
            if 'weight' in name and param.dim() > 1:
                self.masks[name] = torch.ones_like(param.data)

    def compute_sparsity(self, step):
        """计算当前步数下的目标稀疏度（cubic schedule）。"""
        if step < self.begin_step:
            return self.initial_sparsity
        if step > self.end_step:
            return self.final_sparsity

        # Cubic decay schedule
        pct_done = (step - self.begin_step) / (self.end_step - self.begin_step)
        sparsity = (
            self.final_sparsity
            + (self.initial_sparsity - self.final_sparsity)
            * (1 - pct_done) ** 3
        )
        return sparsity

    def step(self, global_step):
        """在训练步骤中更新剪枝。"""
        if (global_step % self.frequency != 0 or
                global_step < self.begin_step or
                global_step > self.end_step):
            return

        target_sparsity = self.compute_sparsity(global_step)
        self._update_masks(target_sparsity)

    def _update_masks(self, sparsity):
        """按照当前稀疏度目标更新 mask。"""
        for name, param in self.model.named_parameters():
            if name not in self.masks:
                continue

            # 只考虑当前存活的权重（排除已被剪枝的部分）
            alive = param.data[self.masks[name].bool()]

            if len(alive) == 0:
                continue

            n_prune = int(sparsity * param.data.numel())
            if n_prune == 0:
                continue

            # 找出全体中排名最低的 n_prune 个
            threshold = param.data.abs().view(-1).kthvalue(n_prune).values.item()
            self.masks[name] = (param.data.abs() > threshold).float()
            param.data *= self.masks[name]

class ALBERTEncoder(nn.Module):
    """
    ALBERT 风格权重共享。
    将同一个 Transformer 层重复执行 N 次。
    """
    def __init__(self, hidden_size=768, num_heads=12,
                 intermediate_size=3072, num_layers=12):
        super().__init__()
        # 只定义一个 Transformer 层
        self.shared_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=intermediate_size,
            batch_first=True,
            norm_first=True,
        )
        self.num_layers = num_layers
        self.norm = nn.LayerNorm(hidden_size)

    def forward(self, x, src_key_padding_mask=None):
        # 将同一层重复执行 num_layers 次
        for _ in range(self.num_layers):
            x = self.shared_layer(x, src_key_padding_mask=src_key_padding_mask)
        return self.norm(x)

    def count_parameters(self):
        # 实际参数量（因共享，只统计一次）
        return sum(p.numel() for p in self.parameters())


# 参数量比较
bert_encoder = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(768, 12, 3072, batch_first=True, norm_first=True),
    num_layers=12
)
albert_encoder = ALBERTEncoder(num_layers=12)

bert_params = sum(p.numel() for p in bert_encoder.parameters())
albert_params = albert_encoder.count_parameters()

print(f"BERT encoder: {bert_params:,}")     # ~85M
print(f"ALBERT encoder: {albert_params:,}") # ~7M
print(f"压缩率: {bert_params/albert_params:.1f}x")  # ~12x

class FactorizedEmbedding(nn.Module):
    """
    ALBERT 的嵌入因式分解。
    vocab_size x H → (vocab_size x E) + (E x H)
    E << H 以大幅减少参数量。
    """
    def __init__(self, vocab_size, embedding_size=128, hidden_size=768):
        super().__init__()
        self.word_embeddings = nn.Embedding(vocab_size, embedding_size)
        # E → H 线性变换
        self.embedding_projection = nn.Linear(embedding_size, hidden_size, bias=False)

    def forward(self, input_ids):
        embed = self.word_embeddings(input_ids)        # (B, seq, E)
        return self.embedding_projection(embed)         # (B, seq, H)

# 参数量比较（vocab_size=30000）
standard_embed = nn.Embedding(30000, 768)             # 23.0M params
factorized_embed = FactorizedEmbedding(30000, 128, 768)  # 3.97M params

s_params = sum(p.numel() for p in standard_embed.parameters())
f_params = sum(p.numel() for p in factorized_embed.parameters())
print(f"Standard: {s_params:,} → Factorized: {f_params:,}")
print(f"节省: {(s_params - f_params)/s_params:.1%}")

import torch
import torch.nn as nn
import torch.nn.functional as F


# 候选运算
PRIMITIVES = [
    'none',           # 无连接
    'skip_connect',   # 恒等函数
    'sep_conv_3x3',   # 3x3 Separable Conv
    'sep_conv_5x5',   # 5x5 Separable Conv
    'dil_conv_3x3',   # 3x3 Dilated Conv
    'dil_conv_5x5',   # 5x5 Dilated Conv
    'avg_pool_3x3',   # 3x3 Average Pool
    'max_pool_3x3',   # 3x3 Max Pool
]


class MixedOperation(nn.Module):
    """
    DARTS：用多种运算的加权和表示一条边。
    alpha（架构参数）决定每种运算的权重。
    """
    def __init__(self, C, stride, primitives):
        super().__init__()
        self.ops = nn.ModuleList([
            self._build_op(primitive, C, stride)
            for primitive in primitives
        ])
        # 可学习的架构参数
        self.arch_params = nn.Parameter(
            torch.randn(len(primitives)) * 1e-3
        )

    def _build_op(self, primitive, C, stride):
        if primitive == 'none':
            return nn.Sequential()  # Zero operation
        elif primitive == 'skip_connect':
            return nn.Identity() if stride == 1 else nn.Sequential(
                nn.AvgPool2d(stride, stride, padding=0),
                nn.Conv2d(C, C, 1, bias=False),
                nn.BatchNorm2d(C)
            )
        elif primitive == 'sep_conv_3x3':
            return nn.Sequential(
                nn.Conv2d(C, C, 3, stride=stride, padding=1, groups=C, bias=False),
                nn.Conv2d(C, C, 1, bias=False),
                nn.BatchNorm2d(C),
                nn.ReLU(inplace=True),
            )
        elif primitive == 'avg_pool_3x3':
            return nn.AvgPool2d(3, stride=stride, padding=1)
        elif primitive == 'max_pool_3x3':
            return nn.MaxPool2d(3, stride=stride, padding=1)
        else:
            return nn.Identity()

    def forward(self, x):
        # 用 Softmax 计算权重，再对各运算加权求和
        weights = F.softmax(self.arch_params, dim=0)
        results = []
        for w, op in zip(weights, self.ops):
            try:
                result = op(x)
                if result.shape == x.shape or len(results) == 0:
                    results.append(w * result)
                else:
                    results.append(w * result)
            except Exception:
                pass

        if results:
            # 所有运算结果的加权和
            output = results[0]
            for r in results[1:]:
                if r.shape == output.shape:
                    output = output + r
            return output
        return x


class DARTSCell(nn.Module):
    """DARTS 单元：由多个 Mixed Operation 组成的 DAG。"""
    def __init__(self, num_nodes, C, stride=1):
        super().__init__()
        self.num_nodes = num_nodes
        self.ops = nn.ModuleList()

        # 为每一对节点生成 Mixed Operation
        for i in range(num_nodes):
            for j in range(i):
                self.ops.append(
                    MixedOperation(C, stride, PRIMITIVES)
                )

    def forward(self, *inputs):
        states = list(inputs)

        op_idx = 0
        for i in range(self.num_nodes):
            # 汇总来自所有先前状态的输入
            node_input = sum(
                self.ops[op_idx + j](states[j])
                for j in range(len(states))
            )
            op_idx += len(states)
            states.append(node_input)

        return states[-1]

    def get_arch_params(self):
        """返回架构参数。"""
        return [op.arch_params for op in self.ops]


def darts_discrete(cell):
    """
    将连续架构转换为离散结构。
    在每条边上选择权重最高的运算。
    """
    for op in cell.ops:
        weights = F.softmax(op.arch_params, dim=0)
        best_op_idx = weights.argmax().item()
        print(f"  Selected: {PRIMITIVES[best_op_idx]} "
              f"(weight: {weights[best_op_idx]:.3f})")

# EfficientNet 的 MBConv 模块（NAS 找到的核心结构）
class MBConvBlock(nn.Module):
    """
    Mobile Inverted Bottleneck Convolution.
    EfficientNet 的基本构建模块。
    """
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, expand_ratio=6, se_ratio=0.25):
        super().__init__()
        self.use_residual = (stride == 1 and in_channels == out_channels)
        hidden_dim = int(in_channels * expand_ratio)

        layers = []
        # Expansion phase (1x1 conv)
        if expand_ratio != 1:
            layers.extend([
                nn.Conv2d(in_channels, hidden_dim, 1, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.SiLU(),
            ])

        # Depthwise conv
        layers.extend([
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size,
                      stride=stride,
                      padding=kernel_size // 2,
                      groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.SiLU(),
        ])

        # Squeeze-and-Excitation（NAS 发现的重要模块）
        se_channels = max(1, int(in_channels * se_ratio))
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(hidden_dim, se_channels, 1),
            nn.SiLU(),
            nn.Conv2d(se_channels, hidden_dim, 1),
            nn.Sigmoid(),
        )

        # Output phase
        layers.extend([
            nn.Conv2d(hidden_dim, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
        ])

        self.conv = nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv[:-2](x) if len(self.conv) > 2 else self.conv(x)
        # 应用 SE 模块
        out_se = self.se(out)
        out = out * out_se
        out = self.conv[-2:](out)

        if self.use_residual:
            return x + out
        return out

class OFALayer(nn.Module):
    """
    Once-for-All：支持多种卷积核大小的弹性层。
    训练时：随机选择卷积核大小 → 训练子网络
    部署时：只使用特定的卷积核大小
    """
    def __init__(self, channels, max_kernel=7):
        super().__init__()
        # 只训练最大卷积核大小的 conv
        self.max_conv = nn.Conv2d(
            channels, channels,
            kernel_size=max_kernel,
            padding=max_kernel // 2,
            groups=channels, bias=False
        )
        self.bn = nn.BatchNorm2d(channels)
        self.act = nn.ReLU(inplace=True)
        self.active_kernel = max_kernel

    def set_active_kernel(self, kernel_size):
        """设置当前激活的卷积核大小。"""
        assert kernel_size <= self.max_conv.kernel_size[0]
        self.active_kernel = kernel_size

    def forward(self, x):
        if self.active_kernel == self.max_conv.kernel_size[0]:
            weight = self.max_conv.weight
        else:
            # 从中心裁剪出 active_kernel 大小
            center = self.max_conv.kernel_size[0] // 2
            half = self.active_kernel // 2
            weight = self.max_conv.weight[
                :, :,
                center - half: center + half + 1,
                center - half: center + half + 1
            ]

        padding = self.active_kernel // 2
        out = F.conv2d(x, weight, padding=padding, groups=x.size(1))
        return self.act(self.bn(out))


# Progressive shrinking 训练（OFA 的核心）
def progressive_shrinking_train(model, dataloader, kernels=[7, 5, 3]):
    """
    第 1 阶段：训练最大卷积核（7）
    第 2 阶段：在 7、5 卷积核中随机采样训练
    第 3 阶段：在 7、5、3 卷积核中随机采样训练
    """
    for stage, active_kernels in enumerate(
        [kernels[:1], kernels[:2], kernels]
    ):
        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
        print(f"Stage {stage+1}: kernels = {active_kernels}")

        for epoch in range(5):
            for images, labels in dataloader:
                # 随机选择卷积核大小
                k = active_kernels[torch.randint(len(active_kernels), (1,)).item()]

                # 设置所有层的激活卷积核
                for module in model.modules():
                    if isinstance(module, OFALayer):
                        module.set_active_kernel(k)

                outputs = model(images)
                loss = F.cross_entropy(outputs, labels)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

class ModelCompressionPipeline:
    """
    知识蒸馏 → 剪枝 → 量化的集成流水线。
    """

    def __init__(self, teacher, student_arch, num_classes):
        self.teacher = teacher
        self.num_classes = num_classes

        # 初始化 Student 模型
        self.student = student_arch

    def step1_distillation(self, train_loader, val_loader,
                           epochs=30, device='cuda'):
        """第 1 步：通过知识蒸馏进行基础训练。"""
        print("=== Step 1: Knowledge Distillation ===")
        self.student = train_with_distillation(
            self.teacher, self.student, train_loader,
            num_epochs=epochs, temperature=4.0, alpha=0.5,
            device=device
        )

    def step2_pruning(self, train_loader, sparsity=0.5,
                      finetune_epochs=10, device='cuda'):
        """第 2 步：渐进式剪枝 + 微调。"""
        print(f"=== Step 2: Pruning (target sparsity: {sparsity:.0%}) ===")
        pruner = GradualMagnitudePruner(
            self.student,
            initial_sparsity=0.0,
            final_sparsity=sparsity,
            begin_step=0,
            end_step=len(train_loader) * finetune_epochs,
            frequency=100
        )

        optimizer = torch.optim.Adam(
            self.student.parameters(), lr=1e-4
        )
        self.student.to(device)
        step = 0

        for epoch in range(finetune_epochs):
            for images, labels in train_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = self.student(images)
                loss = F.cross_entropy(outputs, labels)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                pruner.step(step)
                pruner.apply_masks()  # 保持剪枝 mask
                step += 1

    def step3_quantization(self, calibration_loader, device='cpu'):
        """第 3 步：训练后量化（Post-Training Quantization）。"""
        print("=== Step 3: Quantization ===")
        self.student.eval().to(device)

        # 准备静态量化
        self.student.qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
        torch.ao.quantization.prepare(self.student, inplace=True)

        # 用校准数据确定 scale/zero_point
        with torch.no_grad():
            for images, _ in calibration_loader:
                self.student(images.to(device))

        # 应用量化
        torch.ao.quantization.convert(self.student, inplace=True)
        print("Quantization complete (INT8)")
        return self.student

    def compare_models(self, val_loader, device='cuda'):
        """比较 Teacher、Student、Pruned、Quantized 模型的性能。"""
        def eval_model(model, loader, dev):
            model.eval().to(dev)
            correct, total = 0, 0
            with torch.no_grad():
                for images, labels in loader:
                    images, labels = images.to(dev), labels.to(dev)
                    preds = model(images).argmax(1)
                    correct += (preds == labels).sum().item()
                    total += labels.size(0)
            return correct / total

        teacher_acc = eval_model(self.teacher, val_loader, device)
        student_acc = eval_model(self.student, val_loader, 'cpu')

        t_params = sum(p.numel() for p in self.teacher.parameters())
        s_params = sum(p.numel() for p in self.student.parameters())

        print(f"\n{'='*50}")
        print(f"Teacher  | Params: {t_params:>10,} | Acc: {teacher_acc:.4f}")
        print(f"Student  | Params: {s_params:>10,} | Acc: {student_acc:.4f}")
        print(f"压缩率: {t_params/s_params:.1f}x | 性能保留率: {student_acc/teacher_acc:.1%}")