dmx.compressor.fx.transform

Functions

`cast_input_output_transform`(root[, tracer, ...])	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.
`configure_transform`(gm, scopeDict, cfg)	A function that changes the format of the ops according to the cfg file .. note::.
`export_substitute_transform`(root, kwargs[, ...])
`make_compiler_graph`(root)	A function that transforms the model by substituting CastTos with Q/dQ ops
`prepare_dynamic_shapes`(model, kwargs)
`prepare_module_call_signature`(model)
`prepare_tracing_inputs`(_model, args, kwargs)
`qDq_transform`(root)	A function that transforms the model by substituting CastTos with Q/dQ ops
`substitute_transform`(root[, concrete_args, ...])	A function that transforms the model by substituting torch.nn.modules and activation functions to dmx.nn.modules.

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

dmx.compressor.fx.transform.configure_transform(gm: GraphModule, scopeDict: dict, cfg: str)

A function that changes the format of the ops according to the cfg file .. note:

Configure_transform will only change existing ops and will not add any additional ops.
Hence it is recommened to pass in a cfg file for cast_input_output_transform to make sure all necessary ops are added.

Parameters:

gm (torch.fx.GraphModule) – Graphmodule to apply changes on
scopeDict (dict) – Dictionary that maps node name to scope
cfg (str) – config file for setting the added ops formats.

Returns:

Graphmodule with updated formats

dmx.compressor.fx.transform.export_substitute_transform(root: Module, kwargs, additional_mappings=None)

dmx.compressor.fx.transform.make_compiler_graph(root: GraphModule)

A function that transforms the model by substituting CastTos with Q/dQ ops

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

dmx.compressor.fx.transform.prepare_dynamic_shapes(model, kwargs) → dict

dmx.compressor.fx.transform.prepare_module_call_signature(model: Module) → tuple

dmx.compressor.fx.transform.prepare_tracing_inputs(_model, args, kwargs)

dmx.compressor.fx.transform.qDq_transform(root: GraphModule)

A function that transforms the model by substituting CastTos with Q/dQ ops

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

dmx.compressor.fx.transform.substitute_transform(root: Module, concrete_args: Dict[str, Any] | None = None, input_names: List[str] | None = None, dummy_inputs: Dict[str, Any] | None = None, additional_mappings: Dict[str, Any] | None = None)

A function that transforms the model by substituting torch.nn.modules and activation functions to dmx.nn.modules.

Parameters:

root (torch.nn.Module) – model/module to transform
concrete_args (Dict[str,Any], optional) – concrete arguments to be used for tracer. Defaults to None.
hf (bool, optional) – True if root is a huggingface model. Defaults to False

Returns:

transformed model