aitune.torch.checkpoint.storage | NVIDIA AITune Documentation

Storage abstract base class.

Module Contents

Classes

Name	Description
`Storage`	Base class to save/load state dict.

API

class aitune.torch.checkpoint.storage.Storage(
    save_tasks: list[aitune.torch.checkpoint.storage_tasks.SaveTask],
    load_tasks: list[aitune.torch.checkpoint.storage_tasks.LoadTask]
)

Base class to save/load state dict.

load_tasks

= list(filter(None, load_tasks))

save_tasks

= list(filter(None, save_tasks))

aitune.torch.checkpoint.storage.Storage._get_target_folder_path(
    path: str | pathlib.Path
) -> pathlib.Path

Get the target folder path.

Parameters:

path

str | Path

The folder or a checkpoint file.

Returns: Path

The target folder path.

aitune.torch.checkpoint.storage.Storage.load(
    path: str | pathlib.Path
) -> dict

Load a state dictionary from the specified path.

Parameters:

path

str | Path

The path from where the state dictionary should be loaded- either a directory or a checkpoint file.

Returns: dict

The loaded state dictionary.

aitune.torch.checkpoint.storage.Storage.save(
    path: str | pathlib.Path,
    state_dict: dict
) -> None

Save the state dictionary to the specified path.

Parameters:

path

str | Path

The path where the state dictionary should be saved - either a directory or a checkpoint file.

state_dict

dict

The state dictionary to save.