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+语义分割实战