Averaging models with differentiable weighting coefficients

username: mzimmer
cooked: <p>Suppose I have a couple of models given as state_dicts and instead of optimizing them further, I’d like for every batch to average these models (their weights!) with a weighting coefficient such as 0.1<em>model_1 + 0.4</em>model_2 +…</p>
<p>However, I do not intend to optimize model_i further, I rather like to perform forward- and backward with the averaged model and then update the coefficients.</p>
<p>How can I average the models such that the grad_fn w.r.t. the coefficient tensor persists? I fail at doing this, I always lose the gradients as soon as I assign the averaged parameters (which have valid grad_fns) to the base model used for inference.</p>
<p>Thanks!</p>

username: KFrank
cooked: <p>Hi Max!</p>
<aside class="quote no-group" data-username="mzimmer" data-post="1" data-topic="177240" data-full="true">
<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/mzimmer/48/31510_2.png" class="avatar"> mzimmer:</div>
<blockquote>
<p>How can I average the models such that the grad_fn w.r.t. the coefficient tensor persists? I fail at doing this, I always lose the gradients as soon as I assign the averaged parameters (which have valid grad_fns) to the base model used for inference.</p>
</blockquote>
</aside>
<p>Pytorch layers such as <code>Linear</code> an <code>Conv2</code> are pretty insistent about<br>
having their weights be <code>Parameter</code>s and leaf tensors.  There may be<br>
some way around this that lets you assign a non-leaf tensor, such as<br>
averaged parameters, to be the weight of such a layer, but I don’t<br>
know how to do it.</p>
<p>What you can do is let your “weighting coefficients” be trainable<br>
parameters (with <code>requires_grad = True</code>), compute your averaged<br>
weight tensors as a function of the “weighting coefficients” and the<br>
weight tensors of your <code>model_i</code>.  Then pass the averaged weight<br>
tensors into the <em>functional forms</em> for layers such as <a href="https://pytorch.org/docs/2.0/generated/torch.nn.functional.linear.html" rel="noopener nofollow ugc">linear()</a> and <a href="https://pytorch.org/docs/2.0/generated/torch.nn.functional.conv2d.html" rel="noopener nofollow ugc">conv2()</a>.</p>
<p>Gradients will now pass back properly through the functional-form<br>
layers to your trainable “weighting coefficients” parameters, which<br>
you can then optimize.</p>
<p>This has the disadvantage that you have to rewrite your model to use<br>
functional-form layers – you can’t just instantiate a new instance of<br>
your model and, for example, load its state dict with averaged weights.<br>
But I don’t know of a scheme that avoids such a rewrite.</p>
<p>Best.</p>
<p>K. Frank</p>

username: ptrblck
cooked: <p>That’s a good idea and I believe <code>torch.nn.utils.parametrize</code> could be provide a good approach for this use case:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import torch.nn.utils.parametrize as parametrize


class Average(nn.Module):
    def __init__(self, w1, w2):
        super().__init__()
        self.w1 = w1
        self.w2 = w2
    
    def forward(self, X):
        return (self.w1 + self.w2) / 2.


layer1 = nn.Linear(10, 10)
layer2 = nn.Linear(10, 10)

optimizer = torch.optim.Adam(list(layer1.parameters()) + list(layer2.parameters()), lr=1.)

layer3 = nn.Linear(10, 10)
parametrize.register_parametrization(layer3, "weight", Average(layer1.weight, layer2.weight))

# check if parametrization works
print(((layer1.weight + layer2.weight) / 2. - layer3.weight).abs().max())
# tensor(0., grad_fn=&lt;MaxBackward1&gt;)

# original weights are still different
print((layer1.weight - layer2.weight).abs().max())
# tensor(0.5682, grad_fn=&lt;MaxBackward1&gt;)

x = torch.randn(1, 10)
out = layer3(x)
out.mean().backward()

# check for valid gradients
print(layer1.weight.grad.abs().sum())
# tensor(5.5434)
print(layer2.weight.grad.abs().sum())
# tensor(5.5434)

# update parameters
optimizer.step()

# check again if parametrization works
print(((layer1.weight + layer2.weight) / 2. - layer3.weight).abs().max())
# tensor(0., grad_fn=&lt;MaxBackward1&gt;)

# original weights are still different
print((layer1.weight - layer2.weight).abs().max())
# tensor(0.5682, grad_fn=&lt;MaxBackward1&gt;)
</code></pre>

username: mzimmer
cooked: <p>Thanks, I think I got it to work in a similar fashion!</p>

username: mzimmer
cooked: <p>Unfortunately, I ran into a problem with the current way of doing it. When working on a single GPU, everything works fine. As soon as I get into DataParallel mode, I get errors during inference that input and weight tensors do not lie on the same device - after trying out several things I thought that this makes sense, since my model is now DataParallel and my Parametrization (in your example Average) lies on a single GPU.</p>
<p>I tried the following to also have it in DataParallel mode (adapted to your example):</p>
<pre><code class="lang-auto">reparametrization_instance = Average(layer1.weight, layer2.weight)
reparametrization_instance = torch.nn.DataParallel(reparametrization_instance)
parametrize.register_parametrization(module, 'weight', reparametrization_instance)
</code></pre>
<p>However, I get the following error:</p>
<blockquote>
<p>Registering a parametrization may not change the shape of the tensor, unless <code>unsafe</code> flag is enabled.<br>
unparametrized shape: torch.Size([64, 3, 7, 7])<br>
parametrized shape: torch.Size([128, 3, 7, 7])</p>
</blockquote>
<p>Do you have any advice? Thanks a lot!</p>

username: ptrblck
cooked: <p>I don’t know if parametrizations are supported with <code>nn.DataParallel</code> as it’s deprecated and we generally recommend using <code>DistributedDataParallel</code> instead, which I assume should work fine.</p>