Megengine量化
Megengine量化
量化指的是將浮點數模型（一般是32位浮點數）的權重或激活值用位數更少的數值類型（比如8位整數、16位浮點數）來近似表示的過程。 量化后的模型會占用更小的存儲空間，還能夠利用許多硬件平臺上的專屬算子進行提速。比如在 MegEngine 中使用8位整數來進行量化，相比默認的32位浮點數，模型大小可以減少為1/4，而運行在特定的設備上其計算速度也能提升為2-4倍。
量化的目的是為了追求極致的推理計算速度，為此舍棄了數值表示的精度，直覺上會帶來較大的模型掉點，但是在使用一系列精細的量化處理之后，其掉點可以變得微乎其微，并能支持正常的部署使用。而且近年來隨著專用神經網絡加速芯片的興起，低比特非浮點的運算方式越來越普及，因此如何把一個 GPU 上訓練的浮點數模型轉化為低比特的量化模型，就成為了工業界非常關心的話題。
一般來說，得到量化模型的轉換過程按代價從低到高可以分為以下4種：
圖1. 量化轉換過程分類
? Type1 和 Type2 由于是在模型浮點模型訓練之后介入，無需大量訓練數據，故而轉換代價更低，被稱為后量化（Post Quantization）；
? Type3 和 Type4 則需要在浮點模型訓練時就插入一些假量化（FakeQuantize）算子，模擬計算過程中數值截斷后精度降低的情形，故而稱為量化感知訓練（Quantization Aware Training, QAT）。
本文主要介紹 Type2 和 Type3 在 MegEngine 中的完整流程，事實上，除了 Type2 無需進行假量化，兩者的整體流程完全一致。
整體流程
以 Type3 為例，一般以一個訓練完畢的浮點模型為起點，稱為 Float 模型。包含假量化算子的用浮點操作來模擬量化過程的新模型，稱為 Quantized-Float 模型，或者 QFloat 模型。可以直接在終端設備上運行的模型，稱為 Quantized 模型，簡稱 Q 模型。
而三者的精度一般是 Float > QFloat > Q ，故而一般量化算法也就分為兩步：
? 拉近 QFloat 和 Q，這樣訓練階段的精度可以作為最終 Q 精度的代理指標，這一階段偏工程；
? 拔高 QFloat 逼近 Float，這樣就可以將量化模型性能盡可能恢復到 Float 的精度，這一階段偏算法。
典型的三種模型在三個階段的精度變化如下：
圖2. 三階段模型的精度變化
對應到具體的 MegEngine 接口中，三階段如下：
- 基于 Module 搭建網絡模型，并按照正常的浮點模型方式進行訓練；
- 使用 quantize_qat() 將浮點模型轉換為 QFloat 模型，其中可被量化的關鍵 Module 會被轉換為 QATModule ，并基于量化配置 QConfig 設置好假量化算子和數值統計方式；
- 使用 quantize() 將 QFloat 模型轉換為 Q 模型，對應的 QATModule 則會被轉換為 QuantizedModule ，此時網絡無法再進行訓練，網絡中的算子都會轉換為低比特計算方式，即可用于部署了。
該流程是 Type3 對應 QAT 的步驟，Type2 對應的后量化則需使用不同 QConfig，且需使用 evaluation 模式運行 QFloat 模型，而非訓練模式。更多細節可以繼續閱讀下一節詳細的接口介紹。
接口介紹
在 MegEngine 中，最上層的接口是配置如何量化的 QConfig 和模型轉換模塊里的 quantize_qat() 與 quantize() 。
QConfig
QConfig 包括了 Observer 和 FakeQuantize 兩部分。知道，對模型轉換為低比特量化模型一般分為兩步：一是統計待量化模型中參數和 activation 的數值范圍（scale）和零點（zero_point），二是根據 scale 和 zero_point 將模型轉換成指定的數值類型。而為了統計這兩個值，需要使用 Observer。
Observer 繼承自 Module ，也會參與網絡的前向傳播，但是其 forward 的返回值就是輸入，所以不會影響網絡的反向梯度傳播。其作用就是在前向時拿到輸入的值，并統計其數值范圍，并通過 get_qparams() 來獲取。所以在搭建網絡時把需要統計數值范圍的的 Tensor 作為 Observer 的輸入即可。
forward of MinMaxObserver
def forward(self, x_orig):
if self.enabled:
# stop gradient
x = x_orig.detach()
# find max and min
self.min_val._reset(F.minimum(self.min_val, x.min()))
self.max_val._reset(F.maximum(self.max_val, x.max()))
return x_orig
如果只觀察而不模擬量化會導致模型掉點，于是需要有 FakeQuantize 來根據 Observer 觀察到的數值范圍模擬量化時的截斷，使得參數在訓練時就能提前“適應“這種操作。FakeQuantize 在前向時會根據傳入的 scale 和 zero_point 對輸入 Tensor 做模擬量化的操作，即先做一遍數值轉換再轉換后的值還原成原類型，如下所示：
def fake_quant_tensor(inp: Tensor, qmin: int, qmax: int, q_dict: Dict) -> Tensor:
scale = q_dict[“scale”]
zero_point = 0
if q_dict[“mode”] == QuantMode.ASYMMERTIC:
zero_point = q_dict[“zero_point”]
# Quant
oup = Round()(inp / scale) + zero_point
# Clip
oup = F.minimum(F.maximum(oup, qmin), qmax)
# Dequant
oup = (oup - zero_point) * scale
return oup
目前 MegEngine 支持對 weight/activation 兩部分的量化，如下所示：
ema_fakequant_qconfig = QConfig(
weight_observer=partial(MinMaxObserver, dtype=“qint8”, narrow_range=True),
act_observer=partial(ExponentialMovingAverageObserver, dtype=“qint8”, narrow_range=False),
weight_fake_quant=partial(FakeQuantize, dtype=“qint8”, narrow_range=True),
act_fake_quant=partial(FakeQuantize, dtype=“qint8”, narrow_range=False),
)
這里使用了兩種 Observer 來統計信息，而 FakeQuantize 使用了默認的算子。
如果是后量化，或者說 Calibration，由于無需進行 FakeQuantize，故而其 fake_quant 屬性為 None 即可：
calibration_qconfig = QConfig(
weight_observer=partial(MinMaxObserver, dtype=“qint8”, narrow_range=True),
act_observer=partial(HistogramObserver, dtype=“qint8”, narrow_range=False),
weight_fake_quant=None,
act_fake_quant=None,
)
除了使用在 megengine.quantization.qconfig 里提供的預設 QConfig，也可以根據需要靈活選擇 Observer 和 FakeQuantize 實現自己的 QConfig。目前提供的 Observer 包括：
? MinMaxObserver ，使用最簡單的算法統計 min/max，對見到的每批數據取 min/max 跟當前存的值比較并替換，基于 min/max 得到 scale 和 zero_point；
? ExponentialMovingAverageObserver ，引入動量的概念，對每批數據的 min/max 與現有 min/max 的加權和跟現有值比較；
? HistogramObserver ，更加復雜的基于直方圖分布的 min/max 統計算法，且在 forward 時持續更新該分布，并根據該分布計算得到 scale 和 zero_point。
對于 FakeQuantize，目前還提供了 TQT 算子，另外還可以繼承 _FakeQuant 基類實現自定義的假量化算子。
在實際使用過程中，可能需要在訓練時讓 Observer 統計并更新參數，但是在推理時則停止更新。 Observer 和 FakeQuantize 都支持 enable() 和 disable() 功能，且 Observer 會在 train() 和 train() 時自動分別調用 enable/disable。
所以一般在 Calibration 時，會先執行 net.eval() 保證網絡的參數不被更新，然后再執行 enable_observer(net) 來手動開啟 Observer 的統計修改功能。
模型轉換模塊與相關基類
QConfig 提供了一系列如何對模型做量化的接口，而要使用這些接口，需要網絡的 Module 能夠在 forward 時給參數、activation 加上 Observer 和進行 FakeQuantize。轉換模塊的作用就是將模型中的普通 Module 替換為支持這一系列操作的 QATModule ，并能支持進一步替換成無法訓練、專用于部署的 QuantizedModule 。
基于三種基類實現的 Module 是一一對應的關系，通過轉換接口可以依次替換為不同實現的同名 Module。同時考慮到量化與算子融合（Fuse）的高度關聯，提供了一系列預先融合好的 Module，比如 ConvRelu2d 、 ConvBn2d 和 ConvBnRelu2d 等。除此之外還提供專用于量化的 QuantStub 、 DequantStub 等輔助模塊。
轉換的原理很簡單，就是將父 Module 中可被量化（Quantable）的子 Module 替換為對應的新 Module。但是有一些 Quantable Module 還包含 Quantable 子 Module，比如 ConvBn 就包含一個 Conv2d 和一個 BatchNorm2d，轉換過程并不會對這些子 Module 進一步轉換，原因是父 Module 被替換之后，其 forward 計算過程已經完全不同了，不會再依賴于這些子 Module。
Note
如果需要使一部分 Module 及其子 Module 保留 Float 狀態，不進行轉換，可以使用 disable_quantize() 來處理。
如果網絡結構中涉及一些二元及以上的 ElementWise 操作符，比如加法乘法等，由于多個輸入各自的 scale 并不一致，必須使用量化專用的算子，并指定好輸出的 scale。實際使用中只需要把這些操作替換為 Elemwise 即可，比如 self.add_relu = Elemwise(“FUSE_ADD_RELU”)
由于轉換過程修改了原網絡結構， 網絡的訓練和測試 中提到的模型保存與加載無法直接適用于轉換后的網絡，讀取新網絡保存的參數時，需要先調用轉換接口得到轉換后的網絡，才能用 load_state_dict 將參數進行加載。
實例講解
下面以 ResNet18 為例來講解量化的完整流程，完整代碼見 MegEngine Models 。主要分為以下幾步：
-
修改網絡結構，使用已經 Fuse 好的 ConvBn2d、ConvBnRelu2d、ElementWise 代替原先的 Module；
-
在正常模式下預訓練模型，并在每輪迭代保存網絡檢查點；
-
調用 quantize_qat() 轉換模型，并進行 finetune；
-
調用 quantize() 轉換為量化模型，并執行 dump 用于后續模型部署。
網絡結構見 resnet.py ，相比慣常寫法，修改了其中一些子 Module，將原先單獨的 conv, bn, relu 替換為 Fuse 過的 Quantable Module。
class BasicBlock(Module):
def init(self, in_planes, planes, stride=1):
super(BasicBlock, self).init()
self.conv_bn_relu = ConvBnRelu2d(
in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
)
self.conv_bn = ConvBn2d(
planes, planes, kernel_size=3, stride=1, padding=1, bias=False
)
self.add_relu = Elemwise(“FUSE_ADD_RELU”)
self.shortcut = Sequential()
if stride != 1 or in_planes != planes:
self.shortcut = Sequential(
ConvBn2d(in_planes, planes, kernel_size=1, stride=stride, bias=False)
)def forward(self, x):
out = self.conv_bn_relu(x)
out = self.conv_bn(out)
cut = self.shortcut(x)
out = self.add_relu(out, cut)
return out
然后對該模型進行若干輪迭代訓練，并保存檢查點，這里省略細節：
for step in range(0, total_steps):Linear learning rate decay
epoch = step // steps_per_epoch
learning_rate = adjust_learning_rate(step, epoch)image, label = next(train_queue)
image = tensor(image.astype(“float32”))
label = tensor(label.astype(“int32”))n = image.shape[0]
loss, acc1, acc5 = train_func(image, label, net, gm)
optimizer.step()
optimizer.clear_grad()
調用 quantize_qat() 來將網絡轉換為 QATModule：
from megengine.quantization import ema_fakequant_qconfig
from megengine.quantization.quantize import quantize_qat
model = ResNet18()
if args.mode != “normal”:
quantize_qat(model, ema_fakequant_qconfig)
使用默認的 ema_fakequant_qconfig 來進行 int8 量化。
然后繼續使用上面相同的代碼進行 finetune 訓練。值得注意的是，如果這兩步全在一次程序運行中執行，那么訓練的 trace 函數需要用不一樣的，因為模型的參數變化了，需要重新進行編譯。示例代碼中則是采用在新的執行中讀取檢查點重新編譯的方法。
在 QAT 模式訓練完成后，繼續保存檢查點，執行 inference.py 并設置 mode 為 quantized ，這里需要將原始 Float 模型轉換為 QAT 模型之后再加載檢查點。
from megengine.quantization.quantize import quantize_qat
model = ResNet18()
if args.mode != “normal”:
quantize_qat(model, ema_fakequant_qconfig)
if args.checkpoint:
logger.info(“Load pretrained weights from %s”, args.checkpoint)
ckpt = mge.load(args.checkpoint)
ckpt = ckpt[“state_dict”] if “state_dict” in ckpt else ckpt
model.load_state_dict(ckpt, strict=False)
模型轉換為量化模型包括以下幾步：
from megengine.quantization.quantize import quantize
定義trace函數，打開capture_as_const以進行dump
@jit.trace(capture_as_const=True)
def infer_func(processed_img):
model.eval()
logits = model(processed_img)
probs = F.softmax(logits)
return probs
執行模型轉換
if args.mode == “quantized”:
quantize(model)
準備數據
processed_img = transform.apply(image)[np.newaxis, :]
if args.mode == “normal”:
processed_img = processed_img.astype(“float32”)
elif args.mode == “quantized”:
processed_img = processed_img.astype(“int8”)
執行一遍evaluation
probs = infer_func(processed_img)
將模型 dump 導出
infer_func.dump(output_file, arg_names=[“data”])
至此，便得到了一個可用于部署的量化模型。
總結
以上是生活随笔為你收集整理的Megengine量化的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 腾讯 angel 3.0：高效处理模型
- 下一篇： DeepLabV3+语义分割实战