dmx.compressor.fx.transform.substitute_transform

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