aitune.torch.inspecting.module_info | NVIDIA AITune Documentation

Module Info for inspected modules.

Module Contents

Classes

Name	Description
`DictOfModulesInfo`	Information about a dictionary of PyTorch modules.
`InspectedModulesInfo`	Information about inspected modules.
`ListOfModulesInfo`	Information about a list of PyTorch modules.
`ModuleInfo`	Information about a PyTorch module.
`ObjectOfModulesInfo`	Information about a custom object containing PyTorch modules.

API

class aitune.torch.inspecting.module_info.DictOfModulesInfo()

Bases: ModuleInfo

Information about a dictionary of PyTorch modules.

aitune.torch.inspecting.module_info.DictOfModulesInfo.set_wrapped(
    name: str,
    wrapped: typing.Any
)

Set the wrapped module in a dict under name as key.

class aitune.torch.inspecting.module_info.InspectedModulesInfo(
    total_execution_time: float,
    number_of_batches: int
)

Information about inspected modules.

_modules

dict[str, ModuleInfo] = {}

aitune.torch.inspecting.module_info.InspectedModulesInfo.add_module(
    module: aitune.torch.inspecting.module_info.ModuleInfo
)

Add a module to the specification.

Parameters:

module

ModuleInfo

ModuleInfo object.

aitune.torch.inspecting.module_info.InspectedModulesInfo.describe() -> None

Describe the inspected modules specification.

aitune.torch.inspecting.module_info.InspectedModulesInfo.get_modules(
    min_execution_percentage: float | None = None,
    limit: int | None = None
) -> list[aitune.torch.inspecting.module_info.ModuleInfo]

Get the list of modules.

Parameters:

min_execution_percentage

float | NoneDefaults to None

Minimum execution percentage to include a module.

limit

int | NoneDefaults to None

Maximum number of modules to return.

Returns: list[ModuleInfo]

List of ModuleInfo objects.

class aitune.torch.inspecting.module_info.ListOfModulesInfo()

Bases: ModuleInfo

Information about a list of PyTorch modules.

aitune.torch.inspecting.module_info.ListOfModulesInfo.set_wrapped(
    name: str,
    wrapped: typing.Any
)

Set the wrapped module in a list under name as index.

class aitune.torch.inspecting.module_info.ModuleInfo(
    module: torch.nn.Module,
    name: str = '',
    parent: typing.Optional[aitune.torch.inspecting.module_info.ModuleInfo] = None,
    forward_called: bool = False,
    execution_count: int = 0,
    total_execution_time: float = 0.0,
    input_types: list[dict[str, typing.Any]] = list(),
    output_types: list[dict[str, typing.Any]] = list(),
    object_path: str | None = None,
    depth: int = 0
)

Dataclass

Information about a PyTorch module.

average_execution_time

float

Get average execution time of the forward function.

depth

int = 0

execution_count

int = 0

forward_called

bool = False

input_types

list[dict[str, Any]] = field(default_factory=list)

module

Module

module_type

type

Get the type of the module.

name

str = ''

num_layers

int

Get the number of layers in the module.

num_parameters

int

Get the number of parameters in the module.

object_path

str | None = None

output_types

list[dict[str, Any]] = field(default_factory=list)

parent

Optional[ModuleInfo] = None

precisions

set

Get the layers precisions of the module and return a set of unique precisions.

total_execution_time

float = 0.0

aitune.torch.inspecting.module_info.ModuleInfo.set_wrapped(
    name: str,
    wrapped: typing.Any
)

Set the wrapped module.

class aitune.torch.inspecting.module_info.ObjectOfModulesInfo()

Bases: ModuleInfo

Information about a custom object containing PyTorch modules.

aitune.torch.inspecting.module_info.ObjectOfModulesInfo.set_wrapped(
    name: str,
    wrapped: typing.Any
)

Set the wrapped module.