模型量化技术详解

# 量化误差累积示例
def error_accumulation_simulation():
    layers = 50
    layer_error = 0.01  # 每层1%的量化误差

    # 误差累积
    total_error = 1 - (1 - layer_error) ** layers
    print(f"总误差累积: {total_error:.2%}")
    # 输出: 总误差累积: 39.50%

import numpy as np
import matplotlib.pyplot as plt

def analyze_activation_distribution(activations):
    """分析激活值分布"""
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))

    # 线性尺度
    axes[0].hist(activations.flatten(), bins=100)
    axes[0].set_title('Linear Scale')
    axes[0].set_xlabel('Activation Value')

    # 对数尺度
    axes[1].hist(np.abs(activations.flatten()), bins=100, log=True)
    axes[1].set_title('Log Scale (Long-tail Distribution)')
    axes[1].set_xlabel('|Activation Value|')

    plt.tight_layout()
    plt.show()

    # 统计离群值
    p99 = np.percentile(np.abs(activations), 99)
    p100 = np.abs(activations).max()
    print(f"99th percentile: {p99:.4f}")
    print(f"Max value: {p100:.4f}")
    print(f"Outlier ratio: {p100/p99:.2f}x")

# 实际应用
# activations = model.get_intermediate_outputs()
# analyze_activation_distribution(activations)

def layer_sensitivity_analysis(model, calibration_data):
    """分析各层对量化的敏感度"""
    sensitivity_scores = {}

    for name, layer in model.named_modules():
        if isinstance(layer, (nn.Linear, nn.Conv2d)):
            # 量化单层并评估精度损失
            original_output = layer(calibration_data)

            # 量化权重
            weight_scale = layer.weight.abs().max() / 127
            weight_quant = torch.round(layer.weight / weight_scale).clamp(-128, 127)
            weight_dequant = weight_quant * weight_scale

            # 评估精度损失
            layer.weight.data = weight_dequant
            quantized_output = layer(calibration_data)
            layer.weight.data = weight_dequant  # 恢复原始权重

            # 计算敏感度
            error = torch.norm(original_output - quantized_output) / torch.norm(original_output)
            sensitivity_scores[name] = error.item()

    return sensitivity_scores

# 可视化敏感度
def plot_sensitivity(sensitivity_scores):
    sorted_scores = sorted(sensitivity_scores.items(), key=lambda x: x[1], reverse=True)
    layers, scores = zip(*sorted_scores)

    plt.figure(figsize=(12, 6))
    plt.barh(range(len(layers)), scores)
    plt.yticks(range(len(layers)), layers)
    plt.xlabel('Quantization Sensitivity')
    plt.title('Layer-wise Quantization Sensitivity')
    plt.gca().invert_yaxis()
    plt.tight_layout()
    plt.show()

def activation_range_analysis(model, dataloader):
    """分析激活值的动态范围"""
    activation_ranges = {}

    def hook_fn(name):
        def hook(module, input, output):
            if name not in activation_ranges:
                activation_ranges[name] = []
            activation_ranges[name].append(output.detach())
        return hook

    # 注册hook
    hooks = []
    for name, layer in model.named_modules():
        if isinstance(layer, (nn.Linear, nn.Conv2d, nn.ReLU)):
            hook = layer.register_forward_hook(hook_fn(name))
            hooks.append(hook)

    # 收集激活值
    model.eval()
    with torch.no_grad():
        for batch, _ in dataloader:
            _ = model(batch)

    # 移除hooks
    for hook in hooks:
        hook.remove()

    # 分析范围
    for name, activations in activation_ranges.items():
        activations = torch.cat(activations)
        min_val = activations.min().item()
        max_val = activations.max().item()
        mean_val = activations.mean().item()
        std_val = activations.std().item()

        print(f"{name}:")
        print(f"  Range: [{min_val:.4f}, {max_val:.4f}]")
        print(f"  Mean: {mean_val:.4f}, Std: {std_val:.4f}")
        print(f"  Dynamic Range: {max_val - min_val:.4f}")

import torch
import torch.nn as nn
from torch.quantization import quantize_dynamic

def apply_dynamic_quantization(model):
    """应用动态量化"""
    # 指定要量化的层类型
    model_quantized = quantize_dynamic(
        model,
        {nn.Linear, nn.LSTM, nn.GRU},  # 量化的模块类型
        dtype=torch.qint8  # 量化数据类型
    )
    return model_quantized

# 完整示例
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建并量化模型
model = SimpleModel()
model_quantized = apply_dynamic_quantization(model)

# 对比模型大小
def get_model_size(model):
    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()
    return (param_size + buffer_size) / 1024 / 1024  # MB

original_size = get_model_size(model)
quantized_size = get_model_size(model_quantized)
compression_ratio = original_size / quantized_size

print(f"原始模型大小: {original_size:.2f} MB")
print(f"量化后模型大小: {quantized_size:.2f} MB")
print(f"压缩比: {compression_ratio:.2f}x")

import torch
import torch.nn as nn
from torch.quantization import prepare, convert

def minmax_calibration(model, calibration_loader):
    """MinMax校准"""
    # 配置量化
    model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

    # 准备量化（插入观察者）
    model_prepared = prepare(model, inplace=False)

    # 校准（运行推理收集统计数据）
    model_prepared.eval()
    with torch.no_grad():
        for images, _ in calibration_loader:
            model_prepared(images)

    # 转换为量化模型
    model_quantized = convert(model_prepared, inplace=False)

    return model_quantized

# 完整流程
def static_quantization_workflow():
    # 1. 准备模型
    model = SimpleModel()
    model.eval()

    # 2. 准备校准数据（通常使用训练集的一个子集）
    calibration_data = [...]  # 校准数据集

    # 3. 应用静态量化
    model_quantized = minmax_calibration(model, calibration_data)

    # 4. 评估量化精度
    accuracy = evaluate(model_quantized, test_loader)
    print(f"量化后精度: {accuracy:.2%}")

    return model_quantized

import numpy as np
from scipy.stats import entropy

def kl_divergence(p, q):
    """计算KL散度"""
    return entropy(p, q)

def find_optimal_scale(fp32_tensor, int8_tensor, bins=2048):
    """使用KL散度寻找最优scale"""
    # 计算FP32和INT8的直方图
    fp32_hist, fp32_edges = np.histogram(fp32_tensor, bins=bins, density=True)
    int8_hist, int8_edges = np.histogram(int8_tensor, bins=bins, density=True)

    # 尝试不同的scale值
    best_scale = None
    best_kl = float('inf')

    for scale in np.linspace(0.1, 2.0, 100):
        # 使用当前scale量化
        quantized = np.round(fp32_tensor / scale).clip(-128, 127)
        dequantized = quantized * scale

        # 计算KL散度
        dequant_hist, _ = np.histogram(dequantized, bins=bins, density=True)
        kl = kl_divergence(fp32_hist, dequant_hist)

        if kl < best_kl:
            best_kl = kl
            best_scale = scale

    return best_scale

def kl_calibration(model, calibration_loader):
    """KL散度校准"""
    # 收集激活值统计
    activation_stats = collect_activation_statistics(model, calibration_loader)

    # 为每层计算最优scale
    quantization_params = {}
    for name, stats in activation_stats.items():
        optimal_scale = find_optimal_scale(stats['fp32'], stats['int8'])
        quantization_params[name] = {
            'scale': optimal_scale,
            'zero_point': 0  # 对称量化
        }

    # 应用量化参数
    apply_quantization_params(model, quantization_params)

    return model

def percentile_calibration(activations, percentile=99.99):
    """百分位校准"""
    # 计算指定百分位数的值
    abs_activations = np.abs(activations)
    threshold = np.percentile(abs_activations, percentile)

    # 截断离群值
    clipped_activations = np.clip(activations, -threshold, threshold)

    # 基于截断后的范围计算scale
    scale = threshold / 127.0

    return scale

def apply_percentile_calibration(model, calibration_loader, percentile=99.99):
    """应用百分位校准"""
    # 收集激活值
    activation_ranges = collect_activations(model, calibration_loader)

    # 为每层计算scale
    for name, activations in activation_ranges.items():
        scale = percentile_calibration(activations, percentile)
        set_layer_scale(model, name, scale)

    return model

def activation_equalization(model, calibration_loader):
    """激活值均衡化"""
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            # 获取输入输出通道的权重范围
            weight = module.weight.data
            input_scale = weight.abs().mean(dim=(1, 2, 3))
            output_scale = weight.abs().mean(dim=(0, 2, 3))

            # 计算均衡因子
            equalization_factor = torch.sqrt(input_scale / (output_scale + 1e-6))

            # 应用均衡
            module.weight.data = weight * equalization_factor[None, :, None, None]

            # 在下一层插入反缩放
            next_module = get_next_module(model, name)
            if next_module is not None:
                if hasattr(next_module, 'weight'):
                    next_module.weight.data = next_module.weight.data / equalization_factor[None, :, None, None]

    return model

def bias_correction(fp32_model, int8_model, calibration_loader):
    """偏置校正"""
    fp32_model.eval()
    int8_model.eval()

    for name, module in int8_model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)) and module.bias is not None:
            # 收集FP32和INT8的输出差异
            fp32_module = get_module_by_name(fp32_model, name)

            bias_diff = 0
            count = 0

            with torch.no_grad():
                for batch, _ in calibration_loader:
                    fp32_output = fp32_module(batch)
                    int8_output = module(batch)

                    bias_diff += (fp32_output - int8_output).mean(dim=0)
                    count += 1

            # 校正偏置
            avg_bias_diff = bias_diff / count
            module.bias.data += avg_bias_diff

    return int8_model

def weight_permutation_for_quantization(weight):
    """权重排列优化"""
    # PCA降维排列
    weight_flat = weight.view(weight.size(0), -1)

    # 计算PCA
    U, S, V = torch.svd(weight_flat)

    # 重新排列权重
    weight_permuted = torch.mm(weight_flat, V)

    return weight_permuted.view_as(weight)

import torch
import torch.nn as nn

class FakeQuantOp(torch.autograd.Function):
    """伪量化操作"""

    @staticmethod
    def forward(ctx, x, scale, zero_point, qmin, qmax):
        """
        前向传播：量化 + 反量化
        模拟量化误差，但保持数据类型为float
        """
        # 量化
        x_quant = torch.clamp(
            x / scale + zero_point,
            qmin, qmax
        )
        # 反量化
        x_dequant = (x_quant - zero_point) * scale
        return x_dequant

    @staticmethod
    def backward(ctx, grad_output):
        """
        反向传播：直通估计器（Straight-Through Estimator）
        直接传递梯度，忽略量化操作的梯度
        """
        return grad_output, None, None, None, None

def fake_quantize(x, scale, zero_point, bits=8):
    """应用伪量化"""
    qmin = -2**(bits-1)
    qmax = 2**(bits-1) - 1
    return FakeQuantOp.apply(x, scale, zero_point, qmin, qmax)

class LearnableQuantization(nn.Module):
    """可学习的量化参数"""

    def __init__(self, init_scale=1.0, bits=8):
        super().__init__()
        self.bits = bits
        qmin = -2**(bits-1)
        qmax = 2**(bits-1) - 1
        self.register_buffer('qmin', torch.tensor(qmin, dtype=torch.float32))
        self.register_buffer('qmax', torch.tensor(qmax, dtype=torch.float32))

        # 可学习的scale和zero_point
        self.scale = nn.Parameter(torch.tensor(init_scale))
        self.zero_point = nn.Parameter(torch.tensor(0.0))

    def forward(self, x):
        return fake_quantize(x, self.scale, self.zero_point, self.bits)

import torch
import torch.nn as nn
from torch.quantization import prepare_qat, convert

class QATModel(nn.Module):
    """量化感知训练模型"""

    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model
        self.quantization_config = {}

    def forward(self, x):
        return self.base_model(x)

def qat_training_workflow(base_model, train_loader, val_loader, epochs=10):
    """完整的QAT训练流程"""

    # 1. 准备QAT模型
    model = base_model
    model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

    # 2. 插入伪量化节点
    model_prepared = prepare_qat(model, inplace=False)

    # 3. 微调训练
    optimizer = torch.optim.SGD(model_prepared.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    best_accuracy = 0
    for epoch in range(epochs):
        model_prepared.train()
        train_loss = 0
        train_correct = 0
        train_total = 0

        for batch_idx, (inputs, targets) in enumerate(train_loader):
            optimizer.zero_grad()
            outputs = model_prepared(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = outputs.max(1)
            train_total += targets.size(0)
            train_correct += predicted.eq(targets).sum().item()

        # 验证
        model_prepared.eval()
        val_loss = 0
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for inputs, targets in val_loader:
                outputs = model_prepared(inputs)
                loss = criterion(outputs, targets)
                val_loss += loss.item()
                _, predicted = outputs.max(1)
                val_total += targets.size(0)
                val_correct += predicted.eq(targets).sum().item()

        train_acc = 100. * train_correct / train_total
        val_acc = 100. * val_correct / val_total

        print(f'Epoch {epoch+1}/{epochs}:')
        print(f'  Train Loss: {train_loss/(len(train_loader)):.4f}, Train Acc: {train_acc:.2f}%')
        print(f'  Val Loss: {val_loss/(len(val_loader)):.4f}, Val Acc: {val_acc:.2f}%')

        # 保存最佳模型
        if val_acc > best_accuracy:
            best_accuracy = val_acc
            torch.save(model_prepared.state_dict(), 'best_qat_model.pth')

    # 4. 转换为真正的量化模型
    model_prepared.load_state_dict(torch.load('best_qat_model.pth'))
    model_quantized = convert(model_prepared, inplace=False)

    return model_quantized

def layerwise_finetuning(model, train_loader, sensitive_layers):
    """逐层微调敏感层"""
    for layer_name in sensitive_layers:
        print(f"Fine-tuning layer: {layer_name}")

        # 冻结其他层
        for name, param in model.named_parameters():
            if layer_name not in name:
                param.requires_grad = False

        # 只训练当前层
        optimizer = torch.optim.SGD(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=0.0001
        )

        # 训练几个epoch
        for epoch in range(3):
            for inputs, targets in train_loader:
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = nn.CrossEntropyLoss()(outputs, targets)
                loss.backward()
                optimizer.step()

        # 解冻所有层
        for param in model.parameters():
            param.requires_grad = True

    return model

class LSQQuantization(nn.Module):
    """学习步长量化"""

    def __init__(self, bits=8, gradient_penalty=0.01):
        super().__init__()
        self.bits = bits
        self.gradient_penalty = gradient_penalty
        qmin = -2**(bits-1)
        qmax = 2**(bits-1) - 1
        self.register_buffer('qmin', torch.tensor(qmin, dtype=torch.float32))
        self.register_buffer('qmax', torch.tensor(qmax, dtype=torch.float32))

        # 初始化步长
        self.step_size = nn.Parameter(torch.tensor(1.0))

    def forward(self, x):
        # 计算scale
        scale = self.step_size / self.qmax

        # 量化
        x_quant = torch.clamp(
            torch.round(x / scale) + self.zero_point,
            self.qmin, self.qmax
        )

        # 反量化
        x_dequant = (x_quant - self.zero_point) * scale

        return x_dequant

    def extra_repr(self):
        return f'bits={self.bits}, step_size={self.step_size.item():.4f}'

def mixed_precision_qat(model, layer_precision_config):
    """
    layer_precision_config: {
        'conv1': 8,
        'conv2': 4,
        'fc': 8
    }
    """
    for name, module in model.named_modules():
        if name in layer_precision_config:
            bits = layer_precision_config[name]
            qconfig = torch.quantization.QConfig(
                activation=torch.quantization.MinMaxObserver.with_args(dtype=torch.quint8),
                weight=torch.quantization.MinMaxObserver.with_args(dtype=torch.qint8)
            )
            module.qconfig = qconfig

    return prepare_qat(model)

def distillation_qat(student_model, teacher_model, train_loader, temperature=3.0, alpha=0.5):
    """知识蒸馏辅助QAT"""
    optimizer = torch.optim.SGD(student_model.parameters(), lr=0.001)

    teacher_model.eval()
    student_model.train()

    for inputs, targets in train_loader:
        optimizer.zero_grad()

        # 教师模型输出（soft targets）
        with torch.no_grad():
            teacher_outputs = teacher_model(inputs)

        # 学生模型输出
        student_outputs = student_model(inputs)

        # 蒸馏损失
        distill_loss = nn.KLDivLoss(reduction='batchmean')(
            nn.functional.log_softmax(student_outputs / temperature, dim=1),
            nn.functional.softmax(teacher_outputs / temperature, dim=1)
        ) * (temperature ** 2)

        # 硬标签损失
        hard_loss = nn.CrossEntropyLoss()(student_outputs, targets)

        # 组合损失
        loss = alpha * distill_loss + (1 - alpha) * hard_loss

        loss.backward()
        optimizer.step()

    return student_model

import tensorrt as trt
import torch
from torch.utils.data import DataLoader

class TensorRTEntropyCalibrator(trt.IInt8EntropyCalibrator2):
    """TensorRT INT8校准器"""

    def __init__(self, dataloader, cache_file='calibration.cache'):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.dataloader = dataloader
        self.cache_file = cache_file
        self.data_iter = iter(dataloader)

    def get_batch_size(self):
        return self.dataloader.batch_size

    def get_batch(self, names):
        try:
            batch, _ = next(self.data_iter)
            batch = batch.numpy()
            return [batch]
        except StopIteration:
            return None

    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, 'rb') as f:
                return f.read()
        return None

    def write_calibration_cache(self, cache):
        with open(self.cache_file, 'wb') as f:
            f.write(cache)

def build_tensorrt_engine(onnx_model_path, calibration_dataloader):
    """构建TensorRT INT8引擎"""
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.ONNXParser(network, logger)

    # 解析ONNX模型
    with open(onnx_model_path, 'rb') as f:
        parser.parse(f.read())

    # 创建builder配置
    config = builder.create_builder_config()
    config.set_flag(trt.BuilderFlag.INT8)

    # 设置校准器
    calibrator = TensorRTEntropyCalibrator(calibration_dataloader)
    config.int8_calibrator = calibrator

    # 设置最大工作空间
    config.max_workspace_size = 1 << 30  # 1GB

    # 构建引擎
    engine = builder.build_engine(network, config)

    return engine

def build_tensorrt_dynamic_quant(onnx_model_path):
    """TensorRT动态量化"""
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.ONNXParser(network, logger)

    # 解析ONNX模型
    with open(onnx_model_path, 'rb') as f:
        parser.parse(f.read())

    # 配置
    config = builder.create_builder_config()

    # 启用动态量化
    config.set_flag(trt.BuilderFlag.INT8)
    config.set_flag(trt.BuilderFlag.FP16)  # 混合精度

    # 设置精度
    for i in range(network.num_inputs):
        input_tensor = network.get_input(i)
        network.set_dynamic_range(input_tensor, -127.0, 127.0)

    engine = builder.build_engine(network, config)
    return engine

import onnx
from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantType
from onnxruntime.quantization.calibrate import CalibrationDataReader

def onnx_dynamic_quantization(onnx_model_path, output_path):
    """ONNX动态量化"""
    quantize_dynamic(
        onnx_model_path,
        output_path,
        weight_type=QuantType.QUInt8  # 或QuantType.QInt8
    )
    print(f"量化模型已保存到: {output_path}")

class CalibrationDataReader(CalibrationDataReader):
    """校准数据读取器"""

    def __init__(self, dataloader):
        self.dataloader = dataloader
        self.iterator = iter(dataloader)

    def get_next(self):
        try:
            batch, _ = next(self.iterator)
            return {'input': batch.numpy()}
        except StopIteration:
            return None

def onnx_static_quantization(onnx_model_path, output_path, calibration_dataloader):
    """ONNX静态量化"""
    calibration_data_reader = CalibrationDataReader(calibration_dataloader)

    quantize_static(
        onnx_model_path,
        output_path,
        calibration_data_reader,
        quantization_mode=QuantType.QInt8,
        force_fusions=False
    )
    print(f"量化模型已保存到: {output_path}")

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile

def prepare_for_mobile(model, example_input):
    """准备移动端部署"""
    # 1. 转换为TorchScript
    model.eval()
    traced_model = torch.jit.trace(model, example_input)

    # 2. 量化
    quantized_model = torch.quantization.quantize_dynamic(
        traced_model,
        {nn.Linear, nn.LSTM, nn.GRU},
        dtype=torch.qint8
    )

    # 3. 优化
    optimized_model = optimize_for_mobile(quantized_model)

    return optimized_model

def save_for_mobile(model, output_path='model.ptl'):
    """保存移动端模型"""
    model._save_for_lite_interpreter(output_path)
    print(f"移动端模型已保存到: {output_path}")

import tensorflow as tf

def tflite_quantization(tf_model_path, output_path, representative_data):
    """TFLite量化"""
    converter = tf.lite.TFLiteConverter.from_saved_model(tf_model_path)

    # 设置优化
    converter.optimizations = [tf.lite.Optimize.DEFAULT]

    # 提供代表性数据（用于静态量化）
    def representative_dataset():
        for data in representative_data.take(100):
            yield [data]

    converter.representative_dataset = representative_dataset

    # 确保操作兼容INT8
    converter.target_spec.supported_types = [tf.int8]
    converter.inference_input_type = tf.int8
    converter.inference_output_type = tf.int8

    # 转换
    tflite_quant_model = converter.convert()

    # 保存
    with open(output_path, 'wb') as f:
        f.write(tflite_quant_model)

    print(f"TFLite量化模型已保存到: {output_path}")

def evaluate_quantization(fp32_model, int8_model, test_loader):
    """评估量化前后的精度和性能"""

    # 1. 精度评估
    fp32_accuracy = evaluate_model(fp32_model, test_loader)
    int8_accuracy = evaluate_model(int8_model, test_loader)
    accuracy_drop = fp32_accuracy - int8_accuracy

    print(f"FP32精度: {fp32_accuracy:.2%}")
    print(f"INT8精度: {int8_accuracy:.2%}")
    print(f"精度损失: {accuracy_drop:.2%}")

    # 2. 模型大小
    fp32_size = get_model_size(fp32_model)
    int8_size = get_model_size(int8_model)
    compression = fp32_size / int8_size

    print(f"FP32大小: {fp32_size:.2f} MB")
    print(f"INT8大小: {int8_size:.2f} MB")
    print(f"压缩比: {compression:.2f}x")

    # 3. 推理速度
    fp32_latency = measure_inference_time(fp32_model, test_loader)
    int8_latency = measure_inference_time(int8_model, test_loader)
    speedup = fp32_latency / int8_latency

    print(f"FP32延迟: {fp32_latency:.2f} ms")
    print(f"INT8延迟: {int8_latency:.2f} ms")
    print(f"加速比: {speedup:.2f}x")

    return {
        'accuracy_drop': accuracy_drop,
        'compression': compression,
        'speedup': speedup
    }

def evaluate_model(model, dataloader):
    """评估模型精度"""
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, targets in dataloader:
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    return correct / total

def measure_inference_time(model, dataloader, num_iterations=100):
    """测量推理延迟"""
    model.eval()
    inputs, _ = next(iter(dataloader))

    # 预热
    with torch.no_grad():
        _ = model(inputs)

    # 测量
    start_time = time.time()
    with torch.no_grad():
        for _ in range(num_iterations):
            _ = model(inputs)
    end_time = time.time()

    avg_latency = (end_time - start_time) / num_iterations * 1000  # ms
    return avg_latency

def layerwise_accuracy_analysis(fp32_model, int8_model, test_loader):
    """逐层精度分析"""
    layer_outputs = {}

    # 收集FP32中间输出
    def fp32_hook(name):
        def hook(module, input, output):
            layer_outputs[f'fp32_{name}'] = output.detach()
        return hook

    fp32_hooks = []
    for name, module in fp32_model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            hook = module.register_forward_hook(fp32_hook(name))
            fp32_hooks.append(hook)

    # 收集INT8中间输出
    def int8_hook(name):
        def hook(module, input, output):
            layer_outputs[f'int8_{name}'] = output.detach()
        return hook

    int8_hooks = []
    for name, module in int8_model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            hook = module.register_forward_hook(int8_hook(name))
            int8_hooks.append(hook)

    # 前向传播
    with torch.no_grad():
        for inputs, _ in test_loader:
            _ = fp32_model(inputs)
            _ = int8_model(inputs)
            break  # 只分析第一个batch

    # 移除hooks
    for hook in fp32_hooks + int8_hooks:
        hook.remove()

    # 分析每层的误差
    for name, module in fp32_model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            fp32_out = layer_outputs[f'fp32_{name}']
            int8_out = layer_outputs[f'int8_{name}']

            error = torch.norm(fp32_out - int8_out) / torch.norm(fp32_out)
            cosine_sim = nn.functional.cosine_similarity(
                fp32_out.flatten(),
                int8_out.flatten(),
                dim=0
            )

            print(f"{name}:")
            print(f"  相对误差: {error:.4%}")
            print(f"  余弦相似度: {cosine_sim:.4f}")

def mixed_precision_quantization(model, sensitivity_scores, accuracy_threshold=0.01):
    """
    基于敏感度的混合精度量化

    Args:
        sensitivity_scores: 各层的量化敏感度
        accuracy_threshold: 可接受的精度损失阈值
    """
    precision_config = {}

    for layer_name, sensitivity in sensitivity_scores.items():
        if sensitivity < 0.005:
            # 不敏感层：INT4
            precision_config[layer_name] = 4
        elif sensitivity < 0.02:
            # 中等敏感层：INT8
            precision_config[layer_name] = 8
        else:
            # 高敏感层：FP16
            precision_config[layer_name] = 16

    # 应用混合精度
    for name, module in model.named_modules():
        if name in precision_config:
            bits = precision_config[name]
            if bits == 4:
                apply_int4_quantization(module)
            elif bits == 8:
                apply_int8_quantization(module)
            elif bits == 16:
                apply_fp16_quantization(module)

    return model, precision_config

def non_uniform_quantization(weights, codebook_size=256):
    """
    非均匀量化：基于聚类的量化

    使用k-means聚类找到最优量化点
    """
    from sklearn.cluster import KMeans

    # 展平权重
    weight_flat = weights.reshape(-1, 1)

    # k-means聚类
    kmeans = KMeans(n_clusters=codebook_size, random_state=0).fit(weight_flat)

    # 量化权重
    weight_quantized = kmeans.cluster_centers_[kmeans.labels_].reshape(weights.shape)

    # 计算量化误差
    mse = np.mean((weights - weight_quantized) ** 2)

    return weight_quantized, kmeans, mse

def product_quantization(weights, num_subvectors=8, codebook_size=256):
    """
    乘积量化（Product Quantization）

    将向量分解为子向量，分别量化
    """
    d = weights.shape[1]
    subvector_dim = d // num_subvectors

    codebooks = []
    quantized_subvectors = []

    for i in range(num_subvectors):
        start_idx = i * subvector_dim
        end_idx = start_idx + subvector_dim if i < num_subvectors - 1 else d

        # 提取子向量
        subvectors = weights[:, start_idx:end_idx]

        # 聚类量化
        quantized, kmeans, _ = non_uniform_quantization(subvectors, codebook_size)
        quantized_subvectors.append(quantized)
        codebooks.append(kmeans)

    # 重构量化权重
    weight_quantized = np.concatenate(quantized_subvectors, axis=1)

    return weight_quantized, codebooks

class AdaptiveQuantization(nn.Module):
    """自适应量化"""

    def __init__(self, num_features, bits=8):
        super().__init__()
        self.bits = bits
        self.num_features = num_features

        # 学习量化参数预测网络
        self.param_predictor = nn.Sequential(
            nn.Linear(num_features, num_features // 4),
            nn.ReLU(),
            nn.Linear(num_features // 4, 2)  # 输出scale和zero_point
        )

    def forward(self, x):
        # 动态预测量化参数
        params = self.param_predictor(x.mean(dim=(2, 3)))  # 全局平均池化
        scale = torch.sigmoid(params[:, 0:1]) * 2.0  # scale ∈ [0, 2]
        zero_point = params[:, 1:2]

        # 应用量化
        x_quantized = fake_quantize(x, scale.unsqueeze(-1).unsqueeze(-1),
                                    zero_point.unsqueeze(-1).unsqueeze(-1),
                                    self.bits)

        return x_quantized