大模型压缩与效率优化：量化、剪枝与蒸馏的协同策略

引言：大模型部署的效率困境

当前，GPT-4、LLaMA等百亿甚至万亿参数大模型在各类任务上展现出卓越性能，但巨大的计算开销和内存占用严重限制了其实际部署。单一优化技术往往只能在特定维度带来有限改进，而量化、剪枝与蒸馏的协同策略正在成为解决这一困境的关键突破。本文将深入探讨这三种核心技术的协同优化机制，并提供完整的代码实现。

理论基础：三大压缩技术的互补性分析

1. 量化的数值效率优化

量化通过降低权重和激活值的数值精度来减少存储和计算开销，但可能引入量化误差和精度损失。

2. 剪枝的结构稀疏性优化

剪枝通过移除冗余参数或结构来简化模型架构，但可能破坏模型的连通性和表达能力。

3. 蒸馏的知识传递优化

蒸馏通过将大模型知识迁移到小模型来保持性能，但受限于教师模型的表达能力和学生模型的容量。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.quantization import QuantStub, DeQuantStub, prepare, convert
import numpy as np
from typing import Dict, List, Tuple, Optional, Callable
import copy
import matplotlib.pyplot as plt
from dataclasses import dataclass

@dataclass
class CompressionConfig:
    """统一的压缩配置"""
    # 量化配置
    quantization_bits: int = 8
    quantization_scheme: str = 'symmetric'  # symmetric, asymmetric
    per_channel_quantization: bool = True
    
    # 剪枝配置
    pruning_method: str = 'magnitude'  # magnitude, movement, lottery
    pruning_rate: float = 0.5
    pruning_iterations: int = 10
    pruning_frequency: int = 100
    
    # 蒸馏配置
    distillation_temperature: float = 3.0
    distillation_alpha: float = 0.7
    distillation_loss_weights: Dict[str, float] = None
    
    # 协同策略配置
    execution_order: List[str] = None  # ['prune', 'quantize', 'distill']
    iterative_cycles: int = 3
    
    def __post_init__(self):
        if self.distillation_loss_weights is None:
            self.distillation_loss_weights = {
                'kl_divergence': 1.0,
                'attention_loss': 0.5,
                'hidden_state_loss': 0.3
            }
        if self.execution_order is None:
            self.execution_order = ['prune', 'quantize', 'distill']

协同压缩框架设计与实现

1. 统一压缩框架架构

class UnifiedCompressionFramework:
    """量化、剪枝、蒸馏协同压缩框架"""
    
    def __init__(self, 
                 model: nn.Module,
                 config: CompressionConfig,
                 device: str = 'cuda'):
        self.original_model = model
        self.config = config
        self.device = device
        
        # 初始化各模块
        self.quantizer = AdaptiveQuantizer(config)
        self.pruner = IntelligentPruner(config)
        self.distiller = MultiHeadDistiller(config)
        
        # 状态跟踪
        self.compression_history = []
        self.performance_metrics = []
        
        # 创建压缩模型副本
        self.compressed_model = self._create_compressed_copy()
        
    def _create_compressed_copy(self) -> nn.Module:
        """创建可压缩的模型副本"""
        model_copy = copy.deepcopy(self.original_model)
        model_copy.to(self.device)
        
        # 为量化准备模型
        if 'quantize' in self.config.execution_order:
            model_copy = self.quantizer.prepare_for_quantization(model_copy)
        
        return model_copy
    
    def compress(self, 
                 train_loader: torch.utils.data.DataLoader,
                 val_loader: torch.utils.data.DataLoader,
                 num_epochs: int = 10) -> nn.Module:
        """执行协同压缩流程"""
        
        print("开始协同压缩流程...")
        print(f"执行顺序: {self.config.execution_order}")
        print(f"迭代周期: {self.config.iterative_cycles}")
        
        for cycle in range(self.config.iterative_cycles):
            print(f"\n=== 压缩周期 {cycle + 1}/{self.config.iterative_cycles} ===")
            
            # 按照配置顺序执行压缩操作
            for operation in self.config.execution_order:
                if operation == 'prune':
                    print("执行结构化剪枝...")
                    self.compressed_model = self._prune_phase(
                        self.compressed_model, train_loader, cycle
                    )
                    
                elif operation == 'quantize':
                    print("执行自适应量化...")
                    self.compressed_model = self._quantize_phase(
                        self.compressed_model, train_loader, val_loader
                    )
                    
                elif operation == 'distill':
                    print("执行知识蒸馏...")
                    self.compressed_model = self._distill_phase(
                        self.compressed_model, train_loader, val_loader, num_epochs
                    )
                
                # 评估压缩效果
                metrics = self._evaluate_compression(self.compressed_model, val_loader)
                self.compression_history.append({
                    'cycle': cycle,
                    'operation': operation,
                    'metrics': metrics
                })
                
                print(f"操作后指标 - 准确率: {metrics['accuracy']:.4f}, "
                      f"模型大小: {metrics['model_size_mb']:.2f}MB, "
                      f"推理延迟: {metrics['inference_latency']:.4f}s")
        
        # 最终优化和导出
        self.compressed_model = self._finalize_compression(self.compressed_model)
        
        return self.compressed_model
    
    def _prune_phase(self, 
                     model: nn.Module,
                     data_loader: torch.utils.data.DataLoader,
                     cycle: int) -> nn.Module:
        """剪枝阶段"""
        # 动态调整剪枝率
        current_prune_rate = self.config.pruning_rate * (1.0 - 0.1 * cycle)
        
        # 执行结构化剪枝
        pruned_model = self.pruner.prune_model(
            model=model,
            data_loader=data_loader,
            prune_rate=current_prune_rate,
            method=self.config.pruning_method
        )
        
        # 重训练恢复精度
        pruned_model = self._fine_tune_pruned_model(pruned_model, data_loader)
        
        return pruned_model
    
    def _quantize_phase(self,
                       model: nn.Module,
                       train_loader: torch.utils.data.DataLoader,
                       val_loader: torch.utils.data.DataLoader) -> nn.Module:
        """量化阶段"""
        # 校准量化参数
        calibrated_model = self.quantizer.calibrate(
            model=model,
            data_loader=train_loader,
            num_batches=100
        )
        
        # 执行量化
        quantized_model = self.quantizer.quantize(
            model=calibrated_model,
            scheme=self.config.quantization_scheme,
            bits=self.config.quantization_bits,
            per_channel=self.config.per_channel_quantization
        )
        
        # 量化感知训练
        if self.config.quantization_bits < 8:
            quantized_model = self._quantization_aware_training(
                quantized_model, train_loader, val_loader, epochs=3
            )
        
        return quantized_model
    
    def _distill_phase(self,
                      student_model: nn.Module,
                      train_loader: torch.utils.data.DataLoader,
                      val_loader: torch.utils.data.DataLoader,
                      num_epochs: int) -> nn.Module:
        """蒸馏阶段"""
        # 创建教师模型（使用原始模型或上一阶段模型）
        teacher_model = copy.deepcopy(self.original_model)
        teacher_model.eval()
        
        # 执行多头部蒸馏
        distilled_model = self.distiller.distill(
            teacher_model=teacher_model,
            student_model=student_model,
            train_loader=train_loader,
            val_loader=val_loader,
            temperature=self.config.distillation_temperature,
            alpha=self.config.distillation_alpha,
            num_epochs=num_epochs
        )
        
        return distilled_model
    
    def _fine_tune_pruned_model(self,
                               model: nn.Module,
                               data_loader: torch.utils.data.DataLoader,
                               epochs: int = 3) -> nn.Module:
        """微调剪枝后的模型"""
        model.train()
        optimizer = optim.Adam(model.parameters(), lr=1e-4)
        criterion = nn.CrossEntropyLoss()
        
        for epoch in range(epochs):
            total_loss = 0
            for batch_idx, (data, target) in enumerate(data_loader):
                data, target = data.to(self.device), target.to(self.device)
                
                optimizer.zero_grad()
                output = model(data)
                loss = criterion(output, target)
                loss.backward()
                
                # 对剪枝后的权重应用掩码
                self.pruner.apply_weight_mask(model)
                
                optimizer.step()
                total_loss += loss.item()
        
        return model
    
    def _quantization_aware_training(self,
                                    model: nn.Module,
                                    train_loader: torch.utils.data.DataLoader,
                                    val_loader: torch.utils.data.DataLoader,
                                    epochs: int = 5) -> nn.Module:
        """量化感知训练"""
        model.train()
        optimizer = optim.Adam(model.parameters(), lr=1e-4)
        criterion = nn.CrossEntropyLoss()
        
        for epoch in range(epochs):
            model.train()
            for batch_idx, (data, target) in enumerate(train_loader):
                data, target = data.to(self.device), target.to(self.device)
                
                optimizer.zero_grad()
                output = model(data)
                loss = criterion(output, target)
                loss.backward()
                optimizer.step()
            
            # 验证
            model.eval()
            correct = 0
            total = 0
            with torch.no_grad():
                for data, target in val_loader:
                    data, target = data.to(self.device), target.to(self.device)
                    output = model(data)
                    pred = output.argmax(dim=1, keepdim=True)
                    correct += pred.eq(target.view_as(pred)).sum().item()
                    total += target.size(0)
            
            accuracy = 100. * correct / total
            print(f"QAT Epoch {epoch+1}: Accuracy = {accuracy:.2f}%")
        
        return model
    
    def _evaluate_compression(self,
                            model: nn.Module,
                            val_loader: torch.utils.data.DataLoader) -> Dict:
        """评估压缩效果"""
        model.eval()
        
        # 计算准确率
        correct = 0
        total = 0
        inference_time = 0
        
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(self.device), target.to(self.device)
                
                # 测量推理时间
                start_time = torch.cuda.Event(enable_timing=True)
                end_time = torch.cuda.Event(enable_timing=True)
                
                start_time.record()
                output = model(data)
                end_time.record()
                torch.cuda.synchronize()
                
                inference_time += start_time.elapsed_time(end_time) / 1000.0
                
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()
                total += target.size(0)
        
        # 计算模型大小
        param_size = 0
        for param in model.parameters():
            param_size += param.numel() * param.element_size()
        buffer_size = 0
        for buffer in model.buffers():
            buffer_size += buffer.numel() * buffer.element_size()
        
        model_size_mb = (param_size + buffer_size) / 1024**2
        
        return {
            'accuracy': correct / total,
            'model_size_mb': model_size_mb,
            'inference_latency': inference_time / len(val_loader),
            'compression_ratio': self._calculate_compression_ratio(model)
        }
    
    def _calculate_compression_ratio(self, model: nn.Module) -> float:
        """计算压缩率"""
        original_params = sum(p.numel() for p in self.original_model.parameters())
        compressed_params = sum(p.numel() for p in model.parameters())
        
        # 考虑量化带来的存储节省
        if hasattr(model, '_quantized'):
            # 量化模型参数大小估算
            compressed_params *= self.config.quantization_bits / 32
        
        return original_params / max(compressed_params, 1)
    
    def _finalize_compression(self, model: nn.Module) -> nn.Module:
        """最终优化和导出"""
        # 应用最终优化
        model.eval()
        
        # 对于量化模型，转换为推理模式
        if hasattr(model, '_quantized'):
            model = torch.quantization.convert(model, inplace=False)
        
        # 移除训练特定组件
        if hasattr(model, 'qconfig'):
            delattr(model, 'qconfig')
        
        return model

核心组件深度实现

1. 自适应量化器实现

class AdaptiveQuantizer:
    """自适应量化器：支持混合精度和动态范围量化"""
    
    def __init__(self, config: CompressionConfig):
        self.config = config
        self.observer_stats = {}
        
    def prepare_for_quantization(self, model: nn.Module) -> nn.Module:
        """为量化准备模型"""
        # 设置量化配置
        model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
        
        # 插入量化stub
        model.quant = QuantStub()
        model.dequant = DeQuantStub()
        
        # 准备量化
        model_prepared = torch.quantization.prepare(model, inplace=False)
        
        return model_prepared
    
    def calibrate(self, 
                  model: nn.Module,
                  data_loader: torch.utils.data.DataLoader,
                  num_batches: int = 100) -> nn.Module:
        """校准量化参数"""
        model.eval()
        
        print("开始量化校准...")
        batch_count = 0
        
        with torch.no_grad():
            for data, _ in data_loader:
                data = data.to(next(model.parameters()).device)
                
                # 前向传播收集统计信息
                model(data)
                
                batch_count += 1
                if batch_count >= num_batches:
                    break
        
        print(f"校准完成，处理了 {batch_count} 个批次")
        return model
    
    def quantize(self,
                 model: nn.Module,
                 scheme: str = 'symmetric',
                 bits: int = 8,
                 per_channel: bool = True) -> nn.Module:
        """执行量化"""
        print(f"执行量化: scheme={scheme}, bits={bits}, per_channel={per_channel}")
        
        # 转换为量化模型
        model_quantized = torch.quantization.convert(model, inplace=False)
        model_quantized._quantized = True
        
        # 记录量化信息
        self._analyze_quantization_effect(model_quantized)
        
        return model_quantized
    
    def _analyze_quantization_effect(self, model: nn.Module):
        """分析量化效果"""
        weight_precision = {}
        activation_ranges = {}
        
        for name, module in model.named_modules():
            if hasattr(module, 'weight_fake_quant'):
                # 分析权重量化
                weight_data = module.weight().detach().cpu().numpy()
                weight_precision[name] = {
                    'min': weight_data.min(),
                    'max': weight_data.max(),
                    'std': weight_data.std(),
                    'unique_values': len(np.unique(weight_data))
                }
        
        print("量化效果分析:")
        for name, stats in list(weight_precision.items())[:5]:  # 显示前5层
            print(f"  {name}: unique values = {stats['unique_values']}")
    
    def mixed_precision_quantization(self,
                                   model: nn.Module,
                                   sensitivity_analysis: Dict[str, float]) -> nn.Module:
        """混合精度量化"""
        print("执行混合精度量化...")
        
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
                # 根据敏感度分配比特数
                sensitivity = sensitivity_analysis.get(name, 0.5)
                
                if sensitivity > 0.8:
                    bits = 8  # 高敏感层使用8比特
                elif sensitivity > 0.5:
                    bits = 4  # 中等敏感层使用4比特
                else:
                    bits = 2  # 低敏感层使用2比特
                
                # 应用量化（简化实现）
                self._apply_layer_quantization(module, bits)
        
        return model
    
    def _apply_layer_quantization(self, module: nn.Module, bits: int):
        """对单层应用量化"""
        if bits == 32:
            return  # 保持浮点
        
        # 模拟量化操作
        scale = 2 ** (bits - 1) - 1
        
        if hasattr(module, 'weight'):
            weight = module.weight.data
            # 线性量化
            weight_quantized = torch.round(weight * scale) / scale
            module.weight.data = weight_quantized

2. 智能剪枝器实现

class IntelligentPruner:
    """智能剪枝器：支持结构化剪枝和重要性评分"""
    
    def __init__(self, config: CompressionConfig):
        self.config = config
        self.masks = {}
        self.importance_scores = {}
        
    def prune_model(self,
                   model: nn.Module,
                   data_loader: torch.utils.data.DataLoader,
                   prune_rate: float = 0.5,
                   method: str = 'magnitude') -> nn.Module:
        """执行模型剪枝"""
        print(f"执行剪枝: method={method}, rate={prune_rate}")
        
        # 计算重要性分数
        self._compute_importance_scores(model, data_loader, method)
        
        # 生成剪枝掩码
        self._generate_pruning_masks(model, prune_rate, method)
        
        # 应用剪枝
        pruned_model = self._apply_pruning_masks(model)
        
        # 计算稀疏度
        sparsity = self._calculate_sparsity(pruned_model)
        print(f"剪枝后稀疏度: {sparsity:.2%}")
        
        return pruned_model
    
    def _compute_importance_scores(self,
                                 model: nn.Module,
                                 data_loader: torch.utils.data.DataLoader,
                                 method: str):
        """计算参数重要性分数"""
        self.importance_scores = {}
        
        if method == 'magnitude':
            # 基于权重大小的评分
            for name, param in model.named_parameters():
                if 'weight' in name:
                    self.importance_scores[name] = torch.abs(param.data)
        
        elif method == 'movement':
            # 基于权重变化的评分
            original_params = {name: param.clone() for name, param in model.named_parameters()}
            
            # 执行一步训练
            model.train()
            optimizer = optim.SGD(model.parameters(), lr=0.01)
            criterion = nn.CrossEntropyLoss()
            
            for batch_idx, (data, target) in enumerate(data_loader):
                if batch_idx >= 10:  # 仅使用少量批次
                    break
                
                data, target = data.to(next(model.parameters()).device), target.to(next(model.parameters()).device)
                optimizer.zero_grad()
                output = model(data)
                loss = criterion(output, target)
                loss.backward()
                optimizer.step()
            
            # 计算变化量
            for name, param in model.named_parameters():
                if name in original_params and 'weight' in name:
                    movement = torch.abs(param.data - original_params[name])
                    self.importance_scores[name] = movement
        
        elif method == 'lottery':
            # 彩票假说剪枝
            for name, param in model.named_parameters():
                if 'weight' in name:
                    # 基于梯度信息评分
                    if param.grad is not None:
                        score = torch.abs(param.data * param.grad)
                    else:
                        score = torch.abs(param.data)
                    self.importance_scores[name] = score
    
    def _generate_pruning_masks(self,
                              model: nn.Module,
                              prune_rate: float,
                              method: str):
        """生成剪枝掩码"""
        self.masks = {}
        
        for name, param in model.named_parameters():
            if 'weight' in name and name in self.importance_scores:
                importance = self.importance_scores[name]
                
                # 计算阈值
                if method == 'global':
                    # 全局阈值
                    all_scores = torch.cat([s.flatten() for s in self.importance_scores.values()])
                    threshold = torch.kthvalue(all_scores, int(all_scores.numel() * prune_rate))[0]
                else:
                    # 层独立阈值
                    scores_flat = importance.flatten()
                    k = max(1, int(scores_flat.numel() * prune_rate))
                    threshold = torch.kthvalue(scores_flat, k)[0]
                
                # 生成掩码
                mask = (importance > threshold).float()
                self.masks[name] = mask
    
    def _apply_pruning_masks(self, model: nn.Module) -> nn.Module:
        """应用剪枝掩码"""
        for name, param in model.named_parameters():
            if name in self.masks:
                mask = self.masks[name]
                param.data = param.data * mask
                # 存储掩码用于反向传播
                param.mask = mask
        
        return model
    
    def apply_weight_mask(self, model: nn.Module):
        """在优化步骤后重新应用权重掩码"""
        for name, param in model.named_parameters():
            if hasattr(param, 'mask'):
                param.data = param.data * param.mask
    
    def _calculate_sparsity(self, model: nn.Module) -> float:
        """计算模型稀疏度"""
        total_params = 0
        zero_params = 0
        
        for name, param in model.named_parameters():
            if 'weight' in name:
                total_params += param.numel()
                zero_params += (param.data == 0).sum().item()
        
        return zero_params / total_params if total_params > 0 else 0
    
    def structured_pruning(self,
                         model: nn.Module,
                         prune_rate: float = 0.3) -> nn.Module:
        """结构化剪枝：移除整个通道或神经元"""
        print("执行结构化剪枝...")
        
        # 计算通道重要性
        channel_importance = {}
        
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d):
                # 基于L1范数评估通道重要性
                weights = module.weight.data
                importance = torch.norm(weights, p=1, dim=(1, 2, 3))
                channel_importance[name] = importance
        
        # 移除最不重要的通道
        for name, module in model.named_modules():
            if name in channel_importance:
                importance = channel_importance[name]
                num_channels = len(importance)
                num_prune = int(num_channels * prune_rate)
                
                if num_prune > 0:
                    # 获取要保留的通道索引
                    _, keep_indices = torch.topk(importance, num_channels - num_prune)
                    
                    # 创建新的卷积层
                    old_weight = module.weight.data
                    new_weight = old_weight[keep_indices]
                    
                    if module.bias is not None:
                        old_bias = module.bias.data
                        new_bias = old_bias[keep_indices]
                    
                    # 更新下一层的输入通道
                    # 这里需要复杂的层间连接调整，简化处理
        
        return model

3. 多头部蒸馏器实现

class MultiHeadDistiller:
    """多头部蒸馏器：支持注意力、隐藏状态和特征图蒸馏"""
    
    def __init__(self, config: CompressionConfig):
        self.config = config
        
    def distill(self,
                teacher_model: nn.Module,
                student_model: nn.Module,
                train_loader: torch.utils.data.DataLoader,
                val_loader: torch.utils.data.DataLoader,
                temperature: float = 3.0,
                alpha: float = 0.7,
                num_epochs: int = 10) -> nn.Module:
        """执行知识蒸馏"""
        print("开始知识蒸馏...")
        print(f"温度: {temperature}, alpha: {alpha}")
        
        device = next(teacher_model.parameters()).device
        teacher_model.eval()
        student_model.train()
        
        optimizer = optim.Adam(student_model.parameters(), lr=1e-4)
        
        for epoch in range(num_epochs):
            total_loss = 0
            total_kl_loss = 0
            total_ce_loss = 0
            total_att_loss = 0
            
            for batch_idx, (data, target) in enumerate(train_loader):
                data, target = data.to(device), target.to(device)
                
                optimizer.zero_grad()
                
                # 教师模型输出
                with torch.no_grad():
                    teacher_outputs = self._get_teacher_outputs(teacher_model, data)
                
                # 学生模型输出
                student_outputs = self._get_student_outputs(student_model, data)
                
                # 计算蒸馏损失
                loss, losses_dict = self._compute_distillation_loss(
                    teacher_outputs=teacher_outputs,
                    student_outputs=student_outputs,
                    targets=target,
                    temperature=temperature,
                    alpha=alpha
                )
                
                loss.backward()
                optimizer.step()
                
                total_loss += loss.item()
                total_kl_loss += losses_dict['kl_loss'].item() if losses_dict['kl_loss'] is not None else 0
                total_ce_loss += losses_dict['ce_loss'].item()
                total_att_loss += losses_dict['attention_loss'].item() if losses_dict['attention_loss'] is not None else 0
            
            # 验证
            val_accuracy = self._evaluate_model(student_model, val_loader)
            
            print(f"Epoch {epoch+1}/{num_epochs}: "
                  f"Loss={total_loss/len(train_loader):.4f}, "
                  f"KL Loss={total_kl_loss/len(train_loader):.4f}, "
                  f"CE Loss={total_ce_loss/len(train_loader):.4f}, "
                  f"Val Acc={val_accuracy:.2f}%")
        
        return student_model
    
    def _get_teacher_outputs(self, model: nn.Module, inputs: torch.Tensor) -> Dict:
        """获取教师模型的中间输出"""
        outputs = {}
        
        # 存储中间层输出（Hook机制）
        hooks = []
        layer_outputs = {}
        
        def hook_fn(name):
            def hook(module, input, output):
                layer_outputs[name] = output.detach()
            return hook
        
        # 注册钩子（简化实现）
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
                hooks.append(module.register_forward_hook(hook_fn(name)))
        
        # 前向传播
        final_output = model(inputs)
        
        # 移除钩子
        for hook in hooks:
            hook.remove()
        
        outputs['logits'] = final_output
        outputs['hidden_states'] = layer_outputs
        
        return outputs
    
    def _get_student_outputs(self, model: nn.Module, inputs: torch.Tensor) -> Dict:
        """获取学生模型的中间输出"""
        # 类似教师模型的方法
        return self._get_teacher_outputs(model, inputs)
    
    def _compute_distillation_loss(self,
                                 teacher_outputs: Dict,
                                 student_outputs: Dict,
                                 targets: torch.Tensor,
                                 temperature: float,
                                 alpha: float) -> Tuple[torch.Tensor, Dict]:
        """计算多头部蒸馏损失"""
        losses = {}
        
        # 1. KL散度损失（软标签蒸馏）
        if 'logits' in teacher_outputs and 'logits' in student_outputs:
            teacher_logits = teacher_outputs['logits']
            student_logits = student_outputs['logits']
            
            # 应用温度缩放
            teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
            student_probs = F.log_softmax(student_logits / temperature, dim=-1)
            
            kl_loss = F.kl_div(student_probs, teacher_probs, reduction='batchmean') * (temperature ** 2)
            losses['kl_loss'] = kl_loss
        else:
            kl_loss = torch.tensor(0.0, device=targets.device)
            losses['kl_loss'] = None
        
        # 2. 交叉熵损失（硬标签）
        ce_loss = F.cross_entropy(student_outputs.get('logits', torch.randn_like(targets)), targets)
        losses['ce_loss'] = ce_loss
        
        # 3. 注意力蒸馏损失
        attention_loss = self._compute_attention_loss(teacher_outputs, student_outputs)
        losses['attention_loss'] = attention_loss
        
        # 4. 隐藏状态蒸馏损失
        hidden_loss = self._compute_hidden_state_loss(teacher_outputs, student_outputs)
        losses['hidden_loss'] = hidden_loss
        
        # 总损失
        total_loss = 0
        if kl_loss is not None:
            total_loss += alpha * kl_loss
        total_loss += (1 - alpha) * ce_loss
        total_loss += 0.5 * attention_loss
        total_loss += 0.3 * hidden_loss
        
        return total_loss, losses
    
    def _compute_attention_loss(self, 
                              teacher_outputs: Dict, 
                              student_outputs: Dict) -> torch.Tensor:
        """计算注意力蒸馏损失"""
        # 简化实现：假设模型有关注力输出
        teacher_attentions = teacher_outputs.get('attentions', [])
        student_attentions = student_outputs.get('attentions', [])
        
        if len(teacher_attentions) == 0 or len(student_attentions) == 0:
            return torch.tensor(0.0, device=next(iter(teacher_outputs.values())).device)
        
        att_loss = 0
        for t_att, s_att in zip(teacher_attentions, student_attentions):
            # 计算注意力矩阵的MSE损失
            att_loss += F.mse_loss(s_att, t_att)
        
        return att_loss / len(teacher_attentions)
    
    def _compute_hidden_state_loss(self,
                                 teacher_outputs: Dict,
                                 student_outputs: Dict) -> torch.Tensor:
        """计算隐藏状态蒸馏损失"""
        teacher_hiddens = teacher_outputs.get('hidden_states', {})
        student_hiddens = student_outputs.get('hidden_states', {})
        
        if not teacher_hiddens or not student_hiddens:
            return torch.tensor(0.0, device=next(iter(teacher_outputs.values())).device)
        
        hidden_loss = 0
        count = 0
        
        for name in teacher_hiddens:
            if name in student_hiddens:
                t_hidden = teacher_hiddens[name]
                s_hidden = student_hiddens[name]
                
                # 调整维度匹配
                if t_hidden.shape != s_hidden.shape:
                    # 简单的调整策略
                    if t_hidden.dim() == 4 and s_hidden.dim() == 4:  # Conv层
                        s_hidden = F.adaptive_avg_pool2d(s_hidden, t_hidden.shape[2:])
                    elif t_hidden.dim() == 2 and s_hidden.dim() == 2:  # Linear层
                        if s_hidden.shape[1] < t_hidden.shape[1]:
                            padding = torch.zeros(s_hidden.shape[0], 
                                                 t_hidden.shape[1] - s_hidden.shape[1],
                                                 device=s_hidden.device)
                            s_hidden = torch.cat([s_hidden, padding], dim=1)
                
                # 计算MSE损失
                hidden_loss += F.mse_loss(s_hidden, t_hidden)
                count += 1
        
        return hidden_loss / max(count, 1)
    
    def _evaluate_model(self, model: nn.Module, data_loader: torch.utils.data.DataLoader) -> float:
        """评估模型准确率"""
        model.eval()
        device = next(model.parameters()).device
        
        correct = 0
        total = 0
        
        with torch.no_grad():
            for data, target in data_loader:
                data, target = data.to(device), target.to(device)
                outputs = model(data)
                pred = outputs.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()
                total += target.size(0)
        
        model.train()
        return 100. * correct / total

实验与分析：协同策略效果验证

1. 协同压缩效果对比实验

class CompressionExperiment:
    """压缩策略对比实验"""
    
    def __init__(self):
        self.results = {}
        
    def run_comparison(self, 
                      model: nn.Module,
                      train_loader: torch.utils.data.DataLoader,
                      val_loader: torch.utils.data.DataLoader):
        """运行不同压缩策略的对比"""
        
        strategies = {
            'baseline': {'quantize': False, 'prune': False, 'distill': False},
            'quant_only': {'quantize': True, 'prune': False, 'distill': False},
            'prune_only': {'quantize': False, 'prune': True, 'distill': False},
            'distill_only': {'quantize': False, 'prune': False, 'distill': True},
            'quant_prune': {'quantize': True, 'prune': True, 'distill': False},
            'quant_distill': {'quantize': True, 'prune': False, 'distill': True},
            'prune_distill': {'quantize': False, 'prune': True, 'distill': True},
            'full_coordination': {'quantize': True, 'prune': True, 'distill': True}
        }
        
        for strategy_name, strategy_config in strategies.items():
            print(f"\n=== 测试策略: {strategy_name} ===")
            
            # 创建压缩配置
            config = CompressionConfig(
                execution_order=self._get_execution_order(strategy_config)
            )
            
            # 应用压缩
            framework = UnifiedCompressionFramework(model, config)
            compressed_model = framework.compress(
                train_loader, val_loader, num_epochs=5
            )
            
            # 评估结果
            metrics = self._evaluate_strategy(compressed_model, val_loader)
            self.results[strategy_name] = metrics
            
            print(f"准确率: {metrics['accuracy']:.4f}")
            print(f"压缩率: {metrics['compression_ratio']:.2f}x")
            print(f"推理延迟: {metrics['inference_latency']:.4f}s")
        
        # 可视化结果
        self._visualize_results()
    
    def _get_execution_order(self, strategy_config: Dict) -> List[str]:
        """根据策略配置获取执行顺序"""
        order = []
        if strategy_config['prune']:
            order.append('prune')
        if strategy_config['quantize']:
            order.append('quantize')
        if strategy_config['distill']:
            order.append('distill')
        return order
    
    def _evaluate_strategy(self, 
                          model: nn.Module,
                          val_loader: torch.utils.data.DataLoader) -> Dict:
        """评估策略效果"""
        device = next(model.parameters()).device
        model.eval()
        
        # 推理延迟
        total_time = 0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(device), target.to(device)
                
                start_time = torch.cuda.Event(enable_timing=True)
                end_time = torch.cuda.Event(enable_timing=True)
                
                start_time.record()
                output = model(data)
                end_time.record()
                torch.cuda.synchronize()
                
                total_time += start_time.elapsed_time(end_time) / 1000.0
                
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()
                total += target.size(0)
        
        # 计算模型大小
        param_size = sum(p.numel() * p.element_size() for p in model.parameters())
        buffer_size = sum(b.numel() * b.element_size() for b in model.buffers())
        model_size_mb = (param_size + buffer_size) / 1024**2
        
        return {
            'accuracy': correct / total,
            'model_size_mb': model_size_mb,
            'inference_latency': total_time / len(val_loader),
            'compression_ratio': 100.0 / model_size_mb  # 相对于100MB基准
        }
    
    def _visualize_results(self):
        """可视化对比结果"""
        strategies = list(self.results.keys())
        accuracies = [self.results[s]['accuracy'] for s in strategies]
        compressions = [self.results[s]['compression_ratio'] for s in strategies]
        latencies = [self.results[s]['inference_latency'] for s in strategies]
        
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        # 准确率对比
        axes[0].bar(strategies, accuracies)
        axes[0].set_title('Accuracy Comparison')
        axes[0].set_ylabel('Accuracy')
        axes[0].tick_params(axis='x', rotation=45)
        
        # 压缩率对比
        axes[1].bar(strategies, compressions)
        axes[1].set_title('Compression Ratio Comparison')
        axes[1].set_ylabel('Compression Ratio (x)')
        axes[1].tick_params(axis='x', rotation=45)
        
        # 延迟对比
        axes[2].bar(strategies, latencies)
        axes[2].set_title('Inference Latency Comparison')
        axes[2].set_ylabel('Latency (s)')
        axes[2].tick_params(axis='x', rotation=45)
        
        plt.tight_layout()
        plt.show()

2. 实际案例：BERT模型压缩

class BERTCompressor:
    """BERT模型的协同压缩实现"""
    
    def __init__(self, model_name: str = 'bert-base-uncased'):
        from transformers import BertModel, BertTokenizer
        
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.original_model = BertModel.from_pretrained(model_name)
        
        # 创建分类头
        self.classifier = nn.Linear(self.original_model.config.hidden_size, 2)
    
    def create_compressed_bert(self, 
                              config: CompressionConfig) -> nn.Module:
        """创建压缩后的BERT模型"""
        
        # 创建学生模型（更小的架构）
        student_config = self._create_student_config()
        student_model = BertModel(student_config)
        
        # 创建完整分类模型
        full_model = nn.Sequential(
            student_model,
            nn.Dropout(0.1),
            self.classifier
        )
        
        # 应用协同压缩
        framework = UnifiedCompressionFramework(full_model, config)
        
        # 注意：这里需要实际的数据加载器
        # compressed_model = framework.compress(train_loader, val_loader)
        
        return full_model
    
    def _create_student_config(self):
        """创建学生模型配置"""
        from transformers import BertConfig
        
        # 缩小模型尺寸
        original_config = self.original_model.config
        
        student_config = BertConfig(
            vocab_size=original_config.vocab_size,
            hidden_size=256,  # 原始为768
            num_hidden_layers=6,  # 原始为12
            num_attention_heads=8,  # 原始为12
            intermediate_size=512,  # 原始为3072
            max_position_embeddings=original_config.max_position_embeddings,
            type_vocab_size=original_config.type_vocab_size
        )
        
        return student_config

高级优化：动态协同策略

1. 自适应压缩调度器

class AdaptiveCompressionScheduler:
    """自适应压缩调度器：根据模型状态动态调整压缩策略"""
    
    def __init__(self, 
                 model: nn.Module,
                 target_metrics: Dict[str, float]):
        self.model = model
        self.target_metrics = target_metrics
        self.compression_state = {
            'current_stage': 'initial',
            'prune_completed': False,
            'quantize_completed': False,
            'distill_completed': False
        }
        
    def schedule_compression(self,
                            current_metrics: Dict[str, float]) -> Dict:
        """根据当前指标调度下一步压缩操作"""
        
        # 计算与目标的差距
        accuracy_gap = current_metrics.get('accuracy', 1.0) - self.target_metrics.get('accuracy', 0.9)
        size_gap = current_metrics.get('model_size_mb', 1000) - self.target_metrics.get('model_size_mb', 100)
        
        # 决策逻辑
        if accuracy_gap > 0.1 and size_gap > 200:
            # 精度和大小都差很多，先蒸馏
            action = 'distill'
            intensity = 0.8
        elif accuracy_gap > 0.05 and size_gap > 100:
            # 中等差距，剪枝加蒸馏
            action = 'prune_distill'
            intensity = 0.6
        elif size_gap > 50 and accuracy_gap < 0.02:
            # 大小差距大但精度好，量化
            action = 'quantize'
            intensity = 0.7
        else:
            # 接近目标，精细调整
            action = 'fine_tune'
            intensity = 0.3
        
        return {
            'action': action,
            'intensity': intensity,
            'parameters': self._get_parameters_for_action(action, intensity)
        }
    
    def _get_parameters_for_action(self, action: str, intensity: float) -> Dict:
        """获取对应操作的参数"""
        if action == 'distill':
            return {
                'temperature': 3.0 * intensity,
                'alpha': 0.7 * intensity,
                'epochs': int(10 * intensity)
            }
        elif action == 'prune':
            return {
                'rate': 0.5 * intensity,
                'method': 'magnitude'
            }
        elif action == 'quantize':
            return {
                'bits': max(2, int(8 * (1 - intensity))),
                'scheme': 'symmetric'
            }
        else:
            return {}

结论与展望

1. 协同策略的核心优势

• 互补性优化：量化减少数值精度开销，剪枝减少结构冗余，蒸馏保持知识完整性
• 协同增益：组合策略的效果优于单一策略的简单叠加
• 自适应能力：可根据硬件约束和精度要求动态调整压缩策略

2. 未来研究方向

• 神经架构搜索与压缩的结合：自动寻找最优压缩架构
• 硬件感知压缩：针对特定硬件（如NPU、FPGA）的定制化压缩
• 动态推理压缩：根据输入复杂度动态调整模型大小
• 联邦学习中的压缩：在分布式场景下的高效模型压缩