What is the correct way to use mixed-precision training with OneCycleLR

Anirudh_G: <p>I am using the onecyclelr along with mixed precision training</p>
<pre><code class="lang-auto">model = SwinNet().to(params['device'])
criterion = nn.BCEWithLogitsLoss().to(params['device'])
optimizer = torch.optim.AdamW(model.parameters(), lr=params['lr'])
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, 
                                          2e-3, 
                                          epochs=5, 
                                          steps_per_epoch=len(train_loader)
                                         )

def train(dataloader, model, criterion, optimizer, epoch, params):
    
    model.train()
    
    scaler = torch.cuda.amp.GradScaler() # enable mixed precision training
    
    stream = tqdm(dataloader)
    
    train_loss = 0
    
    for i, (images, target) in enumerate(stream, start=1):

        images = images.to(params['device'], non_blocking=True)
        target = target.to(params['device'], non_blocking=True).float().view(-1, 1)
        images, targets_a, targets_b, lam = mixup_data(images, target.view(-1, 1))
        
        with torch.cuda.amp.autocast(): # wrapper for mixed precision training
            output = model(images).to(params['device'])
            loss = mixup_criterion(criterion, output, targets_a, targets_b, lam)
        
        train_loss += loss
        
        optimizer.zero_grad(set_to_none=True) 
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.step(scheduler)
        scaler.update()
            
    train_loss /= len(dataloader)
        
    return train_loss
       
</code></pre>
<p>I am getting the following error:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/9/990b73cc0e21170b3546bd1cd7aad3edc0ba8681.png" data-download-href="https://discuss.pytorch.org/uploads/default/990b73cc0e21170b3546bd1cd7aad3edc0ba8681" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/990b73cc0e21170b3546bd1cd7aad3edc0ba8681_2_690x427.png" alt="image" data-base62-sha1="lPTAxZq0RFEY0dliJjakOLqsPdL" width="690" height="427" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/990b73cc0e21170b3546bd1cd7aad3edc0ba8681_2_690x427.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/990b73cc0e21170b3546bd1cd7aad3edc0ba8681_2_1035x640.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/990b73cc0e21170b3546bd1cd7aad3edc0ba8681_2_1380x854.png 2x" data-dominant-color="EDD7D7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1942×1202 370 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>What is the correct way to implement mixed precision along with onecyclelr ?</p>

ptrblck: <p>The <code>step</code> method of the <code>GradScaler</code> expects an optimizer object, not a scheduler, so remove this line of code.</p>

Anirudh_G: <p>Thanks for the response, it fixed the error!</p>
<p>I do have a question though -<br>
The PyTorch doc for <a href="https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.OneCycleLR.html" rel="noopener nofollow ugc">OneCycleLR</a> mentions that the scheduler should be called after every batch. However, gradscaler is expecting an optimizer object. Is there a way where the lr schedule &amp; gradscaler can be used together?</p>

ptrblck: <p>You can call the <code>scheduler.step()</code> after <code>scaler.step(optimizer)</code> to adapt the learning rate.</p>