PAI-Blade支持TensorFlow和PyTorch模型在GPU硬件、端側設備的INT8量化。本文主要介紹在GPU硬件上,如何使用PAI-Blade量化優化。
背景信息
量化是模型壓縮常用手段之一,通過將原始的全精度32位浮點數分桶量化為位寬更小的定點整數,可以達到節省訪存開銷、提升指令計算吞吐率的雙重目的(需要底層計算硬件的支持)。
TensorFlow模型量化
對TensorFlow模型進行基本優化的方法詳情請參見優化TensorFlow模型。此外,您還可以開啟量化功能,即使用PAI-Blade優化時指定optimization_level='o2'
。如果硬件支持INT8計算并存在量化性能收益,則PAI-Blade會以默認方式進行量化。其中:
如果不額外提供校正數據集(Calibration_Dataset),則PAI-Blade執行在線INT8量化方案。
如果希望獲得進一步的加速效果,建議您提供一組用于離線計算量化參數的校正數據集,PAI-Blade將自動開啟離線INT8量化。
對于TensorFlow模型,校正數據集是一個包含若干組feed_dict的列表,示例如下。
# 準備校正數據集。
import numpy as np
calib_data = list()
for i in range(10):
# feed_dict中value必須全部為np.ndarray類型。
feed_dict = {'input:0': np.ones((32, 224, 224, 3), dtype=np.float32)}
calib_data.append(feed_dict)
您可以參考如下步驟對TensorFlow模型進行量化優化。
下載示例模型、測試數據及校正數據集。
wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/frozen.pb wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/test_bc32.npy wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/calib_data_test_bc32.npy
加載待優化模型和相應數據。
import numpy as np import tensorflow as tf # 加載模型。 graph_def = tf.GraphDef() with open('frozen.pb', 'rb') as f: graph_def.ParseFromString(f.read()) # 加載測試數據。 test_data = np.load('test_bc32.npy', allow_pickle=True, encoding='bytes',).item() # 加載校正數據集。 calib_data = np.load('calib_data_test_bc32.npy', allow_pickle=True, encoding='bytes',)
執行離線INT8優化。
import blade optimized_model, opt_spec, report = blade.optimize( model=graph_def, optimization_level='o2', device_type='gpu', test_data=test_data, calib_data=calib_data )
驗證優化后的模型精度。
量化優化完成后,您可以在完整的測試數據集上驗證是否存在明顯的精度損失。如果量化后的模型仍然滿足精度要求,則無需繼續執行后續嘗試,反之可以嘗試修改量化配置(quant_config),以緩解精度損失。關于量化配置的詳細信息請參見blade.Config。
目前,TensorFlow模型在GPU上的量化配置僅支持weight_adjustment,啟用該設置后,PAI-Blade將自動調整模型參數以減小量化精度損失。啟用該設置的示例代碼如下所示。
quant_config = { 'weight_adjustment': 'true' # 默認為'false'。 } optimized_model, opt_spec, report = blade.optimize( model=graph_def, optimization_level='o2', device_type='gpu', test_data=test_data, calib_data=calib_data, config=blade.Config(quant_config=quant_config) )
PyTorch模型量化
與TensorFlow模型量化類似,您只需要使用PAI-Blade優化時指定optimization_level='o2'
即可開啟量化功能。但是與TensorFlow模型量化不同,PyTorch模型量化僅支持離線方式。因此,開啟量化功能時,需要提供一組用于離線計算量化參數的校正數據集。
PyTorch模型的校正數據集是一個包含若干組輸入數據的列表,示例如下。
# 準備校正數據集。
import numpy as np
calib_data = list()
for i in range(10):
image = torch.ones(32, 3, 224, 224)
calib_data.append(image)
您可以參考如下步驟對PyTorch模型進行量化優化。
下載示例模型、測試數據及校正數據集。
wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/traced_model.pt wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/test_bc32.pth wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/calib_data_test_bc32.pth
加載待優化模型和相應數據。
import torch # 加載模型。 pt_model = torch.jit.load('traced_model.pt') # 加載測試數據。 test_data = torch.load('test_bc32.pth') # 加載校正數據集。 calib_data = torch.load('calib_data_test_bc32.pth')
執行離線INT8優化。
import blade optimized_model, opt_spec, report = blade.optimize( model=pt_model, optimization_level='o2', device_type='gpu', test_data=test_data, calib_data=calib_data )