一、基本概念
模型压缩是人工智能领域中的一个重要研究方向。顾名思义,模型压缩就是减少模型的大小和复杂度,让模型能够在资源有限的情况下部署和运行。在学术研究的过程中,为了提高建模效果,我们会构建一些很大很复杂的模型,但这样的模型往往较难落地。这是由于模型需要的内存过大,以及模型的推理速度过慢,从而导致了落地这些模型的成本不断升高。模型压缩为我们落地这些高性能模型提供了一个可行的路径。
二、常见的模型压缩方法
1、模型剪枝
模型剪枝(Pruning)可以有效地减少模型中的参数量,从而在模型性能不会下降太多的前提下有效降低模型的计算量和存储需求。模型剪枝实际上就是将一些对模型计算影响不大的参数置零。
(1)非结构化剪枝(Unstructured Pruning)
非结构化剪枝一般是随机地移除模型权重,从而使得权重矩阵变得稀疏,故又称稀疏化剪枝。如果我们使用pytorch搭建了一个模型,那么可以很方便地使用pytorch提供的剪枝工具:
import torch
import torch.nn.utils.prune as prune
# 假设model是你要剪枝的模型
model = torch.nn.Linear(10, 5) # 示例模型
# 非结构化剪枝,移除30%的权重
prune.l1_unstructured(model, name='weight', amount=0.3)
(2)结构化剪枝(Structured Pruning)
结构化剪枝主要是移除一整个神经元或者模型层来减少模型的计算复杂度,它在剪枝的过程中可以有效地保证未被剪枝的层结构的完整性。同样,pytorch也提供了相应的工具:
# 结构化剪枝,移除50%的卷积核
prune.ln_structured(model, name='weight', amount=0.5, n=2, dim=0)
2、模型量化
模型量化主要通过将模型中的浮点数参数(如权重、激活值)转换为低精度的表示来实现模型压缩,通常是整数形式。
(1)动态量化
动态量化在模型推理的过程中动态地将权重从浮点数量化为整数。
import torch
import torch.quantization
# 假设model是你要量化的模型
model = torch.nn.Linear(10, 5) # 示例模型
# 动态量化
model_int8 = torch.quantization.quantize_dynamic(
model, {torch.nn.Linear}, dtype=torch.qint8)
(2)静态量化
静态量化在模型推理之前预先将权重从浮点数量化为整数。
import torch
import torch.quantization
# 静态量化
model_fp32 = torch.nn.Linear(10, 5) # 示例模型
model_fp32.eval() # 设置为评估模式
# 设置量化配置
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model_fp32_prepared = torch.quantization.prepare(model_fp32)
# 假设我们的模型是用于处理10维输入和5维输出的线性模型
input_size = 10
output_size = 5
# 生成随机的校准数据
# 一般来说,校准数据集的大小是整个数据集的一个小比例,例如1%
calibration_data_size = 100 # 你可以根据实际情况调整这个大小
calibration_data = torch.randn(calibration_data_size, input_size)
# 这里我们假设目标输出是5维的
target_output_size = output_size
calibration_targets = torch.randn(calibration_data_size, target_output_size)
# 将数据和目标放入DataLoader中,以便批量处理
calibration_loader = torch.utils.data.DataLoader(
dataset=list(zip(calibration_data, calibration_targets)),
batch_size=16, # 你可以根据实际情况调整批量大小
shuffle=False
)
# 使用校准数据进行校准
with torch.no_grad():
for batch in calibration_loader:
inputs, targets = batch
inputs, targets = inputs, targets
model_fp32_prepared(inputs)
# 转换模型
model_int8 = torch.quantization.convert(model_fp32_prepared)
3、知识蒸馏
知识蒸馏由Hinton老爷子在2015年提出(每天都在感叹老爷子的伟大!),核心思想是使用一个较大的教师模型的输出概率分布来引导一个较小规模的学生模型的训练。
(1)工作流程
- 训练教师模型:首先,训练一个性能达标的大型模型
- 生成软标签:使用教师模型对训练数据生成预测概率(软标签);
- 训练学生模型:用教师模型的软标签以及训练数据的真实标签,训练一个较小的学生模型,使其学习教师模型的知识;
(2)知识蒸馏的类型
- 离线蒸馏:教师模型和学生模型的训练是分开的,教师模型的输出用于指导学生模型的训练;
- 在线蒸馏:教师模型和学生模型同时更新,整个知识蒸馏框架是端到端可训练的;
- 自蒸馏:教师和学生模型使用相同的网络,这可以看作是在线蒸馏的一个特例;
(3)常见损失函数构造
知识蒸馏的损失函数通常由两部分组成:学生模型自身的损失(通常是交叉熵损失)和学生模型输出与教师模型输出之间的相似性损失(通常是KL散度损失)。
首先,我们计算学生模型自身的交叉熵损失:
其中,C是类别数量,是真实标签的one-hot表示。然后,我们来计算学生模型的输出和教师模型的输出之间的KL散度:
其中,S是学生模型,T是教师模型,和
分别是教师模型和学生模型对输入x的输出概率分布。而教师模型的输出就是软目标。
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 数据加载
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 这里我们用经典的图像数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 教师模型和学生模型使用与训练好的resnet
teacher_model = models.resnet34(pretrained=True).to(device)
student_model = models.resnet18().to(device) # 使用resnet18作为学生模型
# 教师模型评估模式
teacher_model.eval()
# 知识蒸馏损失函数
def distillation_loss(student_logits, teacher_logits, temperature, alpha):
soft_student_logits = nn.functional.softmax(student_logits / temperature, dim=1)
soft_teacher_logits = nn.functional.softmax(teacher_logits / temperature, dim=1)
KL_loss = nn.functional.kl_div(soft_student_logits, soft_teacher_logits, reduction='batchmean')
hard_loss = nn.functional.cross_entropy(student_logits, torch.max(soft_teacher_logits, 1)[1])
return (alpha * hard_loss) + ((1 - alpha) * KL_loss * temperature**2)
# 优化器
optimizer = optim.Adam(student_model.parameters(), lr=1e-4)
# 知识蒸馏训练
for epoch in range(10):
for i, (inputs, labels) in enumerate(train_loader):
inputs, labels = inputs.to(device), labels.to(device)
# 教师模型和学生模型前向传播
with torch.no_grad():
teacher_logits = teacher_model(inputs)
student_logits = student_model(inputs)
# 计算损失
loss = distillation_loss(student_logits, teacher_logits, temperature=10, alpha=0.5)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 100 == 0:
print(f'Epoch {epoch}, Step {i}, Loss: {loss.item()}')
三、总结
模型压缩在很多模型的落地过程中发挥了重要作用,但是它却并不是所有建模应用任务的必要步骤。例如在进行在线模型的边端芯片部署过程中,如果一个复杂的神经网络和一个简单的决策树性能差不多,那么直接边端部署决策树即可,又何必大费周章去压缩神经网络呢。因此,对于规模特别大的模型,我们才会考虑到上面的各种模型压缩方法。