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

opacus.utils.tensor_utils.unfold3d(tensor, kernel_size, padding=0, stride=1)[source]

Extracts sliding local blocks from an batched input tensor.

torch.nn.Unfold only supports 4D inputs (batched image-like tensors). This method implements the same action for 5D inputs

  • tensor (Tensor) – An input tensor of shape (B, C, D, H, W).

  • kernel_size (Union[int, Tuple[int, int, int]]) – the size of the sliding blocks

  • padding (Union[int, Tuple[int, int, int]]) – implicit zero padding to be added on both sides of input

  • stride (Union[int, Tuple[int, int, int]]) – the stride of the sliding blocks in the input spatial dimensions


>>> B, C, D, H, W = 3, 4, 5, 6, 7
>>> tensor = torch.arange(1,B*C*D*H*W+1.).view(B,C,D,H,W)
>>> unfold3d(tensor, kernel_size=2, padding=0, stride=1).shape
torch.Size([3, 32, 120])

A tensor of shape (B, C * np.product(kernel_size), L), where L - output spatial dimensions. See torch.nn.Unfold for more details