About using GradScaler() for not only one model in a training progress

username: Hedwig_Huang
cooked: <p>Hi there. For saving the GPU memory, I use FP16 in my work just like nnUNet did, they defined a GradScaler() for updating the gradients and stuff. But I modified the code for my own usage:</p>
<p><em>For updating loss and optimizer for <strong>twice</strong>.</em></p>
<ul>
<li>Compared with the original update order like <code>optimizer1.step()</code>, and <code>optimizer2.step()</code> followed.What’s the difference? Can it work normally as we want?</li>
</ul>
<p>And here, with the example provided <a href="https://pytorch.org/docs/stable/amp.html" rel="noopener nofollow ugc">Pytorch[AUTOMATIC MIXED PRECISION PACKAGE - TORCH.AMP]</a> as below:</p>
<pre><code class="lang-auto">scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()
</code></pre>
<p>for update <code>optimizer1</code> and <code>optimizer2</code>, I change the code like this:</p>
<pre><code class="lang-auto">scaler.scale(loss).backward()
scaler.unscale_(optimizer1)
scaler.unscale_(optimizer2)  # new
torch.nn.utils.clip_grad_norm_(model1.parameters(), max_norm)
torch.nn.utils.clip_grad_norm_(model2.parameters(), max_norm)  # new
scaler.step(optimizer1)
scaler.step(optimizer2)  # new
scaler.update()
</code></pre>
<ul>
<li>But cause such error:</li>
</ul>
<pre><code class="lang-auto">AssertionError: No inf checks were recorded for this optimizer.
</code></pre>
<hr>
<p>The example for <code>loss</code> is in : <a href="https://discuss.pytorch.org/t/runtimeerror-trying-to-backward-through-the-graph-a-second-time-but-the-buffers-have-already-been-freed-specify-retain-graph-true-when-calling-backward-the-first-time/6795/130">For loss update twice in Gradscaler.scale(loss).backward()</a></p>
<hr>
<p><strong>How can I run FP16 with Gradscaler for updating optimizer/loss for twice?</strong><br>
Thank you!</p>

username: Hedwig_Huang
cooked: <p>Or i should just defined two different <code>scaler</code>, but how can I update my loss?</p>

username: ptrblck
cooked: <p>Your general code works correctly as seen here:</p>
<pre><code class="lang-python">model1 = nn.Linear(1, 1).cuda()
model2 = nn.Linear(1, 1).cuda()
optimizer1 = torch.optim.Adam(model1.parameters(), lr=1e-3)
optimizer2 = torch.optim.Adam(model2.parameters(), lr=1e-3)

scaler = torch.cuda.amp.GradScaler()

for _ in range(10):
    x = torch.randn(1, 1).cuda()
    with torch.cuda.amp.autocast():
        out = model1(x)
        out = model2(out)
        loss = out.mean()
    
    
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer1)
    scaler.unscale_(optimizer2)  # new
    max_norm = 1.
    torch.nn.utils.clip_grad_norm_(model1.parameters(), max_norm)
    torch.nn.utils.clip_grad_norm_(model2.parameters(), max_norm)  # new
    scaler.step(optimizer1)
    scaler.step(optimizer2)  # new
    scaler.update()
</code></pre>
<p>and I cannot reproduce the reported error.</p>

username: Hedwig_Huang
cooked: <p>I found your code is work for me, too, when I tried it in the console. But can’t still fix the <code>AssertionError</code>, and I know that it seems work for <code>GradScaler()</code>.</p>
<p>But reported for the other bug not related with it. Now I n-checked my code and I got what’s wrong and no errors report.</p>
<p>Thanks a lot! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <p>Good to hear you’ve found the issue! Would you mind explaining what went wrong in your code as I would be interested what has raised the error?</p>

username: Hedwig_Huang
cooked: <p>It sounds dumb <img src="https://discuss.pytorch.org/images/emoji/apple/smiling_face_with_tear.png?v=12" title=":smiling_face_with_tear:" class="emoji" alt=":smiling_face_with_tear:" loading="lazy" width="20" height="20"> , for the mean-teacher framework, I am training my work with loading weigths for teacher model from trained model, and I forgot to remove the teacher model from the part of <code>torch.no_grad()</code>, so there is not any gradients feed back, might be why this AssertionError cause.<br>
After fixed this, no errors reported anymore.</p>