图像分类模型量化实操

---

前言

最近接到了一个需要将之前训好的模型轻量化的工作，正好实践一下量化和剪枝，持续更新中✌︎( ᐛ )✌︎

---

一、动态量化和非动态量化

在深度学习领域，量化是一种模型优化技术，它通过减少模型中权重和激活的精度来减小模型大小和加速推理速度，同时尽量保持模型性能。量化主要分为两种类型：动态量化和非动态量化。

1. 动态量化（Dynamic Quantization）：

   • 动态量化是一种在模型训练过程中进行的量化方法。
   • 在这种方法中，模型的权重和激活在训练过程中被量化到较低的位宽（如8位或16位），而不是使用传统的32位浮点数。
   • 量化操作是在模型的前向传播过程中动态进行的，即在模型运行时，权重和激活值在计算之前被量化。
   • 动态量化通常可以保持较高的模型精度，因为它允许模型在训练过程中学习量化的表示。
   • 动态量化的模型可以在训练完成后直接部署，无需额外的量化校准步骤。

2. 非动态量化（Static Quantization）：

   • 非动态量化，也称为静态量化，是在模型训练完成后进行的。
   • 在这种量化方法中，模型的权重和激活值被一次性量化到较低的位宽，通常是在模型评估阶段。
   • 静态量化需要一个校准过程，这个过程通常涉及在一组代表性数据上运行模型，以确定量化参数（如量化范围和零点）。
   • 静态量化的模型在推理时速度更快，因为量化操作是离线完成的，不需要在推理时动态进行。
   • 但是，静态量化可能会牺牲一些模型精度，因为量化是在模型训练完成后进行的，模型没有机会在训练过程中适应量化误差。

两种量化方法各有优势和适用场景。动态量化更适合需要保持高精度的应用，而非动态量化则更适合对推理速度要求较高的场景。在实际应用中，选择哪种量化方法取决于具体的需求和资源限制。

二、具体步骤

1.动态量化（GPU版）

在PyTorch中，动态量化通常涉及到将模型的权重和激活值量化到较低的精度，比如16位浮点数（FP16）。动态量化可以在训练过程中进行，也可以在模型推理时进行。对于FP16量化，PyTorch提供了torch.cuda.amp模块，它支持自动混合精度（Automatic Mixed Precision, AMP）训练。

以下是如何在模型中加入动态量化的步骤：

1. 启用自动混合精度：使用torch.cuda.amp中的autocast和GradScaler。

2. 修改训练循环：在训练循环中使用autocast来自动选择使用FP16或FP32进行计算。

3. 保存量化模型：在训练完成后，保存量化后的模型。

下面是代码中加入动态量化的示例：

import torch
from torch.cuda.amp import autocast, GradScaler
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Dataset
from torchvision.datasets import ImageFolder
from tqdm import tqdm
import os

# 移动模型到GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.8)

# 使用GradScaler来自动管理梯度缩放
scaler = GradScaler()

# 训练模型
def train_model(num_epochs):
    best_accuracy = 0  # 保存最佳准确率
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        running_corrects = 0
        for images, labels in tqdm(train_loader):
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            
            # 使用autocast进行FP16量化
            with autocast():
                outputs = model(images)
                loss = criterion(outputs, labels)
            
            scaler.scale(loss).backward()  # 缩放损失，以防止梯度下溢
            scaler.step(optimizer)  # 更新模型参数
            scaler.update()  # 准备下一个优化步骤

            # 统计损失和正确预测的数量
            running_loss += loss.item() * images.size(0)
            _, preds = torch.max(outputs, 1)
            running_corrects += torch.sum(preds == labels.data)

        # 更新学习率
        scheduler.step()

        epoch_loss = running_loss / len(train_loader.dataset)
        epoch_acc = running_corrects.double() / len(train_loader.dataset)

        # 每个epoch结束后打印损失和准确率
        print(f'Epoch {epoch + 1}/{num_epochs}')
        print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

        # 在测试集上评估模型
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for images, labels in tqdm(test_loader):
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        test_acc = correct / total
        print(f'Test Accuracy: {test_acc:.4f}')

        # 如果准确率提高，则保存模型
        if test_acc >= best_accuracy:
            best_accuracy = test_acc
            torch.save(model.state_dict(), 'best_model16.pth')
            print('best Model saved.')

    torch.save(model.state_dict(), 'last_model.pth')
    print('last Model saved.')

train_model(10)

在这个修改后的代码中，autocast上下文管理器用于在模型的前向传播中自动选择FP16或FP32进行计算，而GradScaler用于自动管理梯度缩放，以防止在反向传播过程中发生梯度下溢。这样，你的模型就可以在训练过程中进行动态量化，同时保持较高的训练效率和模型性能。

2.QAT量化（CPU版）

如果你的设备没有GPU，可以采用量化感知训练（QAT）来量化模型，我们需要对您的代码进行一些修改。以下是使用 QAT 进行模型量化的步骤：

1. 导入量化相关的库：
   导入 PyTorch 量化库中的相关模块。

2. 定义量化配置：
   为模型指定量化配置，这将告诉量化流程哪些层需要量化。

3. 准备模型进行 QAT：
   使用 torch.quantization.prepare_qat 函数准备模型进行量化感知训练。

4. 训练模型：
   在训练过程中，模型会模拟量化操作，以便在实际量化时保持更高的性能。

5. 转换模型为量化模型：
   使用 torch.quantization.convert 函数将模型转换为量化模型。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Dataset
from torchvision.datasets import ImageFolder
from tqdm import tqdm
import os
import torch.quantization

# 数据集路径
data_dir = r"D:\game\8_14\hok_pa\data1"

# 图像预处理
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize((17, 10)),
    transforms.ToTensor(),
])

# 创建数据集
dataset = ImageFolder(root=data_dir, transform=transform)

# 划分数据集为训练集和测试集
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

# 英雄头像分类网络
class MyNet(nn.Module):
    def __init__(self, num_classes):
        super(MyNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, stride=2),
            nn.ReLU(),
        )
        self.classifier = nn.Linear(16 * 6 * 5, num_classes)  # Adjusted for the output size

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

num_classes = len(os.listdir(data_dir))
model = MyNet(num_classes)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 移动模型到CPU
device = torch.device("cpu")
model.to(device)

# 定义量化配置
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 准备模型进行 QAT
torch.quantization.prepare_qat(model, inplace=True)

# 训练模型
def train_model(num_epochs):
    best_accuracy = 0
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        running_corrects = 0
        for images, labels in tqdm(train_loader):
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * images.size(0)
            _, preds = torch.max(outputs, 1)
            running_corrects += torch.sum(preds == labels.data)

        epoch_loss = running_loss / len(train_loader.dataset)
        epoch_acc = running_corrects.double() / len(train_loader.dataset)
        print(f'Epoch {epoch + 1}/{num_epochs}')
        print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

        # 在测试集上评估模型
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for images, labels in tqdm(test_loader):
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        test_acc = correct / total
        print(f'Test Accuracy: {test_acc:.4f}')

        if test_acc >= best_accuracy:
            best_accuracy = test_acc
            torch.save(model.state_dict(), 'best_model_qat.pth')
            print('best Model saved.')

    torch.save(model.state_dict(), 'last_model_qat.pth')
    print('last Model saved.')

train_model(20)

# 转换模型为量化模型
model.eval()
torch.quantization.convert(model, inplace=True)

3.测试结果

initial model

1. 1000次调用 single_kda 函数平均每次调用用时: 0.001194 秒

2. Quantized model size: 0.02 MB

3. 准确率：100%(0/3129)

Quantized model

1.1000次调用 single_kda 函数平均每次调用用时: 0.001005 秒

2.Quantized model size: 0.01 MB
3. 准确率：100%(0/3129)

---

总结

这次介绍了动态量化与静态量化，并实际演示了动态量化中GPU和CPU版。最后总结一下他们的特点：

动态量化的特点：

• 在推理时对权重和激活值进行量化。
• 可以减少模型大小和计算需求。
• 通常用于加速模型推理。

静态量化的特点：

• 在训练后对权重进行量化。
• 需要使用校准数据集来确定量化参数。
• 可以显著减少模型大小和提高推理速度。