# Average each weight of two models

**URL:** <https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008>\
**Category:** nlp\
**Created:** [April 16, 2020, 9:12am UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008 "2020-04-16T09:12:02Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![littlemountainman](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/958977/32.png) [@littlemountainman](https://discuss.pytorch.org/u/littlemountainman)\
**Post date:** [April 16, 2020, 9:12am UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/1 "2020-04-16T09:12:02Z")

</div>

Hello,

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 ?

Thanks in advance

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [April 16, 2020, 11:57pm UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/2 "2020-04-16T23:57:56Z")

</div>

Yes, you could get the `state_dicts` of both models, average the parameters and reload the new `state_dict`.  
Here is a small dummy example:

```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)

```

---

<div class="post-metadata">

**Author:** ![littlemountainman](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/958977/32.png) [@littlemountainman](https://discuss.pytorch.org/u/littlemountainman)\
**Post date:** [April 17, 2020, 1:39pm UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/3 "2020-04-17T13:39:38Z")

</div>

Thanks that works. Thank you so much

---

<div class="post-metadata">

**Author:** ![omarfoq](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/omarfoq/32/27473_2.png) [@omarfoq](https://discuss.pytorch.org/u/omarfoq)\
**Post date:** [August 4, 2020, 3:12pm UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/4 "2020-08-04T15:12:43Z")

</div>

Hello @ptrblck,

Is there a proper way of parallelizing the loop over the keys of the state\_dict ?

Thanks in advance

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [August 5, 2020, 7:25am UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/5 "2020-08-05T07:25:42Z")

</div>

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?

---

<div class="post-metadata">

**Author:** ![vdw](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vdw/32/10074_2.png) [@vdw](https://discuss.pytorch.org/u/vdw)\
**Post date:** [August 5, 2020, 7:57am UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/6 "2020-08-05T07:57:46Z")

</div>

@ptrblck Slightly off-topic: Is there any solid evidence that averaging the weights of different runs has any benefits?

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?

---

<div class="post-metadata">

**Author:** ![omarfoq](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/omarfoq/32/27473_2.png) [@omarfoq](https://discuss.pytorch.org/u/omarfoq)\
**Post date:** [August 5, 2020, 9:07am UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/7 "2020-08-05T09:07:53Z")

</div>

Hello

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.

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [August 5, 2020, 9:34am UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/8 "2020-08-05T09:34:57Z")

</div>

At some point I’ve looked into [Stochastic Weight Averaging](https://arxiv.org/abs/1803.05407), which claims that a simple averaging of multiple checkpoints leads to a better generalization.  
If I’m not mistaken, there is also an `SWA` package for PyTorch, which applies this strategy during your training.  
That being said, I don’t know if these experiments still hold true for other models than the one mentioned in the paper.  
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).

@omarfoq I think this operation might benefit from the current ongoing port of `apex'` `multi_tensor_apply` in [this PR](https://github.com/pytorch/pytorch/pull/41554).  
You could have a look at the code changes and check, if you could reuse it.  
However, as mentioned before, I would profile the code beforehand and make sure that a speedup would be visible in your case.

---

<div class="post-metadata">

**Author:** ![vdw](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vdw/32/10074_2.png) [@vdw](https://discuss.pytorch.org/u/vdw)\
**Post date:** [August 5, 2020, 12:34pm UTC](https://discuss.pytorch.org/t/average-each-weight-of-two-models/77008/9 "2020-08-05T12:34:47Z")

</div>

@ptrblck thanks, good points. Yes, when the weights come from the same run at different checkpoints, I can see that this might be beneficial.
