PPQ中optimization使用demo

# 自定义量化优化过程，手动调用优化过程
from typing import Iterable,Callable
import torch
import torchvision
from ppq import QuantizationSettingFactory,TargetPlatform
from ppq.api import (ENABLE_CUDA_KERNEL,QuantizationSettingFactory,
		quantize_torch_model)
from ppq.core import QuantizationStates
from ppq.executor.torch import TorchExecutor
from ppq.IR.quantize import QuantableOperation
BATCHSIZE = 32
INPUT_SHAPE = [BATCHSIZE,3,224,224]
DEVICE = 'cuda'
PLATFORM = TargetPlatform.TRT_INT8
# 创建calibration数据，以及加载模型
def load_calibration_dataset() -> Iterable:
	return [torch.rand(size=INPUT_SHAPE) for _ in range (32)]
CALIBRATION = load_calibration_dataset()
def collate_fn(batch: torch.Tensor) -> torch.Tensor:
	return batch.to(DEVICE)
model = torchvision.models.mobilenet.mobilenet_v2(pretrained=True)
model = model.to(DEVICE)
# 下面将展示如何不借助QSetting来自定义优化过程
# QSetting中包含了PPQ官方量化过程的配置参数，可以借助它来调用所有紫铜内置优化过程
# 但如果设计新的优化过程，必须在合适的时机手动启动它们
QSetting = QuantizationSettingFactory.default_setting()
# 不要进行parameter baking操作，一旦parameter完成baking，后续任何对于参数的修改都是不被允许的
# 你可以设置baking_parameter=True并再次执行这个脚本，PPQ会拒绝后续修改scale的请求
QSetting.quantize_parameter_setting.baking_parameter = False
# 定义自己的优化过程，继承QuantizationOptimizationPass基类，实现optimize接口
# 在optimize接口函数中，可以修改图的属性从而实现特定的目的
# 在这个例子中，将图中所有卷积的输入scale变为原来的两倍
# 同时，解除最后一个Gemm的输入量化
from ppq import BaseGraph,QuantizationOptimizationPass,TorchExecutor
class MyOptim(QuantizationOptimizationPass):
	def optimize(self,graph: BaseGraph,dataloader: Iterable,collate_fn: Callable,executor: TorchExecutor, **kwargs) -> None:
		# graph.operations是一个包含了图中所有op的字典
		for name,op in graph.operations.items():
			# 从图中找出所有已经量化的卷积算子
			# 对于你的网络而言，并非所有的算子最终都会被量化，它们会受到调度策略和Quantizer策略的双重限制
			# 因此要使用isinstance(op,QuantableOperation)来判断它是否是一个量化算子
			if op.type == 'Conv' and isinstance(op, QuantableOperation):
				# 对于卷积算子，它可能有2-3个输入，其中第二个输入为权重，第三个输入为bias
				# 修改权重量化信息的scale
				op.input_quant_config[1].scale *= 2
				print(f'Input scale of Op {name} has been enlarged.')
			# 接下来解除Gemm量化，在这里mobilenet_v2网络只有一个Gemm层
			# 所以将遇到的Gemm层全部解除量化
			if op.type == 'Gemm' and isinstance(op, QuantableOperation):
				# config_with_variable接口将返回量化算子的所有量化信息，包括输入和输出
				for cfg, _ in op.config_with_variable:
					# 在PPQ中有许多方法可以切换算子的量化状态
					# 将量化状态直接设置为FP32，即解除算子的量化
					cfg.state = QuantizationStates.FP32
				# 也可以直接调用算子的dequantize方法
				# op.dequantize()
with ENABLE_CUDA_KERNEL():
	quantized = quantize_torch_model(
		model=model,calib_dataloader=CALIBRATION,
		calib_steps=32,input_shape=INPUT_SHAPE,
		setting=QSetting,collate_fn=collate_fn,platform=PLATFORM,
		onnx_export_file='Output/model.onnx',device=DEVICE,verbose=0)
	# 在完成量化流程之后，调用自定义量化优化过程从而修改量化参数
	optim = MyOptim(name='My Optimization Procrdure')
	optim.optimize(graph=quantized,dataloader=CALIBRATION,
			collate_fn=INPUT_SHAPE,executor=TorchExecutor(quantized,device=DEVICE))