Manually modify gradients of two models, average them and put them back in both models!

Rio1210: <p>I am trying to train two models on two mutually exclusive portions of a datasets.<br>
Now while I am training both the models, I want to manually extract the Gradients from Model A and Model B, after forward propagation, then before updating the weights, I want to average both the model’s gradients and put the average of the models in both the models, and update the weights.</p>
<p>I am trying it in the following way:</p>
<pre><code class="lang-auto">def mergeGrad(modelA, modelB, i):
    listGradA = [] 
    listGradB = [] 
    itr = 0
    for pA, pB in zip(modelA.parameters(), modelB.parameters()):
        listGradA.append(pA)
        listGradB.append(pB)
        avg = (pA.grad + pB.grad)/2
        pA.grad = avg
        pB.grad = avg
        itr += 1
</code></pre>
<p>In the optimization, portion I am doing</p>
<pre><code class="lang-auto">       optimizerA.zero_grad()
        lossA.backward()
        
        optimizerB.zero_grad()
        lossB.backward()
        
        print(gradientA, gradientB)
        mergeGrad(modelA, modelB, i)
        print(gradientA, gradientB)
</code></pre>
<p>Now, I am printing before and after averaging and putting the Gradient back. However, only in the very first iteration the first print statement outputs separate gradients for modelA and modelB, and right after merging them, I get the averaged gradients as I want.</p>
<p>However, from the second iteration onward, both the models are always generating the same gradients after forward propagation (and hence the same averaged values). How is this possible?<br>
The models are seeing different data!</p>
<p>Is the manual gradient modification that I am doing is wrong?<br>
Please do suggest if there’s a better way to manually modify and put them back into the model!<br>
Thanks a lot!</p>

albanD: <p>Hi,</p>
<p>I would bet the issue is with:</p>
<pre><code class="lang-auto">        pA.grad = avg
        pB.grad = avg
</code></pre>
<p>Here you set the same Tensor to be the gradient for both parameters. So the backward passes will accumulate in the Same Tensor.<br>
If you don’t want that, you need to add a <code>.clone()</code> for at least one of them.</p>
<p>Note that it would be a neat way to implement what you want as well: make the <code>.grad </code> fields the same so the two backward passes actually accumulate in the same Tensor. And divide your learning rate by 2 so that you still do the same update (assuming simple SGD).</p>

Rio1210: <p>Thanks a lot. <code>grad.clone()</code> solved the issue. Also appreciate the neat trick you mentioned!</p>

Ohm: <p>Hi,<br>
Do you have any efficient averaging approach?<br>
Since if the number of models increases the ‘for loop’ does not make that much good performance based on time…</p>

albanD: <p>At the moment there is not much no. If your weights are large enough that shouldn’t be a problem.<br>
But we are working on improving that see the work in progress (so subject to heavy change!!) optimizers: <a href="https://github.com/pytorch/pytorch/blob/master/torch/optim/_multi_tensor/sgd.py">https://github.com/pytorch/pytorch/blob/master/torch/optim/_multi_tensor/sgd.py</a></p>

Gurpriya: <p>Can we use clone in line :<br>
avg = (pA.grad + pB.grad)/2</p>
<p>Like:<br>
avg = (pA.grad.clone() + pB.grad.clone())/2</p>
<p>and then simply use that average like:<br>
pA.grad = avg<br>
pB.grad = avg</p>

albanD: <p>That doesn’t change the fact that <code>pA.grad is pB.grad</code> and so the same issue will appear.</p>

Gurpriya: <p>Okay, thank you so much.</p>
<p>And do we have a way other than manually zipping it for say 10 clients?</p>

albanD: <p>What do you mean zip?<br>
The suggestion above was to just replace this to</p>
<pre><code class="lang-python">pA.grad = avg
pB.grad = avg.clone()
</code></pre>