aitune.torch.inspecting.inspecting | NVIDIA AITune Documentation

Inspecting and patching modules.

Module Contents

Functions

Name	Description
`_reset_total_execution_time`	Reset the timings of the modules.
`_warmup`	Warmup the model by running it on a few samples.
`inspect`	Inspect provided callable object searching for nn.Module members executed as part of forward pass.

Data

DEFAULT_INSPECT_ITERATIONS

DEFAULT_WARMUP_ITERATIONS

LOG_FORMAT

logger

API

aitune.torch.inspecting.inspecting._reset_total_execution_time(
    inspector: aitune.torch.inspecting.module_inspector.ModuleInspector
)

Reset the timings of the modules.

aitune.torch.inspecting.inspecting._warmup(
    obj: collections.abc.Callable,
    dataset: aitune.torch.dataloader.DatasetLike | aitune.torch.dataloader.DataLoaderFactory | torch.Tensor,
    number_of_iterations: int
)

Warmup the model by running it on a few samples.

aitune.torch.inspecting.inspecting.inspect(
    obj: collections.abc.Callable | torch.nn.Module,
    dataset: aitune.torch.dataloader.DatasetLike | aitune.torch.dataloader.DataLoaderFactory | torch.Tensor,
    inference_function: collections.abc.Callable | None = None,
    number_of_iterations: int = DEFAULT_INSPECT_ITERATIONS,
    warmup_iterations: int = DEFAULT_WARMUP_ITERATIONS,
    min_depth: int = 0,
    max_depth: int = 5
) -> aitune.torch.inspecting.module_info.InspectedModulesInfo

Inspect provided callable object searching for nn.Module members executed as part of forward pass.

Returns: InspectedModulesInfo object.

Parameters:

obj

Callable | torch.nn.Module

Callable object to inspect.

dataset

DatasetLike | DataLoaderFactory | torch.Tensor

List of tuples with batch size and input.

inference_function

Callable | NoneDefaults to None

Custom inference function to use for inspection, obj is used by default.

number_of_iterations

intDefaults to DEFAULT_INSPECT_ITERATIONS

Number of iterations to run for inference.

warmup_iterations

intDefaults to DEFAULT_WARMUP_ITERATIONS

Number of iterations to run for warmup.

min_depth

intDefaults to 0

Minimum depth of the modules to inspect, if root level modules is not working, try to increase this value

max_depth

intDefaults to 5

Maximum depth of the modules to inspect

aitune.torch.inspecting.inspecting.DEFAULT_INSPECT_ITERATIONS = 10

aitune.torch.inspecting.inspecting.DEFAULT_WARMUP_ITERATIONS = 5

aitune.torch.inspecting.inspecting.LOG_FORMAT = '%(asctime)s - %(levelname)s - %(message)s'

aitune.torch.inspecting.inspecting.logger = getLogger(__name__)