Batch-frequency optim scheduler and step order for torch>=1.1.0

username: austinmw
cooked: <p>Hi, I’ve set up my code to use either epoch-frequency schedulers like <code>MultiStepLR</code> or batch-frequency schedulers like <code>OneCycleLR</code>. In both cases <code>optimizer.step()</code> happens before <code>scheduler.step()</code>. However, when I use a batch-frequency scheduler, I get the warning advising me that the ordering is wrong:</p>
<blockquote>
<p>UserWarning: Detected call of <code>lr_scheduler.step()</code> before <code>optimizer.step()</code>. In PyTorch 1.1.0 and later, you should call them in the opposite order: <code>optimizer.step()</code> before <code>lr_scheduler.step()</code>.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at <a href="https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate" class="inline-onebox" rel="noopener nofollow ugc">torch.optim — PyTorch 2.1 documentation</a></p>
</blockquote>
<p>Is there a chance this is a false warning in the batch-scheduler case, or is there still likely to be something incorrect with my code?</p>

username: ptrblck
cooked: <p>There might always be a bug, we haven’t found yet, but I haven’t seen this error in the latest release yet, so I guess your code might be using the scheduler and optimizer in a wrong order accidentally.<br>
Could you post a minimal, executable code snippet which shows this warning?<br>
Also, are you using automatic mixed-precision training? If so, note that the <code>optimizer.step()</code> call can be skipped, if invalid gradients are found and the <code>GradScaler</code> needs to reduce the scaling factor.</p>

username: austinmw
cooked: <p>Hi, thanks for your reply. I am also using AMP (also using SWA). Here’s the skeleton of what I’m doing:</p>
<pre><code class="lang-python">model = create_model(args)
model = nn.DataParallel(model)
swa_model = torch.optim.swa_utils.AveragedModel(model)
scaler = torch.cuda.amp.GradScaler()
optimizer = torch.optim.AdamW(model.parameters(), lr=7e-3)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer,
                                max_lr=7e-2,
                                epochs=5, steps_per_epoch=len(dataloader), 
                                pct_start=0.3)  
swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, 
                                            swa_lr=0.5, 
                                            anneal_epochs=5)
num_epochs = 20
swa_start = 15

for epoch in (1, range(num_epochs)+1):
    for imgs,labels in dataloader:

        with torch.cuda.amp.autocast():
            output = model(imgs)
            loss = loss_fn(output, labels)

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        
        if epoch &lt; swa_start:
            scheduler.step()
            
    if epoch &gt;= swa_start:
        swa_model.update_parameters(model)
        swa_scheduler.step()
</code></pre>

username: ptrblck
cooked: <p>In that case <code>amp</code> can skip the <code>optimizer.step()</code>, if invalid gradients are found and you should thus also skip the <code>scheduler.step()</code>.<br>
As a workaround, you could check the scale value via <code>scaler.get_scale()</code> and skip the <code>scheduler.step()</code>, if it was decreased. We’ll provide a better API for it so that you could directly check, if the <code>optimizer.step()</code> was skipped.</p>

username: austinmw
cooked: <p>Thanks, something like this?</p>
<pre><code class="lang-python">for epoch in range(epochs):
    for batch in dataloader:
        start_scale = scaler.get_scale()
        # do stuff
        end_scale = scaler.get_scale()
        if not end_scale &lt; start_scale:
            scheduler.step()
</code></pre>
<p>I’m not quite sure if I follow where to place the checks. If I put the start_scale check at the beginning of each batch loop and the end_scale check at the very end, for nearly every batch I get:<br>
start_scale == 65536.0<br>
end_scale == 32.0<br>
And so it would never get to the scheduler.step()</p>
<p><strong>Edit:</strong> Actually by the start of epoch 2 it looks like it settles so that both start_scale and end_scale == 32.0</p>
<p>Should I similarly suppress the swa_scheduler step or is that left alone?</p>

username: ptrblck
cooked: <p>The scale factor should balance itself after a while and not decrease in each iteration. A value of <code>32.0</code> sounds reasonable. The <code>end_scale</code> should be extracted after the <code>scaler.step()</code> operation.</p>
<p>What is <code>swa_scheduler</code> exactly doing? If it’s (re-)creating the “averaged” model using the updated parameters, I think you should also skip it, since the model wasn’t updated.</p>