使用amp训练出现Detected call of lr_scheduler.step() before optimizer.step().

问题

UserWarning: Detected call of lr_scheduler.step() before optimizer.step(). In PyTorch 1.1.0 and later, you should call them in the opposite order: optimizer.step() before lr_scheduler.step(). Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate warnings.warn("Detected call of lr_scheduler.step() before optimizer.step().
正常在训练的时候，对于优化器和参数更新的顺序这样就可以了

        total_loss.backward()
        self.optimizer.step()
        self.scheduler.step()

方案

但是在使用amp进行混合精度训练出现的一个warning，调换代码顺序也不work。
这个警告是怎么产生的？是由于在算梯度的时候出现了nan，所以这一次更新会跳过，但是参数仍旧会更新所以出现了问题，按下面的做法加一个判断就行了。

          self.optimizer.zero_grad()
          scaler.scale(total_loss).backward()
          scaler.step(self.optimizer)
          scale_ = scaler.get_scale()
          scaler.update()
 
          if not ((scale_ != scaler.get_scale)):
              self.scheduler.step()