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
Example
>>> t1 = torch.rand((2, 5)) >>> t2 = torch.rand((2, 5)) >>> calc_sample_norms([("1", t1), ("2", t2)]) [tensor([1.5117, 1.0618])]
-
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
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
- Returns
A tensor of shape
(B, ..., X[n_dims-1])
-
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
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 blockspadding (
Union
[int
,Tuple
[int
,int
,int
]]) – implicit zero padding to be added on both sides of inputstride (
Union
[int
,Tuple
[int
,int
,int
]]) – the stride of the sliding blocks in the input spatial dimensionsdilation (
Union
[int
,Tuple
[int
,int
,int
]]) – the spacing between the kernel points.
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])
- Returns
A tensor of shape
(B, C * np.product(kernel_size), L)
, where L - output spatial dimensions. Seetorch.nn.Unfold
for more details