My program was broken by AdamW and then my console became red!

username: jack_S
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/5/05f32bbbb48f7fc1b932d5a945787bc817664355.png" data-download-href="https://discuss.pytorch.org/uploads/default/05f32bbbb48f7fc1b932d5a945787bc817664355" title="redconsole"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/5/05f32bbbb48f7fc1b932d5a945787bc817664355.png" alt="redconsole" data-base62-sha1="QDnsNmLvh4ibvRlH4sTAhJ05p3" width="690" height="332" data-dominant-color="2D2A2A"><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">redconsole</span><span class="informations">1496×720 60.5 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><br>
This happened when I trained <code>resnet18</code> with <code>CIFAR10</code>. If I use <code>SGD</code>, then the program can run normally. This is how I get <code>AdamW</code>:</p>
<pre><code class="lang-auto">parameters = {'params':model.parameters(), 'lr':learning_rate}
parameters['weight_decay'] = 1e-4
optimizer = optim.AdamW(**parameters)
</code></pre>

username: ptrblck
cooked: <p>Could you post a minimal and executable code snippet to reproduce the issue, please?</p>

username: jack_S
cooked: <p>Thank you! I found the crux. I set the learning rate in <code>AdamW</code> with a <code>tensor</code>. Here is the code:</p>
<pre><code class="lang-auto">import torch
from torch.optim import SGD, AdamW
from torch.cuda import amp
from torch.cuda.amp.grad_scaler import GradScaler
from torchvision import models
from torch.nn import functional as F

model = models.resnet18().cuda()

scaler = GradScaler()
L = torch.nn.CrossEntropyLoss(reduction = 'mean').cuda()

parameters = {'params':model.parameters(), 'lr':0.002}
parameters['weight_decay'] = 2e-5
optimizer = AdamW(**parameters)                 #SGD works fine
for param_group in optimizer.param_groups:
    param_group['lr'] = torch.tensor(0.)
optimizer.zero_grad(True)

images = torch.randn(128, 3, 32, 32).cuda(non_blocking = True)
labels = F.one_hot(torch.randint(0, 9, (128,)), 10).type(torch.float32).cuda(non_blocking = True)
with amp.autocast():
    output = model(images)
    loss = L(output, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
</code></pre>

username: crcrpar
cooked: <p>Hi <a class="mention" href="/u/jack_s">@jack_S</a>,</p>
<p>it looks a bit surprising to me that SGD works with Tensor learning rate considering the type hint of <a href="https://github.com/pytorch/pytorch/blob/fc63d710fee323eee8b135fd193ee37e9f06ed55/torch/optim/sgd.py#L187-L200" class="inline-onebox" rel="noopener nofollow ugc">pytorch/sgd.py at fc63d710fee323eee8b135fd193ee37e9f06ed55 · pytorch/pytorch · GitHub</a> where <code>lr</code>’s type hint is <code>float</code>.</p>
<p>One way to let AdamW with Tensor lr would be to pass <code>fused=False, foreach=False</code> to AdamW constructor.</p>
<p>Would you mind telling me why you want to have lr in Tensor?</p>

username: jack_S
cooked: <p>Well, it’s not very important for me to assign <code>lr</code> with a <code>tensor</code>. I can turn to <code>float</code> in other experiments if necessary. I did it just because I want to see what will happen if I set the <code>lr</code> casually. For example, if I want to set a random learning schedule for my training, I may code like below:</p>
<pre><code class="lang-auto">total_epochs = 200
lr_schedule = torch.normal(mean = 0.05, std = 0.01, size = (total_epochs, )).abs()
optimizer = AdamW(model.parameters(), 0.001)
for i in range(0, total_epochs):
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr_schedule[i]   
    ............
</code></pre>
<p>So this introduces <code>tensor</code> in <code>lr</code>.</p>