How to define a model as a function of its parameters and get its Jacobian?

[homocomputeris]

If one defines a model f, and x is training data, then we have f(x).
How can one concisely define a model f(w) as a function of its parameters w? Will PyTorch be able to differentiate it wrt w, like JAX/STAX can?

[ptrblck]

Could you explain this use case a bit or post some dummy code what should be achieved?

[homocomputeris]

Basically, I just want to linearize a model with respect to its weights to have f(w) = f(w0) + jac(w0)@(w-w0) when x is fixed.

JAX, for example, already has Jacobian and Hessian functions, and it’s quite low-lever, so it allows to set weights directly.

[avish]

Given a model (inherited from nn.Module), calculate Jacobian of parameters of this model with respect to output vector.
This looks like a hack to do this,

Using `autograd.functional.jacobian`/`hessian` with respect to `nn.Module` parameters

I’m also working on this issue when calculating the influence function. My current solution is to follow this workaround: So I write this sample code (which works as a workaround): import torch import torch.nn as nn _input=torch.randn(32,3) layer = nn.Linear(3,4) criterion=nn.CrossEntropyLoss() weight = layer.weight def func(weight): del layer.weight layer.weight=weight return criterion(layer(_input), torch.zeros(len(_input),dtype=torch.long)) torch.autograd.functional.hessia…

[Ren_Pang (Local State)]

Hi, I see that you are using the unsafe hack I previously proposed.
Just want to notice you that there seems to be a new elegant way to do this.
Please take a look at NN Module functional API · Issue #49171 · pytorch/pytorch · GitHub
which shall provide an elegant solution after 1.11 release.