pypose.optim.functional.modjac

class pypose.optim.functional.modjac(model, input=None, create_graph=False, strict=False, vectorize=False, strategy='reverse-mode', flatten=False)

Compute the model Jacobian with respect to the model parameters.

For a parametric model \(\bm{f}(\bm{\theta}, \bm{x})\), where \(\bm{\theta}\) is the learnable parameter and \(\bm{x}\) is the input, it computes the Jacobian of the \(i\)-th output and \(j\)-th parameter as

\[{\displaystyle \mathbf{J}_{i,j} = {\begin{bmatrix} {\dfrac {\partial \bm{f}_{i,1}}{\partial \bm{\theta}_{j,1}}} & \cdots&{\dfrac {\partial \bm{f}_{i,1}}{\partial \bm{\theta}_{j,n}}}\\ \vdots &\ddots &\vdots \\ {\dfrac {\partial \bm{f}_{i,m}}{\partial \bm{\theta}_{j,1}}}& \cdots &{\dfrac {\partial \bm{f}_{i,m}}{\partial \bm{\theta}_{j,n}}} \end{bmatrix}}}. \]

Parameters:

• model (torch.nn.Module) – a PyTorch model that takes Tensor or LieTensor input and returns a tuple of Tensors/LieTensors or a Tensor/LieTensor.

• input (tuple of Tensors/LieTensors or Tensor/LieTensor) – input to the model. Defaults to None.

• create_graph (bool, optional) – If True, the Jacobian will be computed in a differentiable manner. Note that when strict is False, the result can not require gradients or be disconnected from the input. Defaults to False.

• strict (bool, optional) – If True, an error will be raised when we detect that there exists an input such that all the outputs are independent of it. If False, we return a Tensor of zeros as the jacobian for said input, which is the expected mathematical value. Defaults to False.

• vectorize (bool, optional) – When computing the jacobian, usually we invoke autograd.grad once per row of the jacobian. If this flag is True, we perform only a single autograd.grad call with batched_grad=True which uses the vmap prototype feature. Though this should lead to performance improvements in many cases, because this feature is still experimental, there may be performance cliffs. See torch.autograd.grad()’s batched_grad parameter for more information.

• strategy (str, optional) – Set to "forward-mode" or "reverse-mode" to determine whether the Jacobian will be computed with forward or reverse mode AD. Currently, "forward-mode" requires vectorized=True. Defaults to "reverse-mode". If func has more outputs than input, "forward-mode" tends to be more performant. Otherwise, prefer to use "reverse-mode".

• flatten (bool, optional) – If True, all module parameters and outputs are flattened and concatenated to form a single vector. The Jacobian will be computed with respect to this single flattened vectors, thus a single Tensor will be returned.

Returns:

if there is a single parameter and output, this will be a single Tensor containing the Jacobian for the linearized parameter and output. If there are more than one parameters, then the Jacobian will be a tuple of Tensors. If there are more than one outputs (even if there is only one parameter), then the Jacobian will be a tuple of tuple of Tensors where Jacobian[i][j] will contain the Jacobian of the ith output and jth parameter and will have as size the concatenation of the sizes of the corresponding output and the corresponding parameter and will have same dtype and device as the corresponding parameter. If strategy is forward-mode, the dtype will be that of the output; otherwise, the parameters.

Return type:

Jacobian (Tensor or nested tuple of Tensors)

Warning

The function modjac calculate Jacobian of model parameters. This is in contrast to PyTorch’s function jacobian, which computes the Jacobian of a given Python function.

Example

Calculates Jacobian with respect to all model parameters.

>>> model = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1)
>>> input = torch.randn(1, 1, 1)
>>> J = pp.optim.functional.modjac(model, input)
(tensor([[[[[[[0.3365]]]]]]]), tensor([[[[1.]]]]))
>>> [j.shape for j in J]
[torch.Size([1, 1, 1, 1, 1, 1, 1]), torch.Size([1, 1, 1, 1])]

Function with flattened parameters returns a combined Jacobian.

>>> input = torch.randn(2, 2, 2)
>>> model = nn.Conv2d(in_channels=2, out_channels=2, kernel_size=1)
>>> J = pp.optim.functional.modjac(model, input, flatten=True)
tensor([[-0.4162,  0.0968,  0.0000,  0.0000,  1.0000,  0.0000],
        [-0.6042,  1.1886,  0.0000,  0.0000,  1.0000,  0.0000],
        [ 1.4623,  0.7389,  0.0000,  0.0000,  1.0000,  0.0000],
        [ 1.0716,  2.4293,  0.0000,  0.0000,  1.0000,  0.0000],
        [ 0.0000,  0.0000, -0.4162,  0.0968,  0.0000,  1.0000],
        [ 0.0000,  0.0000, -0.6042,  1.1886,  0.0000,  1.0000],
        [ 0.0000,  0.0000,  1.4623,  0.7389,  0.0000,  1.0000],
        [ 0.0000,  0.0000,  1.0716,  2.4293,  0.0000,  1.0000]])
>>> J.shape
torch.Size([8, 6])

Calculate Jacobian with respect to parameter of pypose.LieTensor.

>>> class PoseTransform(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.p = pp.Parameter(pp.randn_so3(2))
...
...     def forward(self, x):
...         return self.p.Exp() * x
...
>>> model, input = PoseTransform(), pp.randn_SO3()
>>> J = pp.optim.functional.modjac(model, input, flatten=True)
tensor([[ 0.4670,  0.7041,  0.0029,  0.0000,  0.0000,  0.0000],
        [-0.6591,  0.4554, -0.2566,  0.0000,  0.0000,  0.0000],
        [-0.2477,  0.0670,  0.9535,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  0.8593,  0.2672,  0.3446],
        [ 0.0000,  0.0000,  0.0000, -0.2417,  0.9503, -0.1154],
        [ 0.0000,  0.0000,  0.0000, -0.3630, -0.0179,  0.9055],
        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]])
>>> J.shape
torch.Size([8, 6])