使用PAI-Blade對TensorFlow與PyTorch模型進行GPU量化最佳化-人工智慧平台 PAI-阿里雲

PAI-Blade支援TensorFlow和PyTorch模型在GPU硬體、端側裝置的INT8量化。本文主要介紹在GPU硬體上，如何使用PAI-Blade量化最佳化。

背景資訊

量化是模型壓縮常用手段之一，通過將原始的全精度32位浮點數分桶量化為位寬更小的定點整數，可以達到節省訪存開銷、提升指令計算吞吐率的雙重目的（需要底層計算硬體的支援）。

TensorFlow模型量化

對TensorFlow模型進行基本最佳化的方法詳情請參見最佳化TensorFlow模型。此外，您還可以通過指定optimization_level='o2'來開啟PAI-Blade的量化功能。如果硬體支援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'即可開啟量化功能。但是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
)