PPQ中targetPlatform使用demo

# 在这个脚本中将创建一个新的量化平台，定义自己的量化规则
# 意味着将创建自己的Quantizer，并调节QuantSetting中的各项属性
from typing import Iterable
import torch
import torchvision
from ppq import (BaseQuantizer,Operation,OperationQuantizationConfig,
		TargetPlatform)
from ppq.api import (ENABLE_CUDA_KERNEL,export_ppq_graph,
		quantize_torch_model,register_network_exporter,
		register_network_quantizer)
from ppq.core import (QuantizationPolicy,QuantizationProperty,
		QuantizationStates,RoundingPolicy,convert_any_to_numpy)
from ppq.IR import BaseGraph,Operation,QuantableOperation

BATCHSIZE = 32
INPUT_SHAPE = [BATCHSIZE,3,224,224]
DEVICE = 'cuda'
PLATFORM = TargetPlatform.TRT_INT8
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)
# 定义自己的量化器，并完成注册
# 自定义量化器需要你完成所有接口函数与接口属性（接口属性必须以@property修饰）
class MyQuantizer(BaseQuantizer):
	# quant_operation_types是一个类型枚举，在这里需要写下所有该量化器所需要量化的算子
	@ property
	def quant_operation_types(self) -> set:
		return {'Conv'}
	# 一旦确定了哪些算子需要量化，则需要在init_quantize_config为他们初始化量化信息
	# 然而需要注意的是，由于手动调度的存在，用户可以强制调度一个类型不在quant_operation_types中的算子来到量化平台
	# 建议针对这种情况进行回应。或者，在探测到算子类型并非可量化类型后进行报错
	def init_quantize_config(self, operation:Operation) -> OperationQuantizationConfig:
		# 为卷积算子初始化量化信息，只量化卷积算子的输出
		if operation.type == 'Conv':
			config = self.create_default_quant_config(
				op = operation,
				num_of_bits = 8,
				quant_max = 127,
				quant_min = -128,
				observer_algorithm = 'percentile',
				policy = QuantizationPolicy(
					QuantizationProperty.PER_TENSOR +
					QuantizationProperty.LINEAR +
					QuantizationProperty.SYMMETRICAL),
				rounding = RoundingPolicy.ROUND_HALF_EVEN)
			# 关闭所有输入量化，状态设置为fp32
			for tensor_quant_config in config.input_quantization_config:
				tensor_quant_config.state = QuantizationStates.FP32
			return config
		else:
			raise TypeError(f'Unsupported Op Type: {operation.type}')
	# 当前量化器进行量化的算子都将被发往一个指定的目标平台
	# 这里选择TargetPlatform.EXTENSION作为目标平台
	@ property
	def target_platform(self) -> TargetPlatform:
		return TargetPlatform.EXTENSION
	@ property
	def activation_fusion_types(self) -> set:
		# 列举此处的算子会与他们之前的卷积和矩阵乘法执行激活函数图融合
		# 从而消去一些错误的量化行为，对于更为复杂的图融合，必须手写一个单独的优化过程
		# 不过这里并没有声明relu，clip是量化类型，因此图融合并不会起作用
		# 图融合起作用的前提是参与融合的所有算子，全部都需要被量化
		return {'Relu', 'Clip'}
# 注册量化器，目标平台为TargetPlatform.EXTENSION
# 后面将以这个平台来调用量化器
register_network_quantizer(MyQuantizer,TargetPlatform.EXTENSION)
with ENABLE_CUDA_KERNEL():
	# 以TargetPlatform.EXTENSION作为目标平台调用量化器
	quantized = quantize_torch_model(
		model=model,calib_dataloader=CALIBRATION,
		calib_steps=32,input_shape=INPUT_SHAPE,
		collate_fn=collate_fn,platform=TargetPlatform.EXTENSION,
		onnx_export_file='Output/onnx.model',device=DEVICE,verbose=0)
import onnx
# 完成量化之后，开始导出量化信息，这里注册一个网络量化信息导出器，
# 对于一些网络信息导出器而言，需要关注一下两件事：
	# 1.对于一些场景而言，需要导出PPQ的量化信息到特定的文件
	# 2.有些时候使用PPQ训练过你的网络，你可能还需要导出训练后的网络权重
# 这里展示如何导出量化信息到文本文档，而网络权重的导出则继承于onnx
from ppq.parser import OnnxExporter
# PPQ可以将计算图导出成onnx标准格式，onnx exporter不会导出QDQ节点
# 如需导出带有QDQ节点的onnx文件，用户需要使用onnxruntime exporter
# 任何导出器的导出逻辑都是原地进行的，他们对将传入的计算图对象进行原地修改，因此在导出之前需要手动克隆计算图
class MyExporter(OnnxExporter):
	def convert_value(self, value: torch.Tensor) -> str:
		if type(value) in {int, float}: return value
		else:
			value = convert_any_to_numpy(value,accept_none=True)
			if value is None: return value # SOI config has None as its scale and
			return value.tolist()
	def export(self, file_path: str, graph: BaseGraph, config_path: str, **kwargs):
		# 接下来导出量化信息，在PPQ中所有量化信息都绑定在Op上
		# 因此需要遍历图中所有的Op，将绑定在其上的量化信息导出到文件
		with open(config_path, 'w') as file:
			# name：算子的名字，op：算子的opset
			for name, op in graph.operations.items():
				# QuantableOperation：量化算子
				if not isinstance(op, QuantableOperation): continue
				# cfg：当前量化算子的量化信息。OperationQuantizationConfig 结构体包含了量化算子的所有输入、输出量化信息
				for cfg, var in op.config_with_variable:
					file.write(f"{name}: {var.name}\n")
					file.write(f"Quant State: {cfg.state.name}\n")
					file.write(f"Scale: {self.convert_value(cfg.scale)}\n")
					file.write(f"Offset: {self.convert_value(cfg.offset)}\n")
		# 最后导出完整的计算图到onnx
		onnx.save(self.export_graph(graph=graph),file_path)
# 注册导出器并导出模型
register_network_exporter(exporter=MyExporter,platform=TargetPlatform.EXTENSION)
export_ppq_graph(graph=quantized,platform=TargetPlatform.EXTENSION,
		graph_save_to='Output/model.onnx',config_save_to='Output/model.json')