Set parameters to model without breaking autograd

username: qq-me
cooked: <p>I am trying to do this:</p>
<pre><code class="lang-auto">def closure(params):
    for p, new_p in zip(optimizer_params, params):
        p.set_(new_p) # this errors when using torch.func.jvp
    preds = model(inputs)
    loss = loss_fn(preds, targets)
    return loss

torch.func.jvp(closure, params, vec)
</code></pre>
<p>The way this could be done is via <code>torch.func.functional_call</code>. This way I can do</p>
<pre><code class="lang-auto">def closure(params):
    params_dict = {k: p for (k, v), p in zip(model.named_parameters(), params) if v.requires_grad}
    preds = torch.func.functional_call(model, params_dict, (inputs, ))
    loss = loss_fn(preds, targets)
    return loss

torch.func.jvp(closure, params_tuple, vec)
</code></pre>
<p>However it would be more convenient if I could somehow use the first way, because second one requires me to use a torch.nn.Module instead of arbitrary list of tensors, I have to rewrite my closure and mess with named_parameters</p>

username: qq-me
cooked: <p>I found this method here <a href="https://github.com/gngdb/pytorch-minimize/blob/master/pytorch_minimize/optim.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch-minimize/pytorch_minimize/optim.py at master · gngdb/pytorch-minimize · GitHub</a></p>
<pre><code class="lang-auto">def f(x):
                        _params = self.unravel_unpack(x)
                        # monkey patch substitute variables
                        named_params = list(model.named_parameters())
                        for _p, (n, _) in zip(_params, named_params):
                            rdelattr(model, n)
                            rsetattr(model, n, _p)
                        return closure.loss(model)
                    def numpyify(x):
                        if x.device != torch.device('cpu'):
                            x = x.cpu()
                        #return x.numpy().astype(np.float64)
                        return self.floatX(x.numpy())
                    return numpyify(torch.autograd.functional.hessian(f, x))
</code></pre>
<p>I am afraid that will break references that other things have to those parameters, and it still requires a torch.nn.Module. But maybe something like this is possible without modifying attributes?</p>

username: soulitzer
cooked: <p>Maybe something like <a href="https://pytorch.org/docs/stable/generated/torch.utils.swap_tensors.html" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.swap_tensors — PyTorch 2.5 documentation</a> would be helpful?</p>

username: qq-me
cooked: <p>thank you, this actually worked! Here is a minimal example of how to make this work</p>
<pre><code class="lang-auto">import torch, torch.nn as nn

model = nn.Sequential(nn.Linear(2, 3), nn.ReLU(), nn.Linear(3, 2))
params = tuple(model.parameters())
inputs = torch.randn(2)
targets = torch.randn(2)

def closure():
    """a typical closure but without backward"""
    preds = model(inputs)
    loss = nn.functional.mse_loss(preds, targets)
    return loss

def param_closure(*new_params):
    """closure that takes in params as inputs to work with torch.func.jvp.
    For example an optimizer can create this from normal closure"""
    # swap params to new params
    for old_p, new_p in zip(params, new_params):
        torch.utils.swap_tensors(old_p, new_p)

    value = closure()

    # swap params back to original ones
    for old_p, new_p in zip(params, new_params):
        torch.utils.swap_tensors(old_p, new_p)

    return value

tangents = tuple(torch.randn_like(i) for i in params)
torch.func.jvp(param_closure, primals = tuple(i.clone().detach_() for i in params), tangents = tangents)
</code></pre>
<p>Two things though, first I have to clone and detach the parameters before passing them to the closure, which is a tiny bit inefficient. Second thing is that on a small 1D convnet this doesn’t seem to be any faster than calling backward and calculating the dot product of gradient with my vector, and on a (badly optimized) box packing task this is actually 2 times slower than backward. I haven’t tested anything large though.</p>

username: qq-me
cooked: <p>UPDATE: works with vmapping as well</p>
<pre data-code-wrap="py"><code class="lang-py">vmapped_param_closure = torch.vmap(param_closure)
stacked_params = [torch.randn((1024, *i.size()), device=i.device, dtype=i.dtype) for i in params]
vmapped_param_closure(*stacked_params)
</code></pre>
<p>thats insane for optimizers. I just tested and trying 1000 params is only 22 times slower</p>