fastdev.utils

Submodules

Package Contents

fastdev.utils.cuda_toolkit_available() → bool[source]

Check if the nvcc is avaiable on the machine.

Return type:

bool

fastdev.utils.current_cuda_arch() → str[source]

Get the current CUDA architecture.

Return type:

str

fastdev.utils.summarize_model(model: torch.nn.Module, max_depth: int = 1)[source]
Parameters:

  • model (torch.nn.Module)

  • max_depth (int)

class fastdev.utils.cuda_timeit(print_tmpl: str | None = None)[source]

Bases: timeit

Measure the time of a block of code that may involve CUDA operations. We use CUDA events and synchronization for the accurate measurements.

Parameters:

print_tmpl (str, optional) – The template to print the time. Defaults to None. Can be a string with a placeholder for the time, e.g., “func foo costs {:.5f} s” or a string without a placeholder, e.g., “func foo”.

__enter__()[source]
__exit__(exec_type, exec_value, traceback)[source]
class fastdev.utils.timeit(fn_or_print_tmpl: Callable | str | None = None)[source]

Measure the time of a block of code.

Parameters:

  • print_tmpl (str, optional) – The template to print the time. Defaults to None. Can be a string with a placeholder for the time, e.g., “func foo costs {:.5f} s” or a string without a placeholder, e.g., “func foo”.

  • fn_or_print_tmpl (Optional[Union[Callable, str]])

Examples

>>> with timeit():
...     time.sleep(1)
it costs 1.00000 s
>>> @timeit
... def foo():
...     time.sleep(1)
foo costs 1.00000 s
>>> @timeit("func foo")
... def foo():
...     time.sleep(1)
func foo costs 1.00000 s
__enter__()[source]
__exit__(exec_type, exec_value, traceback)[source]
__call__(func: T) → T[source]
Parameters:

func (T)

Return type:

T

fastdev.utils.seed_everything(seed: int, deterministic: bool = False)[source]

Seed all random number generators.

Parameters:

  • seed (int) – Seed to be used.

  • deterministic (bool) – Whether to set the deterministic option for CUDNN backend, i.e., set torch.backends.cudnn.deterministic to True and torch.backends.cudnn.benchmark to False. Default: False.

fastdev.utils.list_to_packed(x: List[torch.Tensor]) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor][source]

Transforms a list of N tensors each of shape (Mi, K, …) into a single tensor of shape (sum(Mi), K, …).

Parameters:

x (List[torch.Tensor]) – list of tensors.

Returns:

4-element tuple containing

  • x_packed: tensor consisting of packed input tensors along the 1st dimension.

  • num_items: tensor of shape N containing Mi for each element in x.

  • item_packed_first_idx: tensor of shape N indicating the index of the first item belonging to the same element in the original list.

  • item_packed_to_list_idx: tensor of shape sum(Mi) containing the index of the element in the list the item belongs to.

Return type:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

fastdev.utils.list_to_padded(x: List[torch.Tensor] | Tuple[torch.Tensor], pad_size: Sequence[int] | None = None, pad_value: float | int = 0.0, equisized: bool = False) → torch.Tensor[source]

Transforms a list of N tensors each of shape (Si_0, Si_1, … Si_D) into: - a single tensor of shape (N, pad_size(0), pad_size(1), …, pad_size(D))

if pad_size is provided

  • or a tensor of shape (N, max(Si_0), max(Si_1), …, max(Si_D)) if pad_size is None.

Parameters:

  • x (Union[List[torch.Tensor], Tuple[torch.Tensor]]) – list of Tensors

  • pad_size (Union[Sequence[int], None]) – list(int) specifying the size of the padded tensor. If None (default), the largest size of each dimension is set as the pad_size.

  • pad_value (Union[float, int]) – float value to be used to fill the padded tensor

  • equisized (bool) – bool indicating whether the items in x are of equal size (sometimes this is known and if provided saves computation)

Returns:

tensor consisting of padded input tensors stored

over the newly allocated memory.

Return type:

x_padded

fastdev.utils.packed_to_list(x: torch.Tensor, split_size: Sequence[int] | int) → List[torch.Tensor][source]

Transforms a tensor of shape (sum(Mi), K, L, …) to N set of tensors of shape (Mi, K, L, …) where Mi’s are defined in split_size

Parameters:

  • x (torch.Tensor) – tensor

  • split_size (Union[Sequence[int], int]) – list, tuple or int defining the number of items for each tensor in the output list.

Returns:

A list of Tensors

Return type:

x_list

fastdev.utils.padded_to_list(x: torch.Tensor, split_size: Sequence[int] | None = None, dim: int = 0) → List[torch.Tensor][source]

Transforms a padded tensor of shape (N, S_1, S_2, …, S_D) into a list of N tensors of shape: - (Si_1, Si_2, …, Si_D) where (Si_1, Si_2, …, Si_D) is specified in split_size(i) - or (S_1, S_2, …, S_D) if split_size is None - or (Si_1, S_2, …, S_D) if split_size(i) is an integer.

Parameters:

  • x (torch.Tensor) – tensor

  • split_size (Union[Sequence[int], None]) – optional 1D list/tuple of ints defining the number of items for each tensor.

  • dim (int)

Returns:

a list of tensors sharing the memory with the input.

Return type:

x_list

fastdev.utils.padded_to_packed(x: torch.Tensor, split_size: list | tuple, dim: int = 0)[source]

Transforms a padded tensor of shape (…, N, M, …) into a packed tensor of shape: - (…, sum(split_size), …) if split_size is provided - (…, N * M, …) if split_size is None

Parameters:

  • x (torch.Tensor) – tensor of shape (…, N, M, …)

  • split_size (Union[list, tuple]) – list, tuple defining the number of items for each tensor in the output list.

  • dim (int) – the N dimension in the input tensor

Returns:

a packed tensor

Return type:

x_packed

fastdev.utils.atleast_nd(tensor: None, expected_ndim: int) → None[source]
fastdev.utils.atleast_nd(tensor: numpy.ndarray, expected_ndim: int) → numpy.ndarray
fastdev.utils.atleast_nd(tensor: torch.Tensor, expected_ndim: int) → torch.Tensor

Convert input to at least nD tensor.

Note

Differs from np.atleast_nd and torch.atleast_nd, this function can add dimensions to the front or back of the tensor.

fastdev.utils.auto_cast(fn: Callable | None = None, return_type: Literal['by_input', 'by_func', 'pt', 'np'] = 'by_input') → Callable[source]

Automatically cast input and output of a function to numpy or torch tensors. Since the function simply converts the input and output to numpy or torch tensors, it may introduce overhead. It is recommended to use this function for functions that are not performance critical.

Parameters:

  • fn (Callable) – Function to be wrapped.

  • return_type (Literal["by_input", "by_func", "pt", "np"], optional) – Type of return value. Defaults to “by_input”. - “by_input”: Return type is determined by the input argument type, first found array/tensor type is used. - “by_func”: Return type is determined by the orginal function. - “pt”: Return type is torch.Tensor. - “np”: Return type is np.ndarray.

Returns:

Wrapped function.

Return type:

Callable

fastdev.utils.to_number(x: None) → None[source]
fastdev.utils.to_number(x: int) → int
fastdev.utils.to_number(x: float) → float
fastdev.utils.to_number(x: numpy.ndarray) → int | float
fastdev.utils.to_number(x: torch.Tensor) → int | float

Convert input to number.

Parameters:

x (Any) – Input to be converted.

fastdev.utils.to_numpy(x: torch.Tensor, preserve_list: bool = ...) → numpy.ndarray[source]
fastdev.utils.to_numpy(x: numpy.ndarray, preserve_list: bool = ...) → numpy.ndarray
fastdev.utils.to_numpy(x: numpy.typing.ArrayLike, preserve_list: bool = ...) → numpy.ndarray
fastdev.utils.to_numpy(x: None, preserve_list: bool = ...) → None
fastdev.utils.to_numpy(x: Dict[Any, Any], preserve_list: bool = ...) → Dict[Any, numpy.ndarray]
fastdev.utils.to_numpy(x: List[Any], preserve_list: Literal[True]) → List[numpy.ndarray]
fastdev.utils.to_numpy(x: List[Any], preserve_list: Literal[False] = ...) → numpy.ndarray
fastdev.utils.to_numpy(x: Tuple[Any, Ellipsis], preserve_list: Literal[True]) → Tuple[numpy.ndarray, Ellipsis]
fastdev.utils.to_numpy(x: Tuple[Any, Ellipsis], preserve_list: Literal[False] = ...) → numpy.ndarray

Convert input to numpy array.

Parameters:

  • x (Any) – Input to be converted.

  • preserve_list (bool, optional) – Whether to preserve list or convert to numpy array. Defaults to True.

fastdev.utils.to_torch(x: numpy.ndarray, preserve_list: bool = ...) → torch.Tensor[source]
fastdev.utils.to_torch(x: torch.Tensor, preserve_list: bool = ...) → torch.Tensor
fastdev.utils.to_torch(x: None, preserve_list: bool = ...) → None
fastdev.utils.to_torch(x: Dict[Any, Any], preserve_list: bool = ...) → Dict[Any, torch.Tensor]
fastdev.utils.to_torch(x: List[Any], preserve_list: Literal[True]) → List[torch.Tensor]
fastdev.utils.to_torch(x: List[Any], preserve_list: Literal[False] = ...) → torch.Tensor
fastdev.utils.to_torch(x: Tuple[Any, Ellipsis], preserve_list: Literal[True]) → Tuple[torch.Tensor, Ellipsis]
fastdev.utils.to_torch(x: Tuple[Any, Ellipsis], preserve_list: Literal[False] = ...) → torch.Tensor

Convert input to torch tensor.

Parameters:

  • x (Any) – Input to be converted.

  • preserve_list (bool, optional) – Whether to preserve list or convert to torch tensor. Defaults to True.

fastdev.utils.log_once(message: str, level: str | int = logging.INFO)[source]

Log a message only once (based on the message content and the source code location).

Parameters:

  • message (str) – message to log

  • level (str or int) – log level, could be “critical”, “error”, “warning”, “info”, “debug” or corresponding int value (default: “info”)

fastdev.utils.parallel_track(func: Callable[[T], R], args: List[T], num_workers: int = 8, description: str = 'Processing') → List[R][source]
Parameters:

  • func (Callable[[T], R])

  • args (List[T])

  • num_workers (int)

  • description (str)

Return type:

List[R]