ablator.utils package

Submodules

ablator.utils.base module

class ablator.utils.base.CUDA_PROCESS(process_name, pid, memory)

Bases: tuple

memory

Alias for field number 2

pid

Alias for field number 1

process_name

Alias for field number 0

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

Bases: object

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.

class ablator.utils.base.c_nvmlProcessInfo_t

Bases: _PrintableStructure

computeInstanceId

Structure/Union member

gpuInstanceId

Structure/Union member

pad

Structure/Union member

pid

Structure/Union member

usedGpuMemory

Structure/Union member

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_cuda_processes() → dict[int, list[ablator.utils.base.CUDA_PROCESS]]

Finds the currently running cuda processes on the system. Each process is a CUDA_PROCESS object that contains information on the process name, pid and the memory utilization.

Returns:

dict[int, list[CUDA_PROCESS]]

The key of each dictionary is the device-id, corresponding to a list of running CUDA processes.

ablator.utils.base.get_gpu_mem(mem_type: Literal['used', 'total', 'free'] = 'total') → dict[int, 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:

dict[int, 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.is_oom_exception(err: RuntimeError) → bool

is_oom_exception checks whether the exception is caused by CUDA out of memory errors.

Parameters:

errRuntimeError

the exception to parse

Returns:

bool

whether the exception indicates out of memory error.

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.num_format(value: str | int | float | integer | floating, width: int = 8) → str

Format number to be no larger than width by converting to scientific notation when the value exceeds width either by informative decimal places or size.

Parameters:

valueint | float

the value to format

widthint, optional

the width of the decimal places, by default 5

Returns:

str

The formatted string representation of the value

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.

ablator.utils.file 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:

dict_dict

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:

json_str

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:

dict_dict

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”.

ablator.utils.file.truncate_utf8_chars(filename: Path, last_char: str)

ablator.utils.progress_bar module

class ablator.utils.progress_bar.Display

Bases: object

close()

print_texts(texts)

class ablator.utils.progress_bar.ProgressBar(total_steps, epoch_len: int | None = None, logfile: ~pathlib.Path | None = None, update_interval: int = 1, remote_display: <ablator.utils.progress_bar.ActorClass(RemoteProgressBar) object at 0x7fabfd2b87f0> | None = None, uid: str | None = None)

Bases: object

close()

classmethod make_bar(current_iteration: int, start_time: float, epoch_len: int, total_steps: int, ncols: int | None = None)

classmethod make_metrics_message(metrics: dict[str, Any], nrows: int | None = None, ncols: int | None = None)

make_print_message()

property ncols

property nrows

reset() → None

update_metrics(metrics: dict[str, Any], current_iteration: int)

class ablator.utils.progress_bar.RemoteDisplay(remote_progress_bar: <ablator.utils.progress_bar.ActorClass(RemoteProgressBar) object at 0x7fabfd2b87f0>, update_interval: int = 1)

Bases: Display

refresh(force=False)

ablator.utils.progress_bar.get_last_line(filename: Path)

ablator.utils.progress_bar.in_notebook()

Module contents