Utils package

Submodules

Base Utils module

class ablator.utils.base.Dummy(*args, **kwargs)

Bases: FileLogger

__init__(*args, **kwargs)

Initialize a FileLogger.

Parameters:

pathstr | Path | None, optional

Path to the log file, by default None.

verbosebool, optional

Whether to print messages to the console, by default True.

prefixstr | None, optional

A prefix to add to each logged message, by default None.

ablator.utils.base.apply_lambda_to_iter(iterable, fn: Callable)

Applies a given function fn to each element of an iterable data structure.

This function recursively applies fn to elements within nested dictionaries or lists. It can be used for converting torch.Tensor elements to NumPy arrays or moving tensors to a specified device.

Parameters:

iterableIterable

The input iterable.

fnCallable

The function to apply to each element.

Returns:

any

The resulting data structure after applying fn to each element of the input iterable. The type of the returned object matches the type of the input iterable.

ablator.utils.base.debugger_is_active() → bool

Check if the debugger is currently active.

Returns:

bool

True if the debugger is active, False otherwise.

Notes

Return if the debugger is currently active

ablator.utils.base.get_gpu_max_mem() → List[int]

Get the maximum memory of all available GPUs.

Returns:

ty.List[int]

A list of the maximum memory for each GPU.

ablator.utils.base.get_gpu_mem(mem_type: Literal['used', 'total', 'free'] = 'total') → List[int]

Get the memory information of all available GPUs.

Parameters:

mem_typety.Literal[“used”, “total”, “free”], optional

The type of memory information to retrieve, by default “total”.

Returns:

ty.List[int]

A list of memory values for each GPU, depending on the specified memory type.

ablator.utils.base.get_latest_chkpts(checkpoint_dir: Path) → list[pathlib.Path]

Get a list of all checkpoint files in a directory, sorted from the latest to the earliest.

Parameters:

checkpoint_dirPath

The directory containing checkpoint files.

Returns:

list[Path]

A list of the checkpoint files sorted by filename.

ablator.utils.base.get_lr(optimizer)

Get the learning rate from an optimizer.

Parameters:

optimizertorch.optim.Optimizer or dict

The optimizer.

Returns:

float

The learning rate.

ablator.utils.base.init_weights(module: Module)

Initialize the weights of a module.

Parameters:

modulenn.Module

The input module to initialize.

Notes

• If the module is a Linear layer, initialize weight values from a normal distribution N(mu=0, std=1.0). If biases are available, initialize them to zeros.

• If the module is an Embedding layer, initialize embeddings with values from N(mu=0, std=1.0). If padding is enabled, set the padding embedding to a zero vector.

• If the module is a LayerNorm layer, set all biases to zeros and all weights to 1.

ablator.utils.base.iter_to_device(data_dict, device) → Sequence[Tensor] | dict[str, torch.Tensor]

Moving torch.Tensor elements to the specified device.

Parameters:

data_dictdict or list

The input dictionary or list containing torch.Tensor elements.

devicetorch.device | str

The target device for the tensors.

Returns:

ty.Union[Sequence[torch.Tensor], dict[str, torch.Tensor]]

The input data with tensors moved to the target device.

ablator.utils.base.iter_to_numpy(iterable)

Convert elements of the input iterable to NumPy arrays if they are torch.Tensor objects.

Parameters:

iterableIterable

The input iterable.

Returns:

any

The iterable with torch.Tensor elements replaced with their NumPy array equivalents.

ablator.utils.base.parse_device(device: str | list[str])

Parse a device string, an integer, or a list of device strings or integers.

Parameters:

devicety.Union[str, list[str], int]

The target device for the tensors.

Returns:

any

The parsed device string, integer, or list of device strings or integers.

Raises:

ValueError

If the device string is not one of {‘cpu’, ‘cuda’} or doesn’t start with ‘cuda:’.

AssertionError

If cuda is not found on system or gpu number of device is not available.

Examples

>>> parse_device("cpu")
'cpu'
>>> parse_device("cuda")
'cuda'
>>> parse_device("cuda:0")
'cuda:0'
>>> parse_device(["cpu", "cuda"])
['cpu', 'cuda']
>>> parse_device(["cpu", "cuda:0", "cuda:1", "cuda:2"])
['cpu', 'cuda:0', 'cuda:1', 'cuda:2']

ablator.utils.base.set_seed(seed: int)

Set the random seed.

Parameters:

seedint

The random seed to set.

Returns:

int

The set random seed.

File Utils module

ablator.utils.file.clean_checkpoints(checkpoint_folder: Path, n_checkpoints: int)

Remove all but the n latest checkpoints from the given directory.

Parameters:

checkpoint_folderPath

Directory containing the checkpoint files.

n_checkpointsint

Number of checkpoints to keep.

ablator.utils.file.default_val_parser(val)

Converts the input value to a JSON compatible format.

Parameters:

valty.Any

The value to be converted.

Returns:

ty.Any

The converted value.

ablator.utils.file.dict_to_json(_dict)

Convert a dictionary into a JSON string.

Parameters:

_dictdict

The dictionary to be converted.

Returns:

str

The JSON string representation of the dictionary.

ablator.utils.file.json_to_dict(_json)

Convert a JSON string into a dictionary.

Parameters:

_jsonstr

JSON string to be converted.

Returns:

dict

A dictionary representation of the JSON string.

ablator.utils.file.make_sub_dirs(parent: str | Path, *dir_names) → list[pathlib.Path]

Create subdirectories under the given parent directory.

Parameters:

parentstr | Path

Parent directory where subdirectories should be created.

*dir_namesstr

Names of the subdirectories to create.

Returns:

list[Path]

A list of created subdirectory paths.

ablator.utils.file.nested_set(_dict, keys: list[str], value: Any)

Set a value in a nested dictionary.

Parameters:

_dictdict

The dictionary to update.

keyslist[str]

List of keys representing the nested path.

valuety.Any

The value need to set at the specified path.

Returns:

dict

The updated dictionary with the new value set.

Examples

>>> _dict = {'a': {'b': {'c': 1}}}
>>> nested_set(_dict, ['a', 'b', 'c'], 2)
>>> _dict
{'a': {'b': {'c': 2}}}

ablator.utils.file.save_checkpoint(state, filename='checkpoint.pt')

Save a checkpoint of the given state.

Parameters:

statedict

Model State dictionary to save.

filenamestr, optional

The name of the checkpoint file, by default “checkpoint.pt”.

Module contents