dmx.compressor.fx.transform.cast_input_output_transform

dmx.compressor.fx.transform.cast_input_output_transform(root: ~torch.nn.modules.module.Module, tracer: ~dmx.compressor.fx.tracer.QuantTracer | ~dmx.compressor.fx.tracer.HFQuantTracer = <dmx.compressor.fx.tracer.QuantTracer object>, concrete_args: ~typing.Dict[str, ~typing.Any] | None = None, cfg: str | None = None) Module

A function that transforms the module by adding additional ops, which includes: - casting - approximator - sparsifier An optional config file can be passed to specify the formats for the additional ops, dummy formats would be used otherwise.

Parameters:
  • root (torch.nn.Module) – model/module to transform

  • tracer (Union[QuantTracer, HFQuantTracer], optional) – tracer used for tracing the root. Defaults to QuantTracer.

  • concrete_args (Dict[str,Any], optional) – concrete arguments to be used for tracer. Defaults to None.

  • cfg (Optional[str]) – config file for setting the added ops formats. Defaults to None.

Returns:

transformed model