Pytorch模型量化（一）：快速上手

最新推荐文章于 2024-08-16 09:33:57 发布

Nn_echo

最新推荐文章于 2024-08-16 09:33:57 发布

阅读量774

点赞数

文章标签： pytorch 人工智能 python 机器学习 transformer

什么是模型量化（Quantization）？

模型量化（Quantization）是一种使用更低的数据位宽（bitwidths）进行计算和数据存储的技术，例如，可以将模型中的全部或者部分float32计算和权重转换为int8的计算或者权重。当然，这种转换通常是对已经训练好的模型进行的，换句话说，模型量化常运用在模型推理的过程中。int8所需的存储空间仅为float32的四分之一，并且硬件进行int8计算的速度通常是进行float32计算的2到4倍，这带来了许多的好处：

储存模型所需的硬盘空间或内存空间更小。
从硬盘中加载模型到内存或显存的速度更快。
加快模型的推理速度，显著降低延迟，这在一些延迟敏感的场景非常有意义。

想要得到模型量化的这些优点，只需要付出模型效果轻微下降的代价，这在一些场景中是利远大于弊的。

模型量化的分类

模型量化一般分为三种：

训练后动态量化（Post Training Dynamic Quantization, PTDQ）
静态量化（Post Training Static Quantization, PTSQ）
量化感知训练（Quantization Aware Training, QAT）

考虑到篇幅原因，本文将介绍并演示前两种方式。本系列的后续文章中将介绍第三种方式以及这些方式的深层次原理。

Pytorch中的模型量化

Pytorch提供了两套用于模型量化的API，分别为Eager Mode Quantization和Fx Graph Mode Quantization。其中，Eager Mode已经进入了Beta阶段，但是Fx Graph Mode还在原形开发阶段。不过，因为Fx Graph Mode的API设计更加易用，支持的功能更多，所以本文依旧选择Fx Graph Mode进行演示。你可以在Pytorch Quantization Docs的对应章节中找到这两种API的详细对比。

准备工作

本文选取一个经典的卷积神经网络VGG16进行演示。我们会在CIFAR10数据集上评测原模型和量化后模型的计算速度和预测准确率。首先，我们需要定义一个VGG16模型并加载在CIFAR10数据集上预训练好的VGG16的权重。

%pip install torch torchvision --no-cache
%pip install matplotlib --no-cache

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

from torch.ao.quantization import (
  get_default_qconfig_mapping,
  get_default_qat_qconfig_mapping,
  QConfigMapping,
)
import torch.ao.quantization.quantize_fx as quantize_fx
import copy

# VGG16模型结构
# Modified from https://github.com/kuangliu/pytorch-cifar/blob/master/models/vgg.py

vgg16_cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
class VGG16(nn.Module):
    def __init__(self):
        super(VGG16, self).__init__()
        self.features = self._make_layers(vgg16_cfg)
        self.classifier = nn.Linear(512, 10)

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

    def _make_layers(self, cfg):
        layers = []
        in_channels = 3
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace=True)]
                in_channels = x
        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
        return nn.Sequential(*layers)

# 加载 CIFAR10 数据集
batch_size = 64

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

trainset = torchvision.datasets.CIFAR10(root='/bohr/cifar10-h7hf/v2', train=True, download=False, transform=transform_train)
train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True)

testset = torchvision.datasets.CIFAR10(root='/bohr/cifar10-h7hf/v2', train=False, download=False, transform=transform_test)
test_loader = DataLoader(testset, batch_size=batch_size, shuffle=False)
n_test = len(testset)

这里我们定义一些之后会用到的工具函数。

import time
import os
# 用于测试模型精度以及速度的函数
def test(model, test_loader, debug = False):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        start_time = time.time()
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        end_time = time.time()
    accuracy = 100 * correct / total
    time_cost = end_time - start_time
    if debug:
        print('Accuracy of the network on the %d test images: %.2f %%' % (n_test, accuracy))
        print('Time cost: %.2f s' % time_cost)
    return accuracy, time_cost

# 用于获取模型大小的函数，单位为 MB
def get_model_size(model):
    torch.save(model.state_dict(), "temp.pth")
    size = os.path.getsize("temp.pth")/1e6
    os.remove('temp.pth')
    return size

......

全文点击查看：

Pytorch模型量化（一）：快速上手