第15章 模型部署与优化¶

import torch
import torch.nn as nn

def prune_model(model, pruning_ratio=0.3):
    """剪枝模型"""
    parameters_to_prune = []

    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):  # isinstance检查类型
            parameters_to_prune.append((module, 'weight'))

    # 全局非结构化剪枝
    parameters_to_prune = tuple(parameters_to_prune)
    torch.nn.utils.prune.global_unstructured(
        parameters_to_prune,
        pruning_method=torch.nn.utils.prune.L1Unstructured,
        amount=pruning_ratio
    )

    # 移除剪枝掩码
    for module, name in parameters_to_prune:
        torch.nn.utils.prune.remove(module, name)

    return model

def structured_prune(model, pruning_ratio=0.3):
    """结构化剪枝（剪枝整个通道）"""
    for module in model.modules():
        if isinstance(module, nn.Conv2d):
            # 计算每个输出通道的重要性（对输入通道和空间维度求和）
            importance = module.weight.abs().sum(dim=(1, 2, 3))

            # 选择要剪枝的通道
            num_channels = module.out_channels
            num_pruned = int(num_channels * pruning_ratio)
            _, indices = torch.topk(importance, num_channels - num_pruned, largest=True)

            # 创建掩码
            mask = torch.zeros(num_channels, dtype=torch.bool)
            mask[indices] = True

            # 应用剪枝
            module.weight.data = module.weight.data[mask]
            if module.bias is not None:
                module.bias.data = module.bias.data[mask]

            module.out_channels = num_channels - num_pruned

    return model

import torch

# 动态量化
def dynamic_quantization(model):
    """动态量化"""
    quantized_model = torch.quantization.quantize_dynamic(
        model,
        {nn.Linear},  # 动态量化仅支持Linear/LSTM/GRU，不支持Conv2d
        dtype=torch.qint8
    )
    return quantized_model

# 静态量化
def static_quantization(model, dataloader):
    """静态量化"""
    # 准备量化
    model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
    model_prepared = torch.quantization.prepare(model)

    # 校准
    model_prepared.eval()
    with torch.no_grad():  # 禁用梯度计算，节省内存
        for images, _ in dataloader:
            model_prepared(images)

    # 转换
    quantized_model = torch.quantization.convert(model_prepared)

    return quantized_model

# 量化感知训练
def qat_training(model, dataloader, epochs=10):
    """量化感知训练"""
    # 准备QAT
    model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
    model_prepared = torch.quantization.prepare_qat(model, inplace=True)

    # 训练
    optimizer = torch.optim.SGD(model_prepared.parameters(), lr=0.01)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model_prepared.train()
        for images, labels in dataloader:
            outputs = model_prepared(images)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 更新参数

    # 转换
    quantized_model = torch.quantization.convert(model_prepared)

    return quantized_model

import torch.nn.functional as F

class DistillationLoss(nn.Module):  # 继承nn.Module定义网络层
    def __init__(self, temperature=4.0, alpha=0.7):
        super(DistillationLoss, self).__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.kl_loss = nn.KLDivLoss(reduction='batchmean')
        self.ce_loss = nn.CrossEntropyLoss()

    def forward(self, student_logits, teacher_logits, labels):
        # 软标签（教师模型）
        soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)  # F.xxx PyTorch函数式API

        # 学生模型预测
        student_soft = F.log_softmax(student_logits / self.temperature, dim=1)

        # 蒸馏损失
        distill_loss = self.kl_loss(student_soft, soft_targets) * (self.temperature ** 2)

        # 硬标签损失
        hard_loss = self.ce_loss(student_logits, labels)

        # 加权求和
        loss = self.alpha * distill_loss + (1 - self.alpha) * hard_loss

        return loss

# 训练
def distillation_training(teacher, student, dataloader, epochs=10):
    teacher.eval()
    student.train()

    optimizer = torch.optim.Adam(student.parameters(), lr=0.001)
    criterion = DistillationLoss(temperature=4.0, alpha=0.7)

    for epoch in range(epochs):
        total_loss = 0.0
        for images, labels in dataloader:
            # 教师模型预测（不计算梯度）
            with torch.no_grad():
                teacher_logits = teacher(images)

            # 学生模型预测
            student_logits = student(images)

            # 蒸馏损失
            loss = criterion(student_logits, teacher_logits, labels)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()  # 将单元素张量转为Python数值

        print(f'Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}')

    return student

import torch.onnx

def export_to_onnx(model, dummy_input, onnx_path='model.onnx'):
    """导出ONNX模型"""
    model.eval()  # eval()评估模式

    # 导出
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=17,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size'},
            'output': {0: 'batch_size'}
        }
    )

    # 验证
    import onnx
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)

    print(f"模型已导出到: {onnx_path}")
    return onnx_path

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

def build_tensorrt_engine(onnx_path, engine_path='model.trt'):
    """构建TensorRT引擎"""
    # 创建builder
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(TRT_LOGGER)

    # 创建网络
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)

    # 解析ONNX
    with open(onnx_path, 'rb') as model:  # with自动管理文件关闭
        parser.parse(model.read())

    # 配置builder
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
    config.set_flag(trt.BuilderFlag.FP16)  # 使用FP16

    # 构建序列化引擎（build_engine已弃用，使用build_serialized_network）
    serialized_engine = builder.build_serialized_network(network, config)

    # 保存引擎
    with open(engine_path, 'wb') as f:
        f.write(serialized_engine)

    # 反序列化为可执行引擎
    runtime = trt.Runtime(TRT_LOGGER)
    engine = runtime.deserialize_cuda_engine(serialized_engine)

    print(f"TensorRT引擎已保存到: {engine_path}")
    return engine

# 推理
def tensorrt_inference(engine, input_data):
    """TensorRT推理"""
    import numpy as np

    # 分配内存
    output_shape = engine.get_binding_shape(1)
    output_data = np.empty(output_shape, dtype=np.float32)

    d_input = cuda.mem_alloc(input_data.nbytes)
    d_output = cuda.mem_alloc(output_data.nbytes)

    # 创建流
    stream = cuda.Stream()

    # 创建执行上下文
    context = engine.create_execution_context()

    # 推理（pycuda.autoinit 已初始化 CUDA 上下文，无需手动 push/pop）
    cuda.memcpy_htod_async(d_input, input_data, stream)
    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(output_data, d_output, stream)
    stream.synchronize()

    return output_data

import time

def benchmark_model(model, input_data, num_runs=100):
    """基准测试"""
    model.eval()

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

    # 测试
    start_time = time.time()
    with torch.no_grad():
        for _ in range(num_runs):
            _ = model(input_data)
    end_time = time.time()

    avg_time = (end_time - start_time) / num_runs
    fps = 1.0 / avg_time

    print(f"平均推理时间: {avg_time*1000:.2f} ms")
    print(f"FPS: {fps:.2f}")

    return avg_time, fps

# 完整部署流程
def deploy_pipeline(model_path, onnx_path='model.onnx', trt_path='model.trt'):
    """部署流程"""
    # 1. 加载PyTorch模型
    model = torch.load(model_path, weights_only=False)  # 加载完整模型需 weights_only=False
    model.eval()

    # 2. 量化
    quantized_model = dynamic_quantization(model)

    # 3. 导出ONNX
    dummy_input = torch.randn(1, 3, 224, 224)
    export_to_onnx(quantized_model, dummy_input, onnx_path)

    # 4. 构建TensorRT引擎
    engine = build_tensorrt_engine(onnx_path, trt_path)

    # 5. 基准测试
    print("PyTorch模型:")
    benchmark_model(model, dummy_input)

    print("\n量化模型:")
    benchmark_model(quantized_model, dummy_input)

    print("\nTensorRT引擎:")
    # 使用TensorRT推理
    # ...

# 使用
deploy_pipeline('resnet18.pth')

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

class DynamicBatchProcessor:
    """动态Batch处理器"""

    def __init__(self, engine_path: str):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f:
            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()

    def infer(self, inputs: dict):
        """支持多输入的推理"""
        bindings = []
        for name in self.engine:
            idx = self.engine[name]
            if self.engine.binding_is_input(idx):
                shape = inputs[name].shape
                self.context.set_binding_shape(idx, shape)
                bindings.append(int(inputs[name].data_ptr()))
            else:
                output = np.empty(self.context.get_binding_shape(idx), dtype=np.float32)
                bindings.append(int(output.data_ptr()))

        self.context.execute_v2(bindings)
        return output

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from typing import List

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

    def __init__(self, calibration_images: List[str], batch_size: int = 32, input_shape: tuple = (3, 224, 224)):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.batch_size = batch_size
        self.input_shape = input_shape
        self.calibration_images = calibration_images
        self.current_index = 0

        # 分配GPU内存
        self.device_input = cuda.mem_alloc(self.batch_size * np.prod(input_shape) * 4)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.current_index >= len(self.calibration_images):
            return None

        # 准备batch数据
        batch_images = []
        for i in range(self.batch_size):
            if self.current_index >= len(self.calibration_images):
                break
            # 这里应该加载实际图像并预处理
            img = np.random.randn(*self.input_shape).astype(np.float32)  # 示例：实际应加载真实图像
            batch_images.append(img)
            self.current_index += 1

        if len(batch_images) == 0:
            return None

        # 填充到GPU
        batch_data = np.stack(batch_images, axis=0)
        cuda.memcpy_htod(self.device_input, batch_data)

        return [int(self.device_input)]

    def read_calibration_cache(self):
        return None  # 首次校准返回None

    def write_calibration_cache(self, cache):
        # 保存校准缓存
        with open('calibration.cache', 'wb') as f:
            f.write(cache)

def build_int8_engine(onnx_path, calibrator, engine_path='model_int8.trt'):
    """构建INT8 TensorRT引擎"""
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)

    with open(onnx_path, 'rb') as model:
        parser.parse(model.read())

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
    config.set_flag(trt.BuilderFlag.INT8)  # 启用INT8
    config.int8_calibrator = calibrator  # 设置校准器

    serialized_engine = builder.build_serialized_network(network, config)

    with open(engine_path, 'wb') as f:
        f.write(serialized_engine)

    print(f"INT8引擎已保存到: {engine_path}")
    return serialized_engine

import torch
import torch.onnx
import onnx
import onnxsim

def export_onnx_with_fallback(model, dummy_input, onnx_path, opset_version=17):
    """带错误处理的ONNX导出"""
    try:
        torch.onnx.export(
            model,
            dummy_input,
            onnx_path,
            opset_version=opset_version,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}},
            do_constant_folding=True,
            export_params=True,
        )
        print(f"✓ 导出成功: {onnx_path}")
    except Exception as e:
        print(f"✗ 导出失败: {e}")
        # 常见问题处理
        if "not implemented" in str(e):
            print("  → 建议检查算子兼容性，尝试降低opset_version或使用自定义算子")
        elif "shape" in str(e):
            print("  → 建议检查dynamic_axes设置或输入尺寸")
        return False
    return True

def simplify_onnx(onnx_path, simplified_path=None):
    """简化ONNX模型"""
    if simplified_path is None:
        simplified_path = onnx_path.replace('.onnx', '_sim.onnx')

    try:
        onnx_model = onnx.load(onnx_path)
        onnx.checker.check_model(onnx_model)

        # 使用onnx-simplifier
        simplified_model, check = onnxsim.simplify(onnx_model)
        if check:
            onnx.save(simplified_model, simplified_path)
            print(f"✓ 简化成功: {simplified_path}")
        else:
            print("✗ 简化验证失败")
    except Exception as e:
        print(f"✗ 简化失败: {e}")

# 常见ONNX转换问题速查表
"""
| 问题 | 原因 | 解决方案 |
|------|------|----------|
| 算子不支持 | PyTorch算子无ONNX对应 | 自定义算子/降低opset |
| 动态shape错误 | dynamic_axes配置不当 | 检查并修正dynamic_axes |
| 精度损失 | FP32→FP16转换 | 使用FP32或混合精度 |
| 模型过大 | 冗余算子未消除 | 使用onnx-simplifier |
| 输入尺寸固定 | 未设置动态维度 | 添加dynamic_axes |
"""

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

def analyze_quantization_accuracy(model_fp32, model_quant, test_loader, device='cpu'):
    """分析量化前后的精度差异"""

    def evaluate(model, loader):
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for images, labels in loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        return 100 * correct / total

    acc_fp32 = evaluate(model_fp32, test_loader)
    acc_quant = evaluate(model_quant, test_loader)

    print("=" * 50)
    print("量化精度分析报告")
    print("=" * 50)
    print(f"FP32模型精度: {acc_fp32:.2f}%")
    print(f"量化模型精度: {acc_quant:.2f}%")
    print(f"精度损失: {acc_fp32 - acc_quant:.2f}%")
    print(f"相对损失: {(acc_fp32 - acc_quant) / acc_fp32 * 100:.2f}%")
    print("=" * 50)

    return {
        'fp32_acc': acc_fp32,
        'quant_acc': acc_quant,
        'abs_loss': acc_fp32 - acc_quant,
        'rel_loss': (acc_fp32 - acc_quant) / acc_fp32 * 100
    }

def per_layer_analysis(model_fp32, model_quant, input_tensor):
    """逐层分析量化误差"""
    print("\n逐层量化误差分析:")
    print("-" * 60)
    print(f"{'层名称':<30} {'最大误差':<15} {'平均误差':<15}")
    print("-" * 60)

    # 注册钩子记录中间输出
    fp32_outputs = {}
    quant_outputs = {}

    def get_hook(name, storage):
        def hook(module, input, output):
            storage[name] = output.detach()
        return hook

    # 分析每层（简化示例）
    for (name_fp32, module_fp32), (name_quant, module_quant) in zip(
        model_fp32.named_modules(), model_quant.named_modules()
    ):
        if isinstance(module_fp32, (nn.Conv2d, nn.Linear)):
            # 计算权重差异
            if hasattr(module_quant, 'weight'):
                diff = (module_fp32.weight - module_quant.weight).abs()
                max_err = diff.max().item()
                mean_err = diff.mean().item()
                print(f"{name_fp32:<30} {max_err:<15.6f} {mean_err:<15.6f}")

"""
量化精度损失主要原因及缓解策略：

1. **激活值分布不均匀**
   - 问题：某些层激活值范围过大，导致量化分辨率不足
   - 解决：使用Per-Channel量化而非Per-Tensor量化

2. **权重异常值**
   - 问题：权重中存在离群值影响量化范围
   - 解决：使用对称量化或裁剪异常值

3. **敏感层未特殊处理**
   - 问题：某些层对量化敏感（如第一层、最后一层）
   - 解决：对这些层保持FP16/FP32精度

4. **校准数据不足**
   - 问题：校准数据集不能代表真实数据分布
   - 解决：使用更多样化的校准数据

5. **动态范围估计不准**
   - 问题：静态量化时激活范围估计偏差
   - 解决：使用量化感知训练（QAT）或混合精度
"""