2.11、自定义图融合过程与量化管线

introduction

介绍如何自定义量化优化过程，以及如何手动调用优化过程

code

from typing import Callable, Iterableimport torch
import torchvisionfrom ppq import (BaseGraph, QuantizationOptimizationPass,QuantizationOptimizationPipeline, QuantizationSetting,TargetPlatform, TorchExecutor)
from ppq.api import ENABLE_CUDA_KERNEL
from ppq.executor.torch import TorchExecutor
from ppq.IR.quantize import QuantableOperation
from ppq.IR.search import SearchableGraph
from ppq.quantization.optim import (ParameterQuantizePass,PassiveParameterQuantizePass,QuantAlignmentPass, QuantizeRefinePass,QuantizeSimplifyPass,RuntimeCalibrationPass)
from ppq.quantization.quantizer import TensorRTQuantizer# ------------------------------------------------------------
# 在这个例子中，我们将向你介绍如何自定义量化优化过程，以及如何手动调用优化过程
# ------------------------------------------------------------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)# ------------------------------------------------------------
# 下面，我们将向你展示如何自定义图融合过程