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.

  • 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


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



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.

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

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


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



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