Keeping loaded parameters in computation graph of forward pass

thomas-woehrle: <p>Hi,</p>
<p>my question is whether it is possible to load parameters into a model, such that they are still in the computation graph of the result of the forward function of said model.</p>
<p>Let’s say for example</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.autograd import grad

class SimpleNet(nn.Module):
  def __init__(self):
    super().__init__()

    self.fc1 = nn.Linear(4, 1)

  def forward(self, x):
    return self.fc1(x)

model = SimpleNet()

new_theta = {n: p * 2 for n, p in model.named_parameters()}
model.load_state_dict(new_theta)

x = torch.randn(4)
result = model(x)

grad(result, [p for p in new_theta.values()], allow_unused=True)
</code></pre>
<p>The grad at the end gives None for all tensors, because they are detached from the computation graph that calculates result, i.e. the forward function (that’s what I think at least).</p>
<p>Is it somehow possible to load parameters into the model so that this gradient will get populated ?</p>

ptrblck: <p><code>load_state_dict</code> will load the parameters inplace in a <code>no_grad</code> context and will thus not track the gradients to <code>new_theta</code>. The parameters of the model are still used and show valid gradients:</p>
<pre data-code-wrap="python"><code class="lang-python">grad(result, [p for p in model.parameters()], allow_unused=True)
# (tensor([[ 0.4696, -0.3762, -0.4090, -1.1362]]), tensor([1.]))
</code></pre>
<p>What’s your use case exactly and why do you want to see the gradients in <code>new_theta</code>, but not the actually used parameters (you could still copy the gradients if needed)?</p>

thomas-woehrle: <p>First of all, thanks for the answer.</p>
<p>My use case is an implementation of MAML.</p>
<p>I have a theta (the parameters), which I want to use to calculate forward passes. Basically, I want to tell the model “do a normal forward pass, but use these parameters”, so that the computation graph includes the parameters passed.</p>
<p>One way to achieve what I want, is to do something like:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import grad

class SimpleNet(nn.Module):
  def __init__(self):
    super().__init__()

    self.fc1 = nn.Linear(4, 1)

  def forward(self, x, theta):
    return F.linear(x, theta[0], theta[1])
</code></pre>
<p>or similar.</p>
<p>In this case I would define the model in the self.fc1 to still access things like initial parameters via model.parameters(), but the actual forward pass is calculated via “external” parameters.</p>
<p>But this becomes very cumbersome for complicated models → it would be nice to handle this differently</p>

ptrblck: <p>You might want to try <a href="https://pytorch.org/docs/stable/generated/torch.func.functional_call.html"><code>torch.func.functional_call</code></a>, which allows you to:</p>
<blockquote>
<p>Performs a functional call on the module by replacing the module parameters and buffers with the provided ones.</p>
</blockquote>

thomas-woehrle: <p>Okay, that sounds like just what I need.</p>
<p>Thanks! I will try it out later and come back here to accept the solution or ask again.</p>

thomas-woehrle: <p>I tried it out and it does just what I want.</p>
<p>I have one question regarding buffers/in-place operations however, which was not answered in the documentation.</p>
<p>If I have the following net:</p>
<pre><code class="lang-auto">class SimpleNet(nn.Module):
  def __init__(self):
    super().__init__()

    self.net = nn.Sequential(
        nn.Conv2d(3, 12, 3),
        nn.BatchNorm2d(12),
        nn.ReLU(),
        nn.Conv2d(12, 1, 3),
        nn.BatchNorm2d(1),
        nn.ReLU(),
        nn.Flatten(),
        nn.Linear(36, 1)
    )



  def forward(self, x):
    return self.net(x)
</code></pre>
<p>and use it like this:</p>
<pre><code class="lang-auto">model = SimpleNet()
x = torch.randn(5, 3, 10, 10)

external_params = {n: p*2 for n, p in model.named_parameters()}
external_buffers = {n: torch.ones_like(b) for n, b in model.named_buffers()}

result = functional_call(model, (external_params, external_buffers), x)
</code></pre>
<p>Then, I can see that the external_buffers are updated (as described in the documentation). However, I was surprised that</p>
<pre><code class="lang-auto">result = functional_call(model, external_params, x) 
</code></pre>
<p>works as well. In this case, I assume the buffers of the model are used and updated (f.e. in the case of BatchNorm ?</p>

ptrblck: <p>Yes, the internal parameters and buffers will be used if the <code>strict=False</code> argument is passed (used by default). This example also demonstrates it:</p>
<pre data-code-wrap="python"><code class="lang-python">class SimpleNet(nn.Module):
  def __init__(self):
    super().__init__()

    self.lin1 = nn.Linear(2, 2)
    self.lin2 = nn.Linear(2, 2)

  def forward(self, x):
    return self.lin1(x) + self.lin2(x)


model = SimpleNet()
x = torch.randn(1, 2)

external_params = {n: p*20000 for n, p in model.named_parameters() if "1" in n}
result = torch.func.functional_call(model, external_params, x, strict=False)

result.mean().backward()

print(model.lin1.weight.grad)
# tensor([[-4179.4907, -9056.8799],
#         [-4179.4907, -9056.8799]])
print(model.lin2.weight.grad)
# tensor([[-0.2090, -0.4528],
#         [-0.2090, -0.4528]])
</code></pre>

thomas-woehrle: <p>Okay, got it, perfect!</p>
<p>Thanks und Danke Dir!</p>