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

Return type:

List[Tensor]

Returns:

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

Example

>>> t1 = torch.rand((2, 5))
>>> t2 = torch.rand((2, 5))
>>> norms = calc_sample_norms([("1", t1), ("2", t2)])
>>> norms, norms[0].shape
([tensor([...])], torch.Size([2]))
opacus.utils.tensor_utils.calc_sample_norms_one_layer(param)[source]

Calculates the norm of the given tensor (a single parameter) for each sample.

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

It is equivalent to: calc_sample_norms(named_params=((None, param),))[0]

Parameters:

param (Tensor) – A tensor of shape [B, ...] where B is the size of the batch and is the 0th dimension.

Return type:

Tensor

Returns:

A tensor of norms

Example

>>> t1 = torch.rand((2, 5))
>>> norms = calc_sample_norms_one_layer(t1)
>>> norms, norms.shape
(tensor([...]), torch.Size([2]))
opacus.utils.tensor_utils.filter_dilated_rows(tensor, dilation, dilated_kernel_size, kernel_size)[source]

A helper function that removes extra rows created during the process of implementing dilation.

Parameters:

  • tensor (Tensor) – A tensor containing the output slices resulting from unfolding the input tensor to unfold3d(). Shape is (B, C, D_out, H_out, W_out, dilated_kernel_size[0], dilated_kernel_size[1], dilated_kernel_size[2]).

  • dilation (Tuple[int, int, int]) – The dilation given to unfold3d().

  • dilated_kernel_size (Tuple[int, int, int]) – The size of the dilated kernel.

  • kernel_size (Tuple[int, int, int]) – The size of the kernel given to unfold3d().

Returns:

A tensor of shape (B, C, D_out, H_out, W_out, kernel_size[0], kernel_size[1], kernel_size[2]) For D_out, H_out, W_out definitions see torch.nn.Unfold.

Example

>>> tensor = torch.zeros([1, 1, 3, 3, 3, 5, 5, 5])
>>> dilation = (2, 2, 2)
>>> dilated_kernel_size = (5, 5, 5)
>>> kernel_size = (3, 3, 3)
>>> filter_dilated_rows(tensor, dilation, dilated_kernel_size, kernel_size).shape
torch.Size([1, 1, 3, 3, 3, 3, 3, 3])
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.

Return type:

Tensor

Returns:

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

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])
opacus.utils.tensor_utils.unfold2d(input, *, kernel_size, padding, stride, dilation)[source]

See unfold()

opacus.utils.tensor_utils.unfold3d(tensor, *, kernel_size, padding=0, stride=1, dilation=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

Parameters:
Returns:

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

Example

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