22 模型轻量化：如何让模型运行在低配置设备上？

你好，我是独行。

前面我们从0～1构建大模型的那节课里，最后我通过5MB数据训练出的模型，占用了大概500M的存储空间，参数量约1.2亿，当时为了节省时间，只简单跑了一下，这比较极端，可以说大量的参数是浪费的，这里我简单举个例子说明一下，比如公式：

$$f(x)=k_{1}x_{1}+k_{2}x_{2}+k_{3}x_{3}+…+k_{1000}x_{1000}$$

经过少量数据训练后，我们只确定了 $k_{1}$、$k_{2}$…$k_{100}$，后面的参数都未经过训练，所以其实是没用的。这种情况下，这个公式就有了可优化的空间，比如把 $k_{100}$ 之后的全部砍掉，或者保留 $k_{300}$ 之前的，之后的全部砍掉。

放在模型中，我们知道虽然参数量和训练数据量有很大关系，但是参数量还和模型网络的设计息息相关，当模型设计的层数比较深，那么也是很有可能产生浪费的情况。这种情况下，就可以对模型进行优化，优化的目标是降低模型的复杂度，包括参数量、占用空间。模型复杂度降低了，那么就可以降低资源的消耗，模型就可以运行在更低配置的设备上。

接下来，我带你学习几种模型轻量化方法。

参数剪枝

参数剪枝是一种优化技术，通过减少模型中的参数数量来降低复杂性和运行成本。剪枝通常涉及识别并删除那些对模型性能贡献较小的参数，比如权重。一般来说，是移除某些连接或整个神经元。下面我来详细介绍下参数剪枝步骤。

加载已经训练好的模型，一般包含模型架构和权重。
选择剪枝策略，决定是进行无结构剪枝还是结构剪枝。无结构剪枝就是移除单个权重，它通常对硬件支持较差；而结构剪枝就是移除整个过滤器或通道，它可以更好地利用现代硬件优化。
定义剪枝标准，根据权重的重要性来决定哪些权重被剪枝。常见的标准包括权重的绝对值大小，小的权重可能对输出影响较小，可以被剪枝。
实施剪枝，应用剪枝标准，修改模型的权重。这可以通过修改权重矩阵，将选定的权重设为零来实现。

我们选择无结构剪枝，剪枝之前，先执行以下代码打印一下参数。

for name, param in model.named_parameters():
    if param.requires_grad:
        print(name, param.numel())

positional_encoding 512
embed.weight 50499072
transformer_decoder.layers.0.self_attn.in_proj_weight 786432
transformer_decoder.layers.0.self_attn.in_proj_bias 1536
transformer_decoder.layers.0.self_attn.out_proj.weight 262144
transformer_decoder.layers.0.self_attn.out_proj.bias 512
transformer_decoder.layers.0.multihead_attn.in_proj_weight 786432
transformer_decoder.layers.0.multihead_attn.in_proj_bias 1536
transformer_decoder.layers.0.multihead_attn.out_proj.weight 262144
....
....
....
transformer_decoder.layers.5.norm1.bias 512
transformer_decoder.layers.5.norm2.weight 512
transformer_decoder.layers.5.norm2.bias 512
transformer_decoder.layers.5.norm3.weight 512
transformer_decoder.layers.5.norm3.bias 512
fc.weight 50499072
fc.bias 98631

执行裁剪逻辑：

model = load_model('models/transformer_decoder_model.pth')
linear_layer = model.fc
# 应用无结构剪枝，移除50%的权重
prune.l1_unstructured(linear_layer, name='weight', amount=0.5)
torch.save(model.state_dict(), 'models/transformer_decoder_model_cut.pth')

裁剪之后再打印下参数：

positional_encoding 512
embed.weight 50499072
transformer_decoder.layers.0.self_attn.in_proj_weight 786432
transformer_decoder.layers.0.self_attn.in_proj_bias 1536
transformer_decoder.layers.0.self_attn.out_proj.weight 262144
transformer_decoder.layers.0.self_attn.out_proj.bias 512
...
...
...
transformer_decoder.layers.5.norm1.bias 512
transformer_decoder.layers.5.norm2.weight 512
transformer_decoder.layers.5.norm2.bias 512
transformer_decoder.layers.5.norm3.weight 512
transformer_decoder.layers.5.norm3.bias 512
fc.bias 98631
fc.weight_orig 50499072

fc.weight变成了fc.weight_orig，是因为在PyTorch中使用prune模块进行剪枝后，原始权重会被保留为weight_orig，而weight则变成了一个由原始权重和掩码运算得到的视图（view）。这是为了在训练过程中仍然可以访问和更新原始权重，同时应用掩码来模拟权重被剪枝的效果。

执行完剪枝后，发现模型并没有减少，甚至有可能会增大，那是因为在PyTorch中，剪枝实际上是通过在原有的模型参数上添加掩码来实现的，而不是真正地移除那些参数。这意味着剪枝操作增加了额外的状态信息，比如掩码，从而可能导致模型文件总大小增加。

为了确保剪枝效果永久化并减少模型的大小，你需要在保存模型前将掩码应用到原始权重并移除剪枝带来的额外属性。使用prune.remove函数来把weight_orig替换回weight，并移除掩码。这一步确保了模型权重不再依赖于掩码，然后保存简化模型即可。

prune.remove(model.fc, 'weight')
torch.save(model.state_dict(), 'models/transformer_decoder_model_cut.pth')

剪枝完成，模型的性能可能会有所下降，需要通过再次训练或者微调模型，恢复一部分性能损失，最后，再次评估模型性能，确保它仍然满足应用要求。

量化

量化是把模型中的浮点类型的权重转换为低精度（如INT8或INT16）的过程，不仅可以减少模型的存储空间需求，还能在某些硬件上加速模型的推理速度。这就好比我们开发语言里的精度转换，比如double/long类型转int类型，一般会损失精度，至于影响大不大取决于这个变量本身存储的数值。模型量化也是一样，前面我们讲过模型默认以FP32的精度存储数据，也就是一个参数需要32位4个字节存储，量化成int8，那么只需要8位1个字节来存储。存储量直接减少75%。

你还记不记得我们在学习6B的时候，可以以int4的精度运行，大概只需要6GB的显存。

model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True, ).quantize(4).cuda()

量化一般有三种情况：静态量化、动态量化、感知训练量化，上面的这种应用的是自定义量化技术，有点像静态量化，又不完全是，你可以看一下quantize(4)方法的源码。

def quantize(model, weight_bit_width):
    """Replace fp16 linear with quantized linear"""
    for layer in model.layers:
        layer.attention.query_key_value = QuantizedLinear(
            weight_bit_width=weight_bit_width,
            weight_tensor=layer.attention.query_key_value.weight.to(torch.cuda.current_device()),
            bias_tensor=layer.attention.query_key_value.bias,
            in_features=layer.attention.query_key_value.in_features,
            out_features=layer.attention.query_key_value.out_features,
            bias=True,
            dtype=torch.half,
            device=layer.attention.query_key_value.weight.device,
        )
        layer.attention.dense = QuantizedLinear(
            weight_bit_width=weight_bit_width,
            weight_tensor=layer.attention.dense.weight.to(torch.cuda.current_device()),
            bias_tensor=layer.attention.dense.bias,
            in_features=layer.attention.dense.in_features,
            out_features=layer.attention.dense.out_features,
            bias=True,
            dtype=torch.half,
            device=layer.attention.dense.weight.device,
        )

这个其实是把模型原来的层替换成新的自定义的QuantizedLinear层了，感兴趣的话，你可以继续阅读QuantizedLinear类的源码了解细节。

静态量化

静态量化是在模型的推理阶段将权重和激活数据从浮点数转换为低精度整数的过程。和动态量化不同，静态量化会在推理前对模型的权重和激活数据进行量化，通常需要一个校准步骤，通过这一步可以决定最佳的量化参数。你可以参考我给出的示例代码。

import torch
import torch.nn as nn
import torch.quantization

model = TransformerDecoderModel(vocab_size=10000, embed_size=512, num_heads=8, hidden_dim=2048, num_layers=6)
model.eval()
model_fp32_prepared = torch.quantization.prepare(model)

# 模型校准过程
# 运行一些输入数据来校准模型
for input_batch in calibration_dataset:
    model_fp32_prepared(input_batch)

# 转换模型为量化版本
model_int8 = torch.quantization.convert(model_fp32_prepared)

# 保存量化模型
torch.save(model_int8.state_dict(), 'quantized_transformer_model.pth')

因为模型的权重和激活函数都被量化了，所以模型的整体性能可能会有影响，所以在量化完后需要进行全面的测试。

动态量化

动态量化主要针对模型的权重进行量化，每次输入时都会重新计算量化参数。这种方法通常用于模型的线性层（全连接层）和循环层（如LSTM）。动态量化可以在不需要重新训练模型的情况下实现，适合那些对推理速度有较高要求的应用场景。你可以看一个简单示例。

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

model = TransformerDecoderModel(vocab_size=10000, embed_size=512, num_heads=8, hidden_dim=2048, num_layers=6)
model.eval()  # 设置为评估模式

# 这里指定了量化nn.Linear和nn.LSTM层，使用8-bit的量化整数（torch.qint8）
quantized_model = quantize_dynamic(
    model, 
    {nn.Linear},       # 指定量化的层类型
    dtype=torch.qint8  # 指定量化的数据类型
)

# 保存量化模型
torch.save(quantized_model.state_dict(), 'quantized_transformer_model.pth')

感知训练量化

感知训练量化是指在训练过程中应用量化效果，使模型适应量化带来的精度损失。这种方法通常能保持或仅轻微影响模型的精度。在之前从0～1手敲Transformer的例子中，我们简单修改下模型定义的代码。

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

class TransformerDecoderModel(nn.Module):
    def __init__(self, vocab_size, embed_size, num_heads, hidden_dim, num_layers):
        ...
        self.quant = QuantStub() # 在模型的前端进行激活的量化和反量化。
        self.dequant = DeQuantStub() # 在模型的后端进行激活的量化和反量化。
    ...

# 初始化模型
model = TransformerDecoderModel(vocab_size=10000, embed_size=512, num_heads=8, hidden_dim=2048, num_layers=6)
model.eval()  # 评估模式

# 启用QAT
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model = prepare_qat(model)
model.train()

model.eval()
model = convert(model)
# 保存量化模型
torch.save(model.state_dict(), 'quantized_transformer_model.pth')

不过这只是示例代码，实际会更加复杂一点。除了量化外，还有一种用得比较多的轻量化方法：知识蒸馏，接下来我们一起看下。

知识蒸馏

知识蒸馏是一种训练技术，其中一个小的模型，被叫做学生模型，学习模仿一个大的已经训练好的模型，这个模型通常叫做教师模型。通过这种方式，小模型可以在保持较高性能的同时，显著减少参数数量。我们来看下知识蒸馏的基本步骤。

训练教师模型：首先训练一个大型的高性能的教师模型。这个模型通常是尽可能地大和复杂，以达到高准确率。
定义学生模型：设计一个比教师模型小得多的学生模型。学生模型的架构不必与教师模型相同，但通常需要能够学习相似类型的特征。
蒸馏过程：使用教师模型的输出来训练学生模型。这通常涉及调整学生模型的损失函数，使其不仅要学习真实的标签，还要学习模仿教师模型的输出。
评估学生模型：测试学生模型的性能，确保它能够达到预期的精度，同时具有更小的模型大小或更快的推理速度。

假设我们有一个已经训练好的教师模型和一个较小的学生模型，应该如何使用PyTorch实现知识蒸馏呢？你可以看我给出的示例。

import torch
import torch.nn as nn
import torch.optim as optim

# 假设teacher_model和student_model已经定义并加载
teacher_model.eval()  # 确保教师模型在评估模式

# 定义学生模型
student_model = StudentModel()

# 损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(student_model.parameters(), lr=0.001)

# 数据加载器
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 蒸馏过程
for epoch in range(num_epochs):
    student_model.train()
    for data, target in train_loader:
        optimizer.zero_grad()

        # 教师模型的输出
        with torch.no_grad():
            teacher_output = teacher_model(data)

        # 学生模型的输出
        student_output = student_model(data)
        # 计算损失：可以根据需要混合教师输出和真实标签
        loss = criterion(student_output, teacher_output)  # 以教师模型的输出作为目标
        # 反向传播和优化
        loss.backward()
        optimizer.step()

# 评估学生模型性能
student_model.eval()
# 进行测试和评估...

如果你对前面学习的内容有印象的话，其实知识蒸馏是很好理解的，关键就在下面这一步：

loss = criterion(student_output, teacher_output)  # 以教师模型的输出作为目标

之前我们计算损失的时候，是以真实标签作为对比的，而这里是以教师模型的输出作为对比，为什么教师模型的输出比真实标签更适合知识蒸馏呢？

这是因为教师模型的输出通常是经过Softmax函数处理的概率分布，被叫做“软标签”。相对于硬标签，也就是真实标签（0或1的类别标签），软标签能够提供关于各个类别相对可能性的更多信息。比如，在分类任务中，一个图像可能有90%的概率是猫，10%的概率是狗，而不是简单地标记为“猫”。这种概率分布包含了数据的不确定性和类别之间相似性的信息，有助于学生模型更全面地理解数据。

常见的模型轻量化方法就是这些，当然还有一些其他方式，比如参数共享、低秩因子分解等技术也是可以用来减少模型大小的，有兴趣的话你可以自己了解一下，这里我就不再详细阐述了。

小结

这节课我向你介绍了模型轻量化的一些常用方法，其中参数剪枝是一种非常实用的模型轻量化技术，分为加载模型、选择剪枝策略、定义剪枝标准、实施剪枝、微调训练、评估模型几个步骤。此外我们还可以将模型量化或者进行知识蒸馏。

在实际开发过程中，这些方法都非常有用，大方向上模型可以制作成不同的尺寸，小方向上模型可以做不同精度的量化处理，整体目标是一致的，就是希望模型能完美运行在不同配置的设备上。

这部分内容做起来还是非常有成就感的，有点像我们软件开发过程中的接口性能优化、慢SQL优化，意义非常大，你可以花点时间实际练习一下。

思考题

刚刚我们讲的剪枝操作，就是去掉一些不需要的参数，你可以想一下，我们该如何评估某些参数是否有用？除了我说的计算绝对值的方法，还有别的方法吗？欢迎把你的想法打在评论区，我们一起讨论。如果你觉得这节课的内容对你有帮助的话，也欢迎你分享给其他朋友，我们下节课再见！

精选留言（2）

石云升 👍（3） 💬（0）
基于梯度的剪枝：计算每个参数对损失函数的梯度。梯度小的参数被认为不那么重要,可以被剪掉。这种方法考虑了参数对模型性能的实际影响。基于激活的剪枝：分析神经元的激活情况。很少被激活或激活值很小的神经元可能不那么重要。基于重构误差的剪枝：评估删除某个参数后对网络输出的影响。影响小的参数可以被剪掉。这种方法更准确但计算成本较高。基于敏感度的剪枝：分析每个参数对模型性能的敏感度。敏感度低的参数可以被剪掉。
2024-09-08

潇洒哥66 👍（0） 💬（0）
除了参数对于模型输出重要性的考虑，有其它标准作为剪枝的依据吗？类似于参数对于模型稳定性，甚至是泛化能力的影响。
2024-10-16