一、概述
在深度学习中,量化指的是使用更少的bit来存储原本以浮点数存储的tensor,以及使用更少的bit来完成原本以浮点数完成的计算。这么做的好处主要有如下几点:
- 更少的模型体积,接近4倍的减少;
- 可以更快的计算,由于更少的内存访问和更快的int8计算,可以快2~4倍。
一个量化后的模型,其部分或者全部的tensor操作会使用int类型来计算,而不是使用量化之前的float类型。当然,量化还需要底层硬件支持,x86 CPU(支持AVX2)、ARM CPU、Google TPU、Nvidia Volta/Turing/Ampere、Qualcomm DSP这些主流硬件都对量化提出了支持。
与典型的 FP32 型号相比,PyTorch 支持 INT8 量化,从而可将模型大小减少 4 倍,并将内存带宽要求减少 4 倍。 与 FP32 计算相比,对 INT8 计算的硬件支持通常快 2 到 4 倍。
量化主要是一种加速推理的技术,并且量化算子仅支持前向传递。PyTorch 支持多种方法来量化深度学习模型。
在大多数情况下,该模型在 FP32 中训练,然后将模型转换为 INT8
PyTorch对量化的支持目前有如下三种方式:
- Post Training Dynamic Quantization,模型训练完毕后的动态量化;
- Post Training Static Quantization,模型训练完毕后的静态量化;
- QAT(Quantization Aware Training),模型训练中开启量化。
二、模型量化后保存为jit形式
import torch
from torch import nn
from torch import quantization
class ConvBnReluModel(nn