在深度学习中,量化指的是使用更少的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.Module): def __init__(self) -> None: super(ConvBnReluModel, self).__init__() self.conv = nn.Conv2d(3,5,3, bias=False) self.bn = nn.BatchNorm2d(5) self.relu = nn.ReLU(inplace=True) def forward(self, x): x = self.conv(x) x = self.bn(x) x = self.relu(x) return x m = ConvBnReluModel() m.eval() layers = [['conv','bn','relu']] f = quantization.fuse_modules(m,layers, inplace=True) types_to_quantize = {nn.Conv2d, nn.BatchNorm2d, nn.ReLU} q = quantization.quantize_dynamic(f, types_to_quantize, dtype=torch.qint8) s = torch.jit.script(q) torch.jit.save(s, 'quantize_model.pth')
Pytorch量化(torch.quantization)
最新推荐文章于 2023-10-29 21:59:47 发布