aitune.torch.module.tuned_module | NVIDIA AITune Documentation

Tuned module.

Module Contents

Classes

Name	Description
`TunedModule`	Module that does inference using tuned backend.

API

class aitune.torch.module.tuned_module.TunedModule(
    backends: collections.OrderedDict[aitune.torch.module.sample_metadata.SampleMetadata, aitune.torch.backend.backend.Backend],
    module_name: str,
    forward_signature: aitune.torch.module.forward_signature.ForwardSignature,
    check_graph: bool = True,
    config: aitune.torch.config.AITuneConfig | None = None
)

Module that does inference using tuned backend.

BACKENDS_KEY

= 'backends'

CHECK_GRAPH_KEY

= 'check_graph'

ERROR_NO_BACKENDS

= 'No backends provided'

ERROR_NO_BACKEND_FOUND

= 'No backend found for a graph with metadata: {}'

ERROR_UNKNOWN_BACKEND

= 'Unknown backend type: {}'

FORWARD_SIGNATURE_KEY

= 'forward_signature'

IS_ANY_JIT_KEY

= 'is_any_jit'

MODULE_NAME_KEY

= 'module_name'

TYPE_KEY

= 'type'

_backend_func

= self.safe_call_backend

_config

= config if config is not None else global_config

_unique_backend

= list(self._backends.values())[0]

backends

Get the backends of the module.

device

Get the device of the module.

aitune.torch.module.tuned_module.TunedModule.__call__(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Run the call through the tuned module.

aitune.torch.module.tuned_module.TunedModule.activate()

Activates the module backends.

aitune.torch.module.tuned_module.TunedModule.call_unique_backend(
    sample: aitune.torch.module.recording_module.Sample
)

Calls the only backend without checking metadata.

aitune.torch.module.tuned_module.TunedModule.deactivate()

Deactivates the module backends.

aitune.torch.module.tuned_module.TunedModule.deploy(
    device: torch.device | None
)

Deploys the module backends.

aitune.torch.module.tuned_module.TunedModule.from_dict(
    module: torch.nn.Module,
    state_dict: dict
)

staticmethod

Creates a TunedModule from a state_dict.

aitune.torch.module.tuned_module.TunedModule.safe_call_backend(
    sample: aitune.torch.module.recording_module.Sample
)

Calls the backend according to the metadata of the sample.

aitune.torch.module.tuned_module.TunedModule.to_dict()

Returns the state_dict of the module.