torchelie.utils

class torchelie.utils.AutoStateDict(except_names: List[str] = [])

Inherit this class for automatic state_dict() and load_state_dict() members based on __dict__

Exclusions can be specified via except_names

class torchelie.utils.DetachedModule(m)

Wrap a module to eval model, can’t be turned back to training mode, and make it invisible to recursive calls on `nn.Module`s

Parameters

m (nn.Module) – a module

class torchelie.utils.FrozenModule(m: torch.nn.modules.module.Module)

Wrap a module to eval model, can’t be turned back to training mode

Parameters

m (nn.Module) – a module

torchelie.utils.as_multiclass_shape(preds, as_probs=False)

Manipulate an array of logit predictions so that binary prediction is not a special case anymore. Outputs preds as (Batch size, Num classes (>=2)).

if preds has one dim, another one is added. If this is binary classification, a second column is added as 1 - preds.

Parameters
  • preds (tensor) – predictions

  • as_probs (bool) – whether to return the preds as logits or probs

Returns

the predictions reshaped

torchelie.utils.bgram(m: torch.Tensor)torch.Tensor

Return the batched Gram matrix of m

Parameters

m (tensor) – a matrix of dim 3, first one is the batch

Returns

The batch of Gram matrix

torchelie.utils.constant_init(m: torch.nn.modules.module.Module, val: float)torch.nn.modules.module.Module

Initialize a module with gaussian weights of standard deviation std

Parameters

m (nn.Module) – the module to init

Returns

the initialized module

torchelie.utils.dict_by_key(d: Any, k: str)Any

Recursively index a dict by a hierarchical key

` >>> dict_by_key({'a': [{'b': 42}]}, 'a.0.b') 42 `

Parameters
  • d (dict, list, and any level of nesting) – the data to index

  • k (str) – the key

Returns

The value in d indexed by k

torchelie.utils.dist_setup(rank)

initialize a NCCL process group with default port / address. For internal use.

torchelie.utils.entropy(out: torch.Tensor, dim: int = 1, reduce: str = 'mean')torch.Tensor

Compute the entropy of the categorical distribution specified by the logits out along dimension dim.

Parameters
  • out (tensor) – logits categorial distribution

  • dim (int) – the dimension along which the distributions are specified

  • reduce (str) – “mean”, “none” or “sum”

torchelie.utils.experimental(func)

Decorator that warns about a function being experimental

torchelie.utils.fast_zero_grad(net: torch.nn.modules.module.Module)None

Set .grad to None for all parameters instead of zeroing out. It is faster.

torchelie.utils.forever(iterable: Iterable[T])Iterable[T]

Cycle through iterable forever

Parameters

iterable (iterable) – the iterable

torchelie.utils.freeze(net: T_Module)T_Module

Freeze all parameters of net

torchelie.utils.gram(m: torch.Tensor)torch.Tensor

Return the Gram matrix of m

Parameters

m (tensor) – a matrix of dim 2

Returns

The Gram matrix

torchelie.utils.ilerp(a: Numeric, b: Numeric, t: Numeric)Numeric

Inverse or lerp. For t between a and b, returns the fraction or a and b in t.

\(\frac{t - a}{b - a}\)

Parameters
  • a (number or tensor) – a

  • b (number or tensor) – b

  • t (number or tensor) – t between a and b

Returns

result

torchelie.utils.indent(text: str, amount: int = 4)str

Indent text by amount spaces.

Parameters
  • text (str) – some text

  • amount (int) – an indentation amount

Returns

indented text

torchelie.utils.kaiming(m: T_Module, a: float = 0, nonlinearity: str = 'leaky_relu', mode: str = 'fan_out', dynamic: bool = False)T_Module

Initialize a module with kaiming normal init

Parameters
  • m (nn.Module) – the module to init

  • a (float) – the slope of the nonlinearity

  • nonlinearity (str) – type of the nonlinearity

  • dynamic (bool) – wether to scale the weights on the forward pass for equalized LR such as ProGAN (default: False)

Returns

the initialized module

torchelie.utils.kaiming_gain(m: T_Module, a: float = 0, nonlinearity='leaky_relu', mode='fan_in')float

Return the std needed to initialize a weight matrix with given parameters.

torchelie.utils.layer_by_name(net: torch.nn.modules.module.Module, name: str)Optional[torch.nn.modules.module.Module]

Get a submodule at any depth of a net by its name

Parameters
  • net (nn.Module) – the base module containing other modules

  • name (str) – a name of a submodule of net, like “layer3.0.conv1”.

Returns

The found layer or None

torchelie.utils.lerp(a: float, b: float, t: float)float
torchelie.utils.lerp(a: float, b: float, t: torch.Tensor)torch.Tensor
torchelie.utils.lerp(a: torch.Tensor, b: torch.Tensor, t: torch.Tensor)torch.Tensor
torchelie.utils.lerp(a: torch.Tensor, b: torch.Tensor, t: float)torch.Tensor

Linearly interpolate between a and b according to t.

\((1 - t)a + tb\)

Parameters
  • a (number or tensor) – a

  • b (number or tensor) – b

  • t (number or tensor) – t between 0 and 1

Returns

result between a and b

torchelie.utils.load_recursive_state_dict(x: Any, obj: Any)None

Reload a state dict saved with recursive_state_dict()

Parameters
  • x – the recursive state dict

  • obj – the object that has been recursive_state_dict()ed

torchelie.utils.load_state_dict_forgiving(dst, state_dict: dict, silent: bool = False, fit_dst_size: bool = False)

Loads a state dict, but don’t crash if shapes don’t match.

torchelie.utils.nb_parameters(net: torch.nn.modules.module.Module)int

Counts the number of parameters of net

Parameters

net (nn.Module) – the net

Returns

the number of params

torchelie.utils.normal_init(m: torch.nn.modules.module.Module, std: float = 0.02)torch.nn.modules.module.Module

Initialize a module with gaussian weights of standard deviation std

Parameters

m (nn.Module) – the module to init

Returns

the initialized module

torchelie.utils.parallel_run(fun, *args, n_gpus: int = 0, **kwargs)None

Starts a function in parallel for GPU dispatching.

Parameters
  • fun (callable) – the function to start, with signature fun(*args, **kwargs, rank, world_size). rank is the GPU id tu use on this thread, and world_size is the total number of GPUs.

  • *args – arguments passed to fun

  • n_gpus (int, optional) – number of GPUs to use. Will default to the number of available GPUs.

  • **kwargs – kw arguments passed to fun

torchelie.utils.recursive_state_dict(x: Any)Any

Recursively call state_dict() on all elements contained in a list / tuple / dict so that it can be saved safely via torch.save().

Parameters

x – any nesting of list / tuple / dict containing state_dict()able objects

Returns

the same structure with state dicts

torchelie.utils.send_to_device(x: Any, device, non_blocking: bool = False)Any

Send all tensors contained in x to device, when x is an arbitrary nested datastructure of dicts and lists containing tensors

Parameters
  • x – the tensors

  • device – a torch device

  • non_blocking (bool) – non blocking

Returns

x with device changed

torchelie.utils.slerp(z1: torch.Tensor, z2: torch.Tensor, t: float)torch.Tensor

Spherical linear interpolate between z1 and z2 according to t.

Parameters
  • z1 (torch.Tensor) – ND tensor, interpolating on last dim

  • z2 (torch.Tensor) – ND tensor, interpolating on last dim

  • t (float) – t between 0 and 1

Returns

result between a and b

torchelie.utils.unfreeze(net: T_Module)T_Module

Unfreeze all parameters of net

torchelie.utils.xavier(m: T_Module, a: float = 0, nonlinearity: str = 'relu', dynamic: bool = False)T_Module

Initialize a module with xavier normal init

Parameters
  • m (nn.Module) – the module to init

  • dynamic (bool) – wether to scale the weights on the forward pass for equalized LR such as ProGAN (default: False)

Returns

the initialized module