How to enable grad calculation when using stateless.functional_call

cijerezg: <p>Torch 1.12 introduced the wonderful feature: <code>torch.nn.utils.stateless.functional_call()</code>, whose syntax is as follows:<br>
<code>out = functional_call(model, params, inp)</code></p>
<p>Namely, it allows to evaluate the NN with any parameters. However, how can I calculate the gradient with respect to the evaluated parameters of the output? It seems that it doesn’t save anything to allow backward propagation. Is there a way to enable this? or is this not possible with functional_call and I need to go back to the traditional way of just creating a model per set of parameters I want to evaluate?</p>
<p>To be a little more precise, let me add some code:</p>
<pre><code class="lang-auto">params = [params1, params2, params3]
inputs = [inp1, inp2, inp3]

grads = []

for param, inp in zip(params, inputs):
     out = torch.nn.utils.stateless.functional_call(model, param, inp)
     grad = autograd.grad(out, param)
     grads.append(grad)
</code></pre>
<p>Clearly, that won’t work, but I hope that illustrates what I am trying to do.</p>

ptrblck: <aside class="quote no-group" data-username="cijerezg" data-post="1" data-topic="156048">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cijerezg/48/14945_2.png" class="avatar"> cijerezg:</div>
<blockquote>
<p>Clearly, that won’t work</p>
</blockquote>
</aside>
<p>I’m not sure what the issue is, but it seems to work for me:</p>
<pre data-code-wrap="python"><code class="lang-python">model = nn.Linear(10, 10, bias=False)

param = {"weight": nn.Parameter(torch.ones(10, 10))}
inp = torch.randn(1, 10)

out = torch.nn.utils.stateless.functional_call(model, param, inp)
grad = torch.autograd.grad(out.mean(), param['weight'])
print(grad)
</code></pre>

cijerezg: <p>Yes, that works. I found the problem with my code is that I was doing the following:</p>
<pre><code class="lang-auto">model = nn.Linear(10, 10)

param = dict(model.state_dict())
inp = torch.rand(1, 10)

out = torch.nn.utils.stateless.functional_call(model, param, inp)
</code></pre>
<p>Then when I try to evaluate the grad I get the error: <code>RuntimeError: element 0 of tensors does not require grad</code></p>
<p>With that solved, I have another question. Let’s say the model has a bias, then it looks like this:</p>
<pre><code class="lang-auto">model = nn.Linear(10, 10)

param = {"weight": nn.Parameter(torch.ones(10, 10)),
         "bias": nn.Parameter(torch.ones(10))}

inp = torch.randn(1, 10)

out = torch.nn.utils.stateless.functional_call(model, param, inp)
grad = torch.autograd.grad(out.mean(), param['weight'])
print(grad)
</code></pre>
<p><code>grad</code> will only compute the gradient with respect to the weights and not the bias. What argument should I pass, instead of <code>param['weight']</code>, to compute the gradient with respect to all parameters? I tried <code>param.items()</code>, but that didn’t work.</p>

ptrblck: <p>You can pass a sequence of tensors to <code>autograd.grad</code>:</p>
<pre><code class="lang-python">grad = torch.autograd.grad(out.mean(), [param['weight'], param['bias']])
print(grad)
</code></pre>