PPQ量化torch模型demo

# OpenPPL CUDA 目前的量化方案为：量化位宽为 8bit。
# 其中 activation 和 weight 为 8bit， bias 为 32bit，量化策略为 weight 逐通道 (per-channel) 、activation 逐层 (per-layer) 的线性量化。
from typing import Iterable
import torch
import torchvision
from torch.utils.data import DataLoader
# QuantizationSettingFactory：管理量化过程
from ppq import BaseGraph,QuantizationSettingFactory,TargetPlatform
from ppq.api import export_ppq_graph,quantize_torch_model
BATCHSIZE = 32
INPUT_SHAPE = [3,224,224]
DEVICE = 'cuda'
#选择量化平台，ppq会根据此决定量化策略
PLATFORM = TargetPlatform.PPL_CUDA_INT8
# calibration_dataset需要从训练数据中随机抽样，1024个样本比较好
# 因为scale，offset需要实际跑一些数据才能得出
def load_calibration_dataset() -> Iterable:
	# torch.rand：均匀分布
	return [torch.rand(size=INPUT_SHAPE) for _ in range (32)]
def collate_fn(batch: torch.Tensor) -> torch.Tensor:
	return  batch.to(DEVICE)
model = torchvision.models.mobilenet.mobilenet_v2(pretrained=True)
model = model.to(DEVICE)
# 用cuda设置
quant_setting = QuantizationSettingFactory.pplcuda_setting()
# 采用分层均衡算法
quant_setting.equalization = True
# 量化优化。将网络切分成多个子图，然后将每个子图单独优化
# quant_setting.advanced_optimization = True
# 以保守的方式调度网络，将所有争议区算子视作不可量化的
quant_setting.dispatcher = 'conservative'
calibration_dataset = load_calibration_dataset()
calibration_dataloader = DataLoader(
	dataset = calibration_dataset,
	batch_size = BATCHSIZE,shuffle = True)
# calibration_steps会覆盖数据集长度。就是只取32个batch的数据。如果没有这么多数据，就是来回跑多遍，凑够这么多步
# 在这个例子中，只有一个batch，这个batch会被跑32遍
# calibration_steps：必须在8-512之间
# input_shape：指明网络可接受的数据尺寸
# 如果网络为多输入，需要将input_shape=None，然后设置inputs=从数据中拿一个，比如：calibration_dataloader[0]
# platform：确定量化规则
quantized = quantize_torch_model(
	model=model,calib_dataloader=calibration_dataloader,
	calib_steps=32,input_shape=[BATCHSIZE] + INPUT_SHAPE,
	setting=quant_setting,collate_fn=collate_fn,platform=PLATFORM,
# onnx.model：torch转换出来的onnx模型，没有跑过量化
	onnx_export_file='Output/onnx.model',device=DEVICE,verbose=0)
# PPQ上的量化图带计算信息
assert isinstance(quantized,BaseGraph)
# 导出量化计算图，此时platform的意思是导出到什么平台上，根据导出平台修改导出格式
# 如果两次设置platform不一致，会导致不小的量化误差
export_ppq_graph(graph=quantized,platform=PLATFORM,
	graph_save_to = 'Output/quantized(onnx).onnx',
	config_save_to = 'Output/quantized(onnx).json')