aitune.torch.tuning | NVIDIA AITune Documentation

Main tune function.

Module Contents

Functions

Name	Description
`_activate_tuned_modules`	Activate all tuned modules after tuning is complete.
`_clear_cache`	Clear the cache.
`_describe_module`	-
`_validate_and_normalize_batch_sizes`	Validate and normalize batch sizes for tuning.
`load`	Load the tuned module from a file.
`save`	Save the tuned module to a file.
`tune`	Tune a callable which runs inference on a pipeline or a model.

Data

LOG_FORMAT

logger

API

aitune.torch.tuning._activate_tuned_modules() -> None

Activate all tuned modules after tuning is complete.

aitune.torch.tuning._clear_cache()

Clear the cache.

aitune.torch.tuning._describe_module(
    module: aitune.torch.module.wrapper_module.Module
)

aitune.torch.tuning._validate_and_normalize_batch_sizes(
    batch_sizes: list[int] | None
) -> list[int]

Validate and normalize batch sizes for tuning.

aitune.torch.tuning.load(
    module: torch.nn.Module,
    path: str | pathlib.Path,
    storage: aitune.torch.checkpoint.storage.Storage | None = None,
    device_map: dict[str, torch.device] | None = None,
    disable_external_logging: bool = True
) -> torch.nn.Module | aitune.torch.module.wrapper_module.Module

Load the tuned module from a file.

If storage is not provided, a default storage will be used. Check save function for more details.

Parameters:

module

torch.nn.Module

The module to load.

path

str | Path

The path to load the module from.

storage

Storage | NoneDefaults to None

The storage to use to load the module. If not provided, a local storage will be used.

device_map

dict[str, torch.device] | NoneDefaults to None

The device map to load modules to. Overrides the device of stored in state dict.

disable_external_logging

boolDefaults to True

If True, libraries logging will be suppressed.

aitune.torch.tuning.save(
    module: torch.nn.Module,
    path: str | pathlib.Path,
    storage: aitune.torch.checkpoint.storage.Storage | None = None
) -> None

Save the tuned module to a file.

If storage is not provided, a default storage will be used.

If you would like to save the module to a different folder, you can use the storage parameter. For example, if you want to save the module to the ckpt folder, you can do the following:

Parameters:

module

torch.nn.Module

The module to save.

path

str | Path

The path to save the module to.

storage

Storage | NoneDefaults to None

The storage to use to save the module. If not provided, a local storage will be used.

aitune.torch.tuning.tune(
    func: collections.abc.Callable,
    dataset: aitune.torch.dataloader.DatasetLike | aitune.torch.dataloader.DataLoaderFactory | torch.Tensor,
    batch_sizes: list[int] | None = None,
    max_num_batches_per_batch_size: int | None = None,
    device: str | torch.device | None = DEFAULT_DEVICE,
    dry_run: bool = False,
    disable_external_logging: bool = False,
    clear_cache: bool = False,
    ignore_failing_modules: bool = True
) -> None

Tune a callable which runs inference on a pipeline or a model.

Parameters:

func

Callable

The function to tune.

dataset

DatasetLike | DataLoaderFactory | torch.Tensor

The dataset to tune on. It can be DataLoaderFactory or any dataset/iterable and even torch.Tensor. Tensor will be treated as a single sample dataset.

batch_sizes

list[int] | NoneDefaults to None

The batch sizes to use for tuning. At least 2 different batch sizes are required for determining batch axis.

max_num_batches_per_batch_size

int | NoneDefaults to None

The maximum number of batches to use for tuning per batch size.

device

str | torch.device | NoneDefaults to DEFAULT_DEVICE

The device to use for tuning.

dry_run

boolDefaults to False

If True, only dry run the tuning.

disable_external_logging

boolDefaults to False

If True, libraries logging will be suppressed.

clear_cache

boolDefaults to False

If True, the cache will be cleared before tuning.

ignore_failing_modules

boolDefaults to True

If True, failing modules will be ignored and tuning will continue.

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

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