Skip to content

Utils#

Some useful utility functions.

thunder.torch.utils #

get_device(x) #

Infer device of torch.Tensor or nn.Module instance. Parameters

x: Union[torch.Tensor, nn.Module] Returns

device: torch.device

Source code in thunder/torch/utils.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def get_device(x: Union[torch.Tensor, nn.Module]) -> torch.device:
    """
    Infer device of torch.Tensor or nn.Module instance.
    Parameters
    ----------
    x: Union[torch.Tensor, nn.Module]
    Returns
    -------
    device: torch.device
    """
    if isinstance(x, (torch.Tensor, LightningModule)):
        return x.device
    elif isinstance(x, nn.Module):
        try:
            return next(x.parameters()).device
        except StopIteration as e:
            raise RuntimeError("Can't infer the device, because the module has no parameters") from e

    raise TypeError(f"Can't infer the device of {type(x)}")

last_checkpoint(root) #

Load most fresh last.ckpt file based on time. Parameters

root: Union[Path, str] Path to folder, where last.ckpt or its symbolic link supposed to be. Returns

checkpoint_path: Union[Path, str] If last.ckpt exists - returns Path to it. Otherwise, returns 'last'.

Source code in thunder/torch/utils.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def last_checkpoint(root: Union[Path, str]) -> Union[Path, Literal["last"]]:
    """
    Load most fresh last.ckpt file based on time.
    Parameters
    ----------
    root: Union[Path, str]
        Path to folder, where last.ckpt or its symbolic link supposed to be.
    Returns
    -------
    checkpoint_path: Union[Path, str]
        If last.ckpt exists - returns Path to it. Otherwise, returns 'last'.
    """
    checkpoints = []
    for p in Path(root).rglob("*"):
        if p.is_symlink():
            p = p.resolve(strict=False)
        if p.suffix == ".ckpt":
            checkpoints.append(p)

    return max(checkpoints, key=lambda t: os.stat(t).st_mtime, default="last")

maybe_from_np(*x, device='cpu') #

Recursively converts numpy arrays to torch.Tensor. Parameters

*x: Any device: Union[torch.device, str] Device to move to, default is CPU. Returns

Collection of tensors. Examples

x, y # np.ndarray z = maybe_from_np(x) # convert to torch.Tensor x, y = maybe_from_np(x, y) # x and y are now tensors x, y, z = maybe_from_np(x, y, z) # maybe_from_np converts np arrays and tensors and does not affect other types dict_of_tensors = to_np(dict_of_np) # maybe_from_np converts any collection

Source code in thunder/torch/utils.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def maybe_from_np(*x: Any, device: Union[torch.device, str] = "cpu") -> Any:
    """
    Recursively converts numpy arrays to torch.Tensor.
    Parameters
    ----------
    *x: Any
    device: Union[torch.device, str]
        Device to move to, default is CPU.
    Returns
    -------
    Collection of tensors.
    Examples
    -------
    >>> x, y  # np.ndarray
    >>> z = maybe_from_np(x) # convert to torch.Tensor
    >>> x, y = maybe_from_np(x, y) # x and y are now tensors
    >>> x, y, z = maybe_from_np(x, y, z) # maybe_from_np converts np arrays and tensors and does not affect other types
    >>> dict_of_tensors = to_np(dict_of_np) # maybe_from_np converts any collection
    """

    def to_tensor(x):
        if isinstance(x, torch.Tensor):
            return x.to(device)
        return torch.from_numpy(x).to(device)

    return squeeze_first(apply_to_collection(x, (np.ndarray, np.generic, torch.Tensor), to_tensor))

tensor2np(x) #

Detaches, moves torch.Tensor to CPU and converts into numpy array. Parameters

x: torch.Tensor Returns

np.ndarray

Source code in thunder/torch/utils.py 
55
56
57
58
59
60
61
62
63
64
65
def tensor2np(x: torch.Tensor) -> np.ndarray:
    """
    Detaches, moves torch.Tensor to CPU and converts into numpy array.
    Parameters
    ----------
    x: torch.Tensor
    Returns
    -------
    np.ndarray
    """
    return x.detach().cpu().numpy()

to_np(*x) #

Converts collection of tensors into numpy arrays. Parameters

*x: Any Returns

Collection of numpy arrays Examples

x, y # torch.Tensor z = to_np(x) # convert to numpy array x, y = to_np(x, y) # x and y are now numpy arrays x, y, z = to_np(x, y, z) # to_np converts only tensors and does not affect other types dict_of_np = to_np(dict_of_tensors) # to_np converts any collection

Source code in thunder/torch/utils.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def to_np(*x: Any) -> Any:
    """
    Converts collection of tensors into numpy arrays.
    Parameters
    ----------
    *x: Any
    Returns
    ----------
    Collection of numpy arrays
    Examples
    --------
    >>> x, y  # torch.Tensor
    >>> z = to_np(x) # convert to numpy array
    >>> x, y = to_np(x, y) # x and y are now numpy arrays
    >>> x, y, z = to_np(x, y, z) # to_np converts only tensors and does not affect other types
    >>> dict_of_np = to_np(dict_of_tensors) # to_np converts any collection
    """
    return squeeze_first(apply_to_collection(x, torch.Tensor, tensor2np))