HVP w.r.t model parameters

username: dedeswim
cooked: <p>Hi!</p>
<p>I have seen that within the 1.5.0 release, the possibility to compute HVP of a function has been added.</p>
<p>As far as I understand from the documentation, the HVP (as well as the VHP, VJP and so on) can be computed w.r.t. the input only, and not w.r.t. some other variable (such as, for instance, the parameters of a model). Is my understanding correct? Is there any way to use the built-in HVP (or VHP) function differentiating w.r.t. the model parameters?</p>
<p>Thanks!</p>

username: albanD
cooked: <p>Hi,</p>
<p>They calculate the gradient wrt of the input of the function you give. So if your functions takes as input the parameters you want and evaluate it on a fixed input. THen it will compute the value wrt to the parameters.</p>

username: dedeswim
cooked: <p>Hi,</p>
<p>Thanks for your helpful reply! So, in case I have a model <code>model</code>, and I want to compute the VHP of the loss w.r.t. the parameters, how do I create a function that computes the loss and takes the parameters as input? Right now I have:</p>
<pre><code class="lang-auto">f = lambda x: criterion(model(img), label)
vhps = [vhp(f, params, grad)[1] for params, grad in zip(model.parameters(), grads)]
</code></pre>
<p>Where the model is a logistic regression trained for MNIST, so its parameters are a list of this shape <code>[(10, 784), (10)]</code>. Instead, <code>grads</code> is the gradient w.r.t. to the parameters, so it has the same shape as the params.</p>
<p>However, with this, I get a list of the same shape as the parameters, but it is all zeros, which makes me suppose that the function doesn’t really depend on <code>model.parameters()</code>. Do you have any suggestions? I’d like the solution to be scalable to an kind of model, so creating a function that uses the parameters to compute the logistic regression by hand wouldn’t be much helpful.</p>
<p>Thanks in advance!</p>

username: albanD
cooked: <p>Hi,</p>
<p>You don’t have to do each parameter one by one, you can give all the params/grads as tuples.</p>
<p>You get all zeros because your function <code>f</code> does not use the inputs <code>x</code> to compute the output.<br>
You can do something like this to use the autograd API with torch.nn:</p>
<pre><code class="lang-python"># Utilities to make nn.Module functional
def del_attr(obj, names):
    if len(names) == 1:
        delattr(obj, names[0])
    else:
        del_attr(getattr(obj, names[0]), names[1:])
def set_attr(obj, names, val):
    if len(names) == 1:
        setattr(obj, names[0], val)
    else:
        set_attr(getattr(obj, names[0]), names[1:], val)

def make_functional(mod):
    orig_params = tuple(mod.parameters())
    # Remove all the parameters in the model
    names = []
    for name, p in list(mod.named_parameters()):
        del_attr(mod, name.split("."))
        names.append(name)
    return orig_params, names

def load_weights(mod, names, params):
    for name, p in zip(names, params):
        set_attr(mod, name.split("."), p)

N = 10
model = models.resnet18(pretrained=False)
criterion = torch.nn.CrossEntropyLoss()

params, names = make_functional(model)
# Make params regular Tensors instead of nn.Parameter
params = tuple(p.detach().requires_grad_() for p in params)

inputs = torch.rand([N, 3, 224, 224], device=device)
labels = torch.rand(N, device=device).mul(10).long()
def f(*new_params):
    load_weights(model, names, new_params)
    out = model(inputs)

    loss = criterion(out, labels)
    return loss

vhp(f, params, grads)</code></pre>

username: dedeswim
cooked: <p>Thanks a lot, it seems to be working!</p>
<p>I only had to change</p>
<pre><code class="lang-python">def load_weights(mod, names, params):
    for name, p in zip(names, params):
        set_attr(mod, name.split("."), p)
</code></pre>
<p>to</p>
<pre><code class="lang-python">def load_weights(mod, names, params):
    for name, p in zip(names, params):
        set_attr(mod, name.split("."), torch.nn.Parameter(p))
</code></pre>
<p>Because otherwise, after calling <code>make_functional</code>, and then <code>f</code>, the model parameters were not working properly and, when calling vhp a second time, I get</p>
<pre><code class="lang-auto">v is a tuple of invalid length: should be 0 but got 2.
</code></pre>
<p>Which I guess is due to the fact that when we set the weights back, they’re not set as <code>Parameter</code>s. Does it make sense?</p>

username: albanD
cooked: <p>This is expected!<br>
Parameters are only leaf Tensors, but here your new module works with any input that is given. They can be leafs or not.</p>
<p>Making these Parameters might break the graph with the other layers above, you should not do that.</p>

username: dedeswim
cooked: <p>I have just realized that, actually, using <code>torch.nn.Parameter</code> makes the VHP 0 again (and with <code>strict=True</code> an exception is raised.</p>
<p>The problem is solved if I load the params with <code>torch.nn.Parameter</code> after the computation.</p>
<p>So I changed <code>load_weights</code> to</p>
<pre><code class="lang-python">def load_weights(model, names, params, as_params=False):
    for name, p in zip(names, params):
        if not as_params:
            set_attr(model, name.split("."), p)
        else:
            set_attr(model, name.split("."), torch.nn.Parameter(p))
</code></pre>
<p>And I call</p>
<pre><code class="lang-python">load_weights(model, names, params, as_params=True)
</code></pre>
<p>I have just read your answer, though, and I guess this would be wrong too, right?</p>
<p>Do you have any suggestion to have back the <code>nn</code> model after making it functional?</p>
<p>Thanks!</p>

username: albanD
cooked: <p>Loading them back as <code>nn.Paramter()</code> looks good yes to recover the original <code>nn.Module</code> back.<br>
Just be careful that you need to call <code>make_functional</code> again before being able to use it in a functional way again.</p>

username: dedeswim
cooked: <p>Sure, thank you very much!</p>

username: peria1
cooked: <p>This is a little old, but I think others may have this same question.</p>
<p>Can you explain why it is necessary to remove the parameters from the model prior to calling vhp?  I have been playing with your example, and I can see that yes, it is necessary. Furthermore, I see what you did to restore your model to usefulness using your newest load_weights().  But I would like to understand why removing the parameters from the model is necessary in the first place.</p>
<p>Thanks very much!</p>

username: umarbek_nasimov
cooked: <p><a class="mention" href="/u/dedeswim">@dedeswim</a> i am getting the same error that hvp is 0 (and the exception thrown). i think it’s because once <code>p</code> is casted to a Parameter, it is copied so the reference to the original <code>p</code> is lost. did you find a way around this?</p>
<pre><code class="lang-auto">parameter = torch.nn.Parameter(tensor)
print(parameter.data)
tensor = torch.zeros(tensor.shape)
print(parameter.data)
</code></pre>
<p>here ^ print outputs the same result</p>

username: tengerye
cooked: <aside class="quote no-group" data-username="peria1" data-post="11" data-topic="83520">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/ac8455/48.png" class="avatar"> peria1:</div>
<blockquote>
<p>Can you explain why it is necessary to remove the parameters from the model prior to calling vhp?</p>
</blockquote>
</aside>
<p>I have similar confusion here.</p>

username: albanD
cooked: <p>Note that today, you can simply use <code>torch.func.functional_call()</code> to achieve this!</p>