Average each weight of two models

username: littlemountainman
cooked: <p>Hello,</p>
<p>I have two models trained exactly the same way just with a different learning rate now I would like to average each weight of every single layer in the model an create a new one with the weight averages. Is that possible ?</p>
<p>Thanks in advance</p>

username: ptrblck
cooked: <p>Yes, you could get the <code>state_dicts</code> of both models, average the parameters and reload the new <code>state_dict</code>.<br>
Here is a small dummy example:</p>
<pre><code class="lang-python"># Setup
modelA = nn.Linear(1, 1)
modelB = nn.Linear(1, 1)

sdA = modelA.state_dict()
sdB = modelB.state_dict()

# Average all parameters
for key in sdA:
    sdB[key] = (sdB[key] + sdA[key]) / 2.

# Recreate model and load averaged state_dict (or use modelA/B)
model = nn.Linear(1, 1)
model.load_state_dict(sdB)
</code></pre>

username: littlemountainman
cooked: <p>Thanks that works. Thank you so much</p>

username: omarfoq
cooked: <p>Hello <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Is there a proper way of parallelizing the loop over the keys of the state_dict ?</p>
<p>Thanks in advance</p>

username: ptrblck
cooked: <p>I guess Python might provide some multiprocessing utilities for this use case. However, how often are you using this code snippet, as it seems it’s the bottleneck you would like to optimize?</p>

username: vdw
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Slightly off-topic: Is there any solid evidence that averaging the weights of different runs has any benefits?</p>
<p>In principle, both runs may converge two very different minima w.r.t. the the loss. It’s not obvious to my, why averaging the weights would be a good idea. Admittedly, this might depend what “trained exactly the same way” means, but even a different order of batches might effect this, wouldn’t it?</p>

username: omarfoq
cooked: <p>Hello</p>
<p>I use this averaging after each optimizer step, I agree that using multiprocessing utilities from python can be a good idea, however in general this kind of utilities creates jobs pool on the cpu, then execute them, and I am not sure if this the best possible way of  doing it. I am wondering if there is a possibility to create the jobs pool on the gpu directly. I believe that it won’t make any significant difference, but I am not sure about it.</p>

username: ptrblck
cooked: <p>At some point I’ve looked into <a href="https://arxiv.org/abs/1803.05407">Stochastic Weight Averaging</a>, which claims that a simple averaging of multiple checkpoints leads to a better generalization.<br>
If I’m not mistaken, there is also an <code>SWA</code> package for PyTorch, which applies this strategy during your training.<br>
That being said, I don’t know if these experiments still hold true for other models than the one mentioned in the paper.<br>
Generally, I would assume that this method would only work, if all training checkpoints are close to the global minimum, which seems to be the case if you use the checkpoints from a single run (if I remember the paper correctly).</p>
<p><a class="mention" href="/u/omarfoq">@omarfoq</a> I think this operation might benefit from the current ongoing port of <code>apex'</code> <code>multi_tensor_apply</code> in <a href="https://github.com/pytorch/pytorch/pull/41554">this PR</a>.<br>
You could have a look at the code changes and check, if you could reuse it.<br>
However, as mentioned before, I would profile the code beforehand and make sure that a speedup would be visible in your case.</p>

username: vdw
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thanks, good points. Yes, when the weights come from the same run at different checkpoints, I can see that this might be beneficial.</p>