torch.fx 量化支持——FX GRAPH MODE QUANTIZATION

torch.fx目前支持的量化方式：

• Post Training Quantization
  • Weight Only Quantization
  • Dynamic Quantization
  • Static Quantization
• Quantization Aware Training
  • Static Quantization

其中，Post Training Quantization中的Static Quantization和Dynamic Quantization提供了demo。

與Eager模式對比

簡單來說，fx提供一個Graph模式：

• 可以自動插入量化節點（如quantize和dequantize），不需要手動修改當前的network及forward
• 這個模式下可以看到forward是怎么被自動構建的，可以進行更精細的調整

Graph模式

局限：只有可以轉換為符號的部分（symbolically traceable）可以被量化，Data dependent control flow是不支持的。如果模型有些部分無法被符號化，則量化只能在模型的部分上工作，不能被符號化的部分會被跳過。

如果需要這些部分被量化：

• 重寫代碼讓這些部分symbolically traceable
• 將這些部分轉換成observed和quantized的子模塊

相關的具體操作見(PROTOTYPE) FX GRAPH MODE QUANTIZATION USER GUIDE。

訓練后量化嘗試

環境準備：

import torch import copy from torch.quantization import get_default_qconfig from torch.quantization.quantize_fx import prepare_fx, convert_fx, fuse_fx

步驟

準備訓練好的權重、數據及網絡模型

初始化網絡，加載訓練好的權重（一般用copy.deepcopy保留原始模型），并將其置于eval模式：

float_model = load_model(saved_model_dir + float_model_file).to("cpu") float_model.eval() model_to_quantize = copy.deepcopy(float_model) model_to_quantize.eval()

指定量化模型的qconfig_dict

qconfig = get_default_qconfig("fbgemm") qconfig_dict = {"": qconfig}

qconfig是QConfig的一個實例，QConfig這個類就是維護了兩個observer，一個是activation所使用的observer，一個是op權重所使用的observer。

backendactivationweight

fbgemm (x86)	HistogramObserver (reduce_range=True)	PerChannelMinMaxObserver (default_per_channel_weight_observer)
qnnpack (arm)	HistogramObserver (reduce_range=False)	MinMaxObserver (default_weight_observer)
default	MinMaxObserver (default_observer)	MinMaxObserver (default_weight_observer)

準備模型并打印模型：

prepared_model = prepare_fx(model_to_quantize, qconfig_dict) print(prepared_model.graph)

模型較準

def calibrate(model, data_loader):model.eval()with torch.no_grad():for image, target in data_loader:model(image) calibrate(prepared_model, data_loader_test) # run calibration on sample data

量化模型

quantized_model = convert_fx(prepared_model) print(quantized_model)

對比量化前后，評估量化效果，包括模型大小、性能、時延等

總結

以上是生活随笔為你收集整理的利用torch.fx进行后量化的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。