ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

神经网络量化代码实践:INT8推理显存压缩与精度调优

神经网络量化代码实践:INT8推理显存压缩与精度调优 上周帮朋友排查一个推理服务的性能问题,他那个模型在离线服务器上跑,权重文件不大,但推理时显存占用高得离谱, batch 稍微开大一点就报 OOM。我看了下配置,模型全程 float32 推理,一点没做优化。让他试了模型量化,权重压到 INT8 之后,显存直接降到原来的四分之一左右,单条推理延迟也从 40 多毫秒降到 20 毫秒出头,精度掉了不到一个百分点。这件事让我意识到,神经网络量化这件事,概念大家都听过,但真到代码层面能自己动手跑通、能说清楚参数是怎么算出来的人,其实没那么多。模型量化、神经网络量化、代码实践,这三个词单看是三个方向,串起来才是一条完整的链路——量化是目标,神经网络是对象,代码是落地手段。这篇文章写给两类人:一类是手里已经有能跑起来的模型,想在推理侧省显存、压延迟的工程同学;另一类是刚接触模型压缩,对量化只有把 float32 变成 int8这种模糊印象,想搞清楚底层到底在算什么的人。我不打算只停在概念层面,会把缩放因子怎么推导、零点为什么存在、逐通道和逐张量差在哪、代码为什么这么写、实际部署踩了哪些坑,都尽量摊开讲清楚。1. 拆解量化的核心逻辑:为什么精度换空间这笔账划算1.1 从一次显存告急说起很多人第一次接触量化,动机都很朴素:模型太大,跑不动。一个 7B 参数量的 transformer,光是权重用 float32 存就要 28GB 左右,别说消费级显卡,连一些专业卡都得掂量。把数据类型换成 float16 能砍一半,换成 INT8 再砍一半,理论上直接降到 7GB 上下。这个账算起来很直观,但真正让人犹豫的是下面这个顾虑:精度会不会崩?我实测过几个不同规模的模型,在合理量化配置下,分类任务的 top-1 精度通常掉 0.3 到 1.5 个百分点,检测任务的 mAP 掉 1 到 3 个点,语言模型的任务指标波动更明显一些。这个损失值不值得,取决于你的场景——如果是推荐排序这种对绝对精度没那么敏感、但对延迟和成本极其敏感的业务,量化几乎是无脑选的;如果是医疗影像诊断这类容错率极低的场景,那就要慎重评估,或者上量化感知训练去把精度找补回来。提示:量化不是万能的。它本质是一种用可控的精度损失换取显存、带宽和计算效率的工程取舍,不存在既不掉点又省一半资源的免费午餐,想清楚自己的容忍边界再动手。1.2 量化在数学上到底做了什么抛开框架的各种封装,量化的数学内核其实就一句话:用一个低比特整数去近似表示一个高比特浮点数,并且记录下这个近似关系。这个关系由两个东西确定,一个是缩放因子,通常记作 scale,另一个是零点,记作 zero_point。假设我们把一段浮点数据映射到 8 位整数,整数能表示的范围是 0 到 255,或者对称场景下是 -128 到 127。缩放因子的作用是把浮点的动态范围压缩到整数的动态范围里,零点的作用是保证浮点里的 0 在整数域里也有一个精确对应的表示。为什么零这么重要?因为在卷积和矩阵乘里,大量的 padding、mask 操作都依赖 0,如果 0 映射过去变成了别的数,那 padding 的地方就会引入系统性偏差。用生活里的例子类比,量化的过程就像把一把精度很高的游标卡尺(浮点)换成一把普通钢尺(整数)。钢尺的刻度是固定的、离散的,你量出来的值只能落在最近的那条刻度线上,所以必然有舍入误差。scale 就是每两条刻度之间代表多少毫米,zero_point 就是钢尺的零点是否和真实零点对齐。理解了这一层,后面所有的量化方案,其实都是在回答同一个问题:怎么让这把钢尺的误差最小。2. 量化方案怎么选:几种主流路线的取舍逻辑2.1 对称量化与非对称量化对称量化,顾名思义,它的映射区间是关于原点对称的,只用一个 scale 和一个隐含为 0 的零点。对于 8 位有符号整数,映射范围是 -127 到 127。它的计算简单,在硬件上做乘法的时候不用额外处理零点偏移,所以像英伟达的一些加速库对这种形式支持最好。非对称量化则要把实际的 min 和 max 都考虑进来,映射到 0 到 255 这样的无符号区间, scale 和 zero_point 都要算。它的优势是对于那种数据分布严重偏斜、零点不在中间的激活值,比如 ReLU 之后的输出全是正数,非对称量化能把整个整数范围都利用起来,精度通常更好一点。我用一个具体数字感受一下区别。假设某层激活值范围是 [0, 6],用对称量化的话,范围要先扩展到 [-6, 6], scale 6/127 ≈ 0.0472,能利用的整数档位只有正半轴的 127 档。换成非对称量化, scale 6/255 ≈ 0.0235,足足多了一倍的档位,分辨率直接翻倍。所以对于 ReLU 之后的激活,非对称几乎是默认选择。反过来,权重通常正负都有,分布接近对称,这时候对称量化既省事又不掉精度。2.2 训练后量化、量化感知训练、动态量化的差别按介入时机分,量化主要分三大类,这块选错了,后面怎么调参都白搭。训练后量化是最省事的,拿一个已经训练好的 float 模型,直接转成整型。它不需要重新训练,只需要一小批校准数据跑一遍,统计各层激活值的动态范围,算出 scale 和 zero_point。优点是快,几分钟到几十分钟就能搞定;缺点是精度损失不好控制,对于那些对量化敏感的模型,可能掉点掉得比较厉害。量化感知训练是在训练阶段就模拟量化的舍入误差,让模型在训练过程中学会去适应这种误差。前向传播时插入伪量化节点,把值先量化再反量化,反向传播时用直通估计器把梯度近似传过去。训练结束后,模型已经见过量化带来的扰动,推理时直接转换精度损失就小很多。代价是要重新训练,算力和时间成本都上去了。动态量化是个折中方案,它只量化权重,激活值在推理时动态地量化。对于 LSTM、Transformer 里那种激活分布变化比较大的层,这种模式比较友好,常见于 NLP 场景。Python 里用 PyTorch 一行代码就能跑,后面实操部分我会详细演示。2.3 量化粒度:逐张量、逐通道、逐组粒度这个问题经常被忽略,但它对精度的影响比很多人想象的大。逐张量是整层共享一个 scale,最省空间,一个张量只需要存一个 scale;逐通道是每个输出通道一个 scale,存储稍微多一点,但精度提升明显;逐组是把通道再细分,每几个通道一组,介于两者之间,在权重量化里很常见。我做过一组对比,同一个卷积层,逐张量量化的权重误差大约是逐通道的 4 到 6 倍。原因很好理解,卷积核不同通道的权重量级差异可能很大,一个通道的值范围是 0.01 量级,另一个通道是 0.5 量级,共用一个 scale,小量级的通道精度直接被牺牲掉。所以权重量化基本都默认逐通道,激活量化因为逐通道统计成本高,多数还是逐张量。实操心得:选量化方案的时候,有个粗略的判断顺序——先看激活是不是全正(是就非对称),再看权重是否逐通道(是就逐通道),最后看精度能不能达标(不能就上量化感知训练)。别一上来就追求最复杂的方案,很多场景训练后量化就够了。3. 手撸代码:从零实现对称量化与逐通道量化3.1 定点表示与缩放因子的计算在动框架之前,我强烈建议先自己用 numpy 或者原生 Python 实现一遍最基础的量化,把 scale 和 zero_point 的推导过程走一遍。这一步花不了半小时,但能让你后面调框架参数时心里有底,不再是把参数当黑盒乱试。对称量化的 scale 计算公式很直接。假设我们要量化到 b 位有符号整数,整数能表示的最大绝对值是 2^(b-1) - 1,当 b8 时就是 127。那 scale 就等于数据绝对值的最大值除以 127。import numpy as np def compute_symmetric_scale(x, bits8): qmax 2 ** (bits - 1) - 1 abs_max np.max(np.abs(x)) # 防止全零张量导致除零 scale abs_max / qmax scale max(scale, 1e-8) return scale非对称量化的 scale 和 zero_point 要一起算,零点必须落在整数范围内,所以算完还要做一次裁剪。def compute_asymmetric_params(x, bits8): qmin, qmax 0, 2 ** bits - 1 xmin, xmax np.min(x), np.max(x) scale (xmax - xmin) / (qmax - qmin) scale max(scale, 1e-8) zero_point qmin - xmin / scale zero_point np.clip(np.round(zero_point), qmin, qmax) return scale, zero_point这里有三个细节值得单独说。第一,scale 必须兜底一个极小值,否则当某一层激活全为零(比如某些 padding 后的层)时会除零,得到 inf 或 nan,推理直接崩。第二, zero_point 一定要四舍五入成整数再裁剪,它本身必须是整数,因为量化后的值要通过 q x/scale zero_point 得到,如果 zero_point 不是整数,那个加法的结果就不是定点数了。第三, xmin 和 xmax 的取值方式很关键,用真实极值容易被离群点带偏,实践中更常用分位数统计或者滑动平均,后面校准部分会展开。3.2 量化与反量化的完整实现有了 scale 和 zero_point,量化和反量化就是两个方向的一来一回。量化是把浮点转成整数,反量化是把整数还原成近似的浮点。def quantize(x, scale, zero_point, bits8): if zero_point 0: # 对称量化 qmin, qmax -(2 ** (bits - 1)), 2 ** (bits - 1) - 1 else: qmin, qmax 0, 2 ** bits - 1 q np.round(x / scale zero_point) q np.clip(q, qmin, qmax) return q.astype(np.int32) def dequantize(q, scale, zero_point): return (q.astype(np.float32) - zero_point) * scale反量化的公式是 反量化值 (量化值 - zero_point) * scale,这个减零点、乘缩放的操作在推理里是不可避免的,它引入的运算量相对小,但如果是逐通道,就得按照通道维度广播 scale 和 zero_point,广播写错了就会出现通道错位,这是新手常犯的错。测试的时候,如果发现量化前后整体误差不大,但某几个通道偏差特别明显,大概率就是广播维度对不上。3.3 逐通道量化的改造与精度对比把逐张量改成逐通道,核心改动是 scale 的计算维度。以卷积权重为例,权重的形状通常是 [out_channels, in_channels, kh, kw],逐通道就是沿着 out_channels 这一个维度,把后面三个维度压成一个最大值。def per_channel_symmetric_quantize(w, bits8): qmax 2 ** (bits - 1) - 1 # w: [oc, ic, kh, kw], 沿 (1,2,3) 求最大绝对值 scales np.max(np.abs(w), axis(1, 2, 3), keepdimsTrue) / qmax scales np.maximum(scales, 1e-8) q np.round(w / scales) q np.clip(q, -qmax - 1, qmax) return q.astype(np.int32), scales我拿一个真实的卷积层做过对比,整个权重张量的动态范围是 -1.8 到 1.8,但其中有将近三成的输出通道,值范围只有 -0.06 到 0.06。逐张量量化的 scale 是 1.8/127 ≈ 0.01417,那些小通道的量化误差达到了它们自身量级的 20% 以上,基本等于信息被抹平了。换成逐通道之后,小通道用自己约 0.00047 的 scale,相对误差降到百分之一以内。这组对比说明了为什么权重量化几乎清一色用逐通道,那点 scale 的额外存储开销换来的是精度上的质变。4. 真实模型上的量化实操:以 PyTorch 为例4.1 环境准备与基线模型搭建从手写实现切到真实框架,先要准备环境。PyTorch 从 1.x 之后就把量化 API 内置了,主流的量化后端有 fbgemm(针对 x86 服务器 CPU)和 qnnpack(针对 ARM 移动端)。如果你在服务器上跑,默认用 fbgemm,可以通过 torch.backends.quantized.engine 来切换。pip install torch torchvision我为了演示方便,用一个小的卷积分类网络做例子,真实场景里换成你自己的模型结构即可。import torch import torch.nn as nn class SmallNet(nn.Module): def __init__(self, num_classes10): super().__init__() self.conv1 nn.Conv2d(3, 16, 3, padding1) self.relu nn.ReLU() self.pool nn.AdaptiveAvgPool2d((4, 4)) self.fc nn.Linear(16 * 4 * 4, num_classes) def forward(self, x): x self.relu(self.conv1(x)) x self.pool(x) x torch.flatten(x, 1) return self.fc(x) model SmallNet() model.eval() # 量化前必须切到 eval,否则 BN 和 dropout 会干扰统计这里有个坑我必须提前强调,量化之前模型一定要调用 eval(),让它进入推理模式。如果你在 train 模式下做校准, BatchNorm 层会更新滑动均值, dropout 也还在起作用,统计出来的激活范围是错的,量化后的精度会莫名其妙地差。4.2 动态量化实操动态量化是最简单的一种,一行代码就能完成。它只量化权重,激活在运行时动态量化,特别适合 LSTM、GRU、Linear 密集的模型。import torch.quantization q_model torch.quantization.quantize_dynamic( model, {nn.Linear, nn.LSTM}, # 指定要量化的模块类型 dtypetorch.qint8 )跑完之后你对比一下模型大小,通常能压到原来的四分之一。用 dynamic 的好处是无需校准数据,直接就能转,坏处是它不量化卷积,所以对 CNN 为主的模型效果有限,主要收益点在线性层。我自己测过一个以 LSTM 为主的序列模型,动态量化后模型体积从 12MB 降到 3.2MB,单条推理延迟降了大约 35%。4.3 静态量化实操与校准静态量化是真正能发挥 CNN 性能优势的方案,它把权重和激活都量化成 INT8,但需要一小批校准数据来统计激活的动态范围。校准数据的量不用多,几百张通常就够,关键是要有代表性,覆盖到模型实际会遇到的输入分布。model.qconfig torch.quantization.get_default_qconfig(fbgemm) # 也可以指定逐通道权重量化 # model.qconfig torch.quantization.get_default_qconfig(fbgemm) # model.qconfig.weight torch.quantization.per_channel_weight_observer model_prepared torch.quantization.prepare(model, inplaceFalse) # 校准:让真实数据流经模型,收集激活分布 with torch.no_grad(): for images, _ in calib_loader: model_prepared(images) model_int8 torch.quantization.convert(model_prepared, inplaceFalse) torch.save(model_int8.state_dict(), model_int8.pth)校准这一步是整个静态量化的灵魂,校准集选得好不好,直接决定精度掉多少。我的经验是,校准数据量在 500 到 1000 条之间是个比较舒服的区间,太少统计不稳定,太多收益递减还费时间。另外注意校准阶段不用计算 loss,也不用反传,纯粹是前向跑数据让 observer 记录 min/max。4.4 精度与速度实测数据我把上面的流程完整跑了一遍,用一个虚拟数据集做对照,得到的数据不一定和你的场景完全一致,但趋势有参考价值。配置模型体积单条延迟(CPU)精度float32 基线1.0x1.0x100%动态量化0.27x0.68x98.9%静态量化(逐张量)0.25x0.41x97.6%静态量化(逐通道)0.26x0.43x99.1%从这张表能读出几个信息。动态量化体积压得最狠但速度提升一般,因为它没优化卷积;静态量化的加速最明显,延迟降到四成左右;逐通道相比逐张量,体积只多了一点点,但精度几乎追平了 float32,这就是逐通道的价值所在。需要注意的是,加速比和你模型里卷积的占比强相关,如果模型全是矩阵乘,动态量化可能反而更划算。5. 常见问题排查实录与避坑清单5.1 精度掉点严重怎么定位量化后精度大幅下降,先别急着上量化感知训练,按顺序排查通常能解决大半。第一步,确认模型在量化前是 eval 模式,且推理结果本身正常;第二步,检查校准数据是否和真实推理分布一致,用训练集里偏差很大的数据去校准是大忌;第三步,逐层对比量化前后的输出,定位到具体是哪一层掉得厉害。定位到敏感层之后,有两种处理思路。一种是给这个层保留 float 精度,在 PyTorch 里可以通过 qconfig 单独设置某些层不量化;另一种是回到训练阶段做微调,只微调量化敏感层或者整个模型,通常几轮就能把精度拉回来。我遇到过最极端的案例是某层激活范围里混进了一个巨大的离群值,导致 scale 被撑得很大,其他正常值全被压到很窄的整数区间里,处理办法就是改用分位数统计激活范围,把离群点的影响屏蔽掉。5.2 算子不支持与回退不是所有算子都支持量化,一些自定义算子、复杂的激活函数、某些归一化操作,框架会直接回退到 float 实现。回退本身不会报错,但它会打断整个量化图,导致局部还在跑 float,反而比不量化更慢,因为多了量化反量化的转换开销。排查办法是打印量化后的模型结构,把那些没有变成量化版本的算子揪出来。print(model_int8) # 关注有没有出现 fp32 后缀,或看起来没被转换的层如果你的模型里有大量不支持的算子,量化收益可能很小,这种情况下要么换等价的可量化算子重写这部分结构,要么放弃整模型量化,只对支持的子模块做局部量化。5.3 常见问题速查表现象可能原因处理方向量化后精度暴跌校准数据分布不对 / 敏感层换校准集 / 敏感层保持 float推理速度反而变慢算子回退 / 量化开销大于收益检查回退算子 / 换量化粒度输出全为同一类zero_point 或 scale 异常检查是否有 inf/nan / 全零层逐通道后某几通道偏差大广播维度错误核对 scale 的 shape模型体积没变小只量化了部分层确认 qconfig 覆盖范围运行时形状报错输入尺寸和校准时不符保证校准与推理输入一致注意:量化是一个工程方法和业务场景强绑定的活,没有一套参数能打遍天下。别人博客里的配置只能当起点,最终一定要在你的数据和硬件上实测。6. 踩坑之后的一些个人体会走了这么多弯路,我越来越觉得量化这件事的难点不在数学,而在对模型和数据分布的直觉。同样是 INT8,同样是逐通道,换个校准集精度能差好几个点,这时候靠的不是公式,是你对这个模型到底在算什么的理解。我的习惯是先在几个不同代表性子集上分别校准,看看精度波动范围,如果波动很小,说明这个模型对量化本身不敏感,可以放心推上线;如果波动很大,那就说明精度对数据分布高度依赖,得谨慎评估甚至补充训练。另外说说容易被混淆的一个点。搜索量化这个词的时候,会碰到大量关于量化交易、量化策略的内容,那是金融领域里的另一码事,和神经网络量化没有任何关系。初学者如果搜错了方向,看半天 K 线策略的代码会非常迷惑。神经网络量化讨论的是数值精度和计算效率,出发点是让模型跑得更省更快,这两个量化只是中文翻译撞了车,理解的时候务必分开。最后分享一个我个人觉得性价比很高的做法:如果你不想碰训练流程,又想让静态量化的精度稳一点,可以先把逐通道权重配置打开,再用分位数观察器限制激活范围,这两步几乎不增加成本,但通常能挽回一半以上的精度损失。真正需要上量化感知训练的场景,往往是那些本身就对数值极其敏感的模型,比如检测框回归头、语音合成里的声学特征。搞清楚自己的模型属于哪一类,再决定投入多少精力,比盲目上重方案要聪明得多。
返回列表