地瓜官网
地瓜开发者社区
地瓜官网
地瓜开发者社区

quantization.prepare

horizon_plugin_pytorch.quantization.prepare (model: Module, example_inputs: Any | None = None, qconfig_setter: Tuple[QconfigSetterBase, ...] | QconfigSetterBase | QconfigSetter | None = None, method: PrepareMethod = PrepareMethod.JIT_STRIP, example_kw_inputs: Any | None = None, check_result_dir: str | None = None, *, fuse_mode: FuseMode | None = None)

Prepare model.

Prepare and check a copy of the model for QAT.

Parameters:

model (Module) – Model to be prepared.

example_inputs (Optional[Any]) – Model inputs. Used to trace and check model.

qconfig_setter (Union[Tuple[QconfigSetterBase, ...], QconfigSetterBase, QconfigSetter, None]) – Qconfig setter. Used to set qconfig.

method (PrepareMethod) –

Method used to trace model, availiable options are:

PrepareMethod.EAGER: Don’t trace.

PrepareMethod.JIT_STRIP: Use jit trace and strip the graph outside QuantStub and Dequantstub.

example_kw_inputs (Optional[Any]) – Model keyword inputs. Used to trace and check model.

check_result_dir (Optional[str]) – Directory to save qat check result txt.

fuse_mode (Optional[FuseMode]) –

Control op fusion on compute graph, availiable options are:

FuseMode.OnlyBN: Only fuse conv + bn, add and relu can be handled by qconfig template.

None: Automatically choose from FuseMode.OnlyBN and FuseMode.BNAddReLU according to current qconfig template.

FuseMode.BNAddReLU: Fuse conv + bn + add + relu.

FuseMode.NoFuse: Not do any fusion.

Return type: Module