第十七章：模型压缩与加速

目标：量化/剪枝/蒸馏流程，torch.compile 基础，张量核与算子融合直觉。

一个在 NVIDIA A100 上训练出的百亿参数大模型，固然性能强大，但它无法直接运行在手机、嵌入式设备甚至普通的云服务器上。高昂的推理成本、巨大的内存占用和不可接受的延迟，是阻碍许多深度学习模型从“实验室”走向“生产线”的“三座大山”。模型压缩与加速技术，就是为了移开这三座大山而生。本章将系统性地介绍三大经典的压缩技术——量化、剪枝、蒸馏，并带你初步探索 PyTorch 2.0 时代最激动人心的性能优化工具——torch.compile。我们将为你建立起关于底层硬件（如 Tensor Cores）和算子融合的直觉，让你明白模型加速不仅仅是算法层面的游戏，更是软件与硬件协同优化的艺术。

知识卡片：核心概念 & API

17.1 模型量化：从浮点到整数的“降维打击”

量化是目前工业界应用最广泛、效果最显著的加速技术之一。

为什么量化有效？

1. 模型尺寸减小: 一个 int8 权重占用的空间只有 float32 的 1/4。一个 4GB 的模型量化后可能只有 1GB。
2. 内存带宽降低: 推理时，从内存中读取 int8 数据比读取 float32 数据快得多，这在内存带宽受限的设备上尤为重要。
3. 计算加速: 现代 CPU 和 GPU（尤其是 NVIDIA 的 Tensor Cores）对 int8 计算有专门的硬件指令集优化，其计算吞吐量远高于 float32。

核心原理:
量化的本质是一种映射。对于一个 float32 张量 x，我们希望找到一个缩放因子 S 和一个零点 Z，使得 x ≈ S * (q - Z)，其中 q 是量化后的 int8 张量。

• S (Scale): 浮点数范围到整数范围的缩放比例。
• Z (Zero-point): 浮点数中的 0 对应到整数范围中的哪个值。

17.1.1 动态量化 (Dynamic Quantization)

• 特点: 最简单，无需校准数据，开箱即用。
• 做法:
  • 权重: 离线时，将模型的权重（主要是 nn.Linear, nn.LSTM 等）从 float32 转换为 int8。
  • 激活值: 在推理时，当数据流经被量化的层时，动态地将浮点激活值量化为 int8，进行 int8 计算，然后再将结果反量化回 float32。
• 适用场景: 模型中权重占计算密集型操作的大头，而激活值的内存传输不是瓶颈。非常适合 NLP 中的 Transformer 和 RNN 模型。

最小可运行示例：动态量化一个 LSTM 模型

import torch
import torch.nn as nn

# 1. 创建一个浮点模型
model_fp32 = nn.LSTM(input_size=20, hidden_size=10, num_layers=1)

# 2. 应用动态量化
# dtype=torch.qint8 指定了量化类型
model_quantized_dynamic = torch.quantization.quantize_dynamic(
    model_fp32, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)

# 3. 比较模型大小和性能
def print_model_size(model, label):
    torch.save(model.state_dict(), "temp.p")
    size = os.path.getsize("temp.p")/1e6
    print(f"Size of {label}: {size:.2f} MB")
    os.remove("temp.p")

print_model_size(model_fp32, "FP32 Model")
print_model_size(model_quantized_dynamic, "Dynamic Quantized Model")
# 你会看到模型大小显著减小

# 性能对比 (需要运行足够多次才能看到差异)
# dummy_input = torch.randn(5, 3, 20) # (seq_len, batch, input_size)
# %timeit model_fp32(dummy_input)
# %timeit model_quantized_dynamic(dummy_input)```

#### **17.1.2 静态量化 (Post-Training Quantization, PTQ)**

*   **特点**: 性能通常优于动态量化，因为激活值也被静态量化了，避免了运行时的动态计算开销。
*   **做法**:
    1.  **准备 (Prepare)**: 在浮点模型中插入“观察者”(`Observer`) 模块，用于收集激活值的统计信息（最大值、最小值等）。
    2.  **校准 (Calibrate)**: 将一小部分有代表性的、未经训练的样本数据（校准集）喂给插入了观察者的模型。观察者会记录下激活值的分布范围。
    3.  **转换 (Convert)**: 根据观察者收集到的信息，计算出最佳的 `S` 和 `Z`，然后将模型（包括权重和激活值的计算路径）正式转换为量化版本。
*   **适用场景**: 卷积神经网络（CNN）等激活值内存带宽成为瓶颈的模型。

#### **17.1.3 量化感知训练 (Quantization-Aware Training, QAT)**

*   **特点**: 精度最高的量化方法，但需要重新训练或微调。
*   **做法**: 在训练循环中，在前向和反向传播时**模拟**量化过程。具体来说，权重和激活值在计算后会被“伪量化”（即 `quantize -> dequantize`），这个过程会引入量化噪声。模型在训练时会学会对这种噪声的鲁棒性，从而在最终转换为真正的量化模型时，精度损失最小。
*   **API**: `torch.quantization.prepare_qat` -> `training_loop` -> `torch.quantization.convert`。

### **17.2 模型剪枝：为网络“瘦身”**

剪枝的核心思想是，大型神经网络中的许多参数实际上是冗余的，对最终性能贡献很小。移除这些参数可以减小模型尺寸并可能加速推理。

*   **非结构化剪枝 (Unstructured Pruning)**: 移除单个的、独立的权重（通常是将它们的数值设为 0）。这会产生一个**稀疏权重矩阵**。
    *   **优点**: 灵活，精度损失通常较小。
    *   **缺点**: 产生的稀疏矩阵在通用硬件（CPU/GPU）上难以获得实际的加速，因为它破坏了规整的矩阵结构。需要专门的稀疏计算库或硬件支持。
*   **结构化剪枝 (Structured Pruning)**: 移除整个的结构单元，如整个卷积核、通道或甚至是层。
    *   **优点**: 剪枝后的模型仍然是规整的、稠密的，可以直接在现有硬件上获得加速。
    *   **缺点**: 更粗粒度，可能会导致较大的精度下降。

**`torch.nn.utils.prune` 模块**:
PyTorch 提供了 `prune` 模块来实现剪枝。它通过一种称为“前向重参数化”的技巧工作：它不会直接修改权重张量，而是在模块上添加一个名为 `weight_mask` 的缓冲区和一个 `prune` 的前向钩子。在每次前向传播时，钩子会将 `weight` 和 `weight_mask` 相乘，得到有效的剪枝后权重。

**一个典型的迭代剪枝流程**:
1.  **训练**: 正常训练一个稠密模型。
2.  **剪枝**: 根据某种重要性标准（如权重的大小 `L1Norm`），剪掉 `p%` 的权重。
3.  **微调**: 继续在剪枝后的模型上进行几个 epoch 的微调，以恢复因剪枝造成的精度损失。
4.  **重复**: 重复步骤 2 和 3，逐步提高稀疏度。

### **17.3 知识蒸馏：让“学霸”带“学渣”**

知识蒸馏是一种模型压缩技术，但它不直接修改模型本身，而是改变训练方式。

*   **教师模型 (Teacher Model)**: 一个已经训练好的、性能强大但结构复杂的模型。
*   **学生模型 (Student Model)**: 一个参数量更少、结构更简单的模型，这是我们最终想要部署的模型。

**训练过程**:
1.  **硬标签损失 (Hard Label Loss)**: 学生模型像往常一样，计算其预测与真实标签（ground truth）之间的损失（如交叉熵）。
2.  **软标签损失 (Soft Label Loss)**:
    *   我们将同一个训练样本输入教师模型，得到其输出的 **logits**（未经 softmax 的原始输出）。这些 logits 包含了比硬标签（如 `[0, 1, 0, 0]`）更丰富的信息，例如，教师模型可能认为一张“猫”的图片，其 logits 可能是 `[10.2, 2.5, 0.1, ...]`，这表明它“非常确定是猫”，但“也有一点点像狗”。
    *   我们将教师和学生的 logits 都通过一个带有**温度 (Temperature, T)** 的 Softmax 函数进行平滑处理：`softmax(logits / T)`。较高的 `T` 会产生更平滑的概率分布（软标签），鼓励学生学习教师的“思考过程”。
    *   学生模型的软标签损失就是其平滑后的输出与教师的软标签之间的 KL 散度或 MSE。
3.  **总损失**: `Total_Loss = alpha * Hard_Loss + (1 - alpha) * Soft_Loss`。学生模型在两个损失的共同指导下进行训练。

通过这种方式，学生模型能够从教师那里学到“暗知识”，从而在相同的模型容量下，达到比单独训练更高的性能。

### **17.4 `torch.compile`：PyTorch 2.0 的“涡轮增压”**

`torch.compile` 是 PyTorch 2.0 引入的、可能是近年来 PyTorch 最重要的性能特性。它承诺在**只需一行代码** `optimized_model = torch.compile(model)` 的情况下，就能为你的模型带来显著的（通常是 30%-200%）性能提升。

**它是如何工作的？**
`torch.compile` 是一个 JIT (Just-In-Time) 编译器，它将你的 PyTorch 代码从解释执行的 Python，转换成一个更底层的、高度优化的表示。
1.  **图捕获 (Graph Capture)**: 它首先会“追踪”你的 `forward` 方法的执行，将动态的 Python 操作捕获成一个静态的计算图。
2.  **图降低 (Graph Lowering)**: 将高层的 PyTorch 算子（如 `nn.Conv2d`）分解成更底层的计算原语。
3.  **图编译 (Graph Compilation)**: 这是魔法发生的地方。它使用不同的后端（如 TorchInductor）来对计算图进行深度优化：
    *   **算子融合 (Operator Fusion)**: 这是最重要的优化之一。考虑 `y = x + bias; z = nn.ReLU(y)` 这样一个序列。在标准的 PyTorch 中，这需要两次独立的 GPU 内核调用和一次 GPU 全局内存的读写（将 `y` 写回，再读出）。`torch.compile` 可以将这两个操作**融合**成一个单一的 GPU 内核，这个内核直接在 GPU 的寄存器或高速缓存中完成 `add` 和 `ReLU`，无需与全局内存进行交互。这极大地减少了内存带宽瓶颈和内核启动开销。
    *   **硬件特定的代码生成**: 编译器可以根据你的具体硬件（如 NVIDIA Ampere 架构）生成最优化的代码，例如，自动调用 Tensor Cores 来加速 `float16` 矩阵乘法。

**直觉：张量核 (Tensor Cores) 与算子融合**
*   **张量核**: 现代 NVIDIA GPU 上的专用硬件单元，可以极其高效地执行小型的 `4x4` 矩阵乘法-累加操作。当你的计算（如卷积或全连接层）能够被分解成这种形式，并且数据类型是 `float16` 或 `int8` 时，就能获得数倍的性能提升。
*   **算子融合的价值**: GPU 的计算速度远快于其内存访问速度。大多数模型的性能瓶颈在于**内存带宽**，即数据在 GPU 全局内存和计算单元之间的来回搬运。算子融合通过减少这种“来回跑”的次数，让数据尽可能地停留在高速缓存中被连续处理，从而最大化计算单元的利用率。

**最小可运行示例：体验 `torch.compile` 的加速**

```python
import torch
import torch.nn as nn
import time

# 定义一个简单的模型
def create_model():
    return nn.Sequential(
        nn.Conv2d(3, 32, 3), nn.ReLU(),
        nn.Conv2d(32, 64, 3), nn.ReLU(),
        nn.Conv2d(64, 128, 3), nn.ReLU(),
        nn.Flatten(),
        nn.Linear(128 * 58 * 58, 10)
    ).cuda()

model_eager = create_model()

# --- 【核心】一行代码应用编译 ---
# mode="reduce-overhead": 适用于小输入尺寸，减少框架开销
# mode="max-autotune": 适用于大输入，花费更多时间编译以寻找最佳内核
model_compiled = torch.compile(create_model(), mode="max-autotune")

dummy_input = torch.randn(16, 3, 64, 64).cuda()

# 预热，确保 CUDA 内核被加载
for _ in range(5):
    model_eager(dummy_input)
    model_compiled(dummy_input)

# 性能测试
torch.cuda.synchronize()
start = time.time()
for _ in range(50):
    model_eager(dummy_input)
torch.cuda.synchronize()
eager_time = time.time() - start

torch.cuda.synchronize()
start = time.time()
for _ in range(50):
    model_compiled(dummy_input)
torch.cuda.synchronize()
compiled_time = time.time() - start

print(f"Eager mode time: {eager_time:.4f} seconds")
print(f"Compiled mode time: {compiled_time:.4f} seconds")
print(f"Speedup: {eager_time / compiled_time:.2f}x")

在支持的硬件上，你会看到 torch.compile 带来的显著加速。

练习题

1. 量化精度对比: 使用 torchvision 加载一个预训练的 ResNet-18。分别使用动态量化和静态量化（需要你创建一个小的校准数据集）对其进行量化。比较原始模型、动态量化模型和静态量化模型在验证集上的准确率，并记录下它们的模型大小。
2. 剪枝实践: 使用 torch.nn.utils.prune.l1_unstructured，对一个训练好的 MLP 模型进行 50% 的非结构化剪枝。剪枝后，检查模型的一个权重矩阵，确认其中约一半的值变为了零。然后，对剪枝后的模型进行微调，观察其精度能恢复到什么水平。
3. 蒸馏损失实现: 编写一个名为 DistillationLoss 的 nn.Module，它的 __init__ 接收 teacher_model, alpha 和 temperature。它的 forward 方法接收 student_logits, labels 和 inputs，并在内部完成计算教师 logits、软化 logits、计算硬标签损失和软标签损失，并最终返回加权的总损失。
4. torch.compile 的动态性限制: torch.compile 在图捕获时，对 Python 的动态控制流（如依赖于张量值的 if 语句）有限制。创建一个模型，其 forward 方法中包含一个 if x.sum() > 0: 的分支。尝试用 torch.compile 编译它，观察会发生什么。（提示：可能会触发 “Graph Break”）。

延伸阅读

1. PyTorch 官方文档与教程:
   • Quantization Recipes: 官方的量化教程，涵盖了所有量化模式的详细示例。
   • Pruning Tutorial: 官方的剪枝教程。
   • PyTorch 2.0 Get Started: torch.compile 的官方介绍和性能基准。
2. 深度学习模型压缩综述:
   • 搜索 “A Survey of Model Compression and Acceleration for Deep Neural Networks” 等关键词，可以找到许多对该领域技术进行系统性梳理的优秀综述论文。
3. 博客和文章:
   • “The Hardware Lottery” (by Sara Hooker): 一篇引人深思的文章，讨论了硬件的发展如何影响了深度学习算法的演进，这与模型加速的思想息息相关。
   • Hugging Face 的 optimum 库: 一个专门用于对 transformers 模型进行压缩和加速的库，集成了 ONNX, OpenVINO, Quantization 等多种技术。