quantization.prepare

horizon_plugin_pytorch.quantization.prepare(model: Module, example_inputs: Any = None, qconfig_setter: Tuple[QconfigSetterBase, ...] | QconfigSetterBase | None = None, method: PrepareMethod = PrepareMethod.JIT_STRIP)

Prepare model.

Prepare and check a copy of the model for QAT.

  • Parameters:
    • model – Model to be prepared.
    • example_inputs – Model inputs. Used to trace and check model.
    • qconfig_setter – Qconfig setter. Used to set qconfig.
    • method – Method used to trace model, availiable options are: ‘eager’: Don’t trace. ‘symbolic’: Use symbolic trace. ‘jit’: Use jit trace. ‘jit-strip’: Use jit trace and strip the graph outside QuantStub and Dequantstub. # noqa: E501