Tensor Utils

Utils for generating stats from torch tensors.

opacus.utils.tensor_utils.calc_sample_norms(named_params, flat=True)[source]

Calculates the norm of the given tensors for each sample.

This function calculates the overall norm of the given tensors for each sample, assuming the each batch’s dim is zero.

Parameters
  • named_params (Iterator[Tuple[str, Tensor]]) – An iterator of tuples <name, param> with name being a string and param being a tensor of shape [B, ...] where B is the size of the batch and is the 0th dimension.

  • flat (bool) – A flag, when set to True returns a flat norm over all layers norms

Example

>>> t1 = torch.rand((2, 5))
>>> t2 = torch.rand((2, 5))
>>> calc_sample_norms([("1", t1), ("2", t2)])
    [tensor([1.5117, 1.0618])]
Return type

List[Tensor]

Returns

A list of tensor norms where length of the list is the number of layers

opacus.utils.tensor_utils.sum_over_all_but_batch_and_last_n(tensor, n_dims)[source]

Calculates the sum over all dimensions, except the first (batch dimension), and excluding the last n_dims.

This function will ignore the first dimension and it will not aggregate over the last n_dims dimensions.

Parameters
  • tensor (Tensor) – An input tensor of shape (B, ..., X[n_dims-1]).

  • n_dims (int) – Number of dimensions to keep.

Example

>>> tensor = torch.ones(1, 2, 3, 4, 5)
>>> sum_over_all_but_batch_and_last_n(tensor, n_dims=2).shape
torch.Size([1, 4, 5])
Return type

Tensor

Returns

A tensor of shape (B, ..., X[n_dims-1])