AdamW is an adaptive gradient-based optimizer and a variant of Adam. It updates model parameters to reduce training error while applying weight decay directly to the parameters after optimization, rather than directly to gradients. section 3.4.1 section 3.1.2
A typical training step connects AdamW to the surrounding optimization process:
- Run a forward pass to produce predictions, compare them with targets through a loss function, and backpropagate the gradients. section 3.1.2
- Call AdamW to update the network weights in the direction that reduces the error. section 3.1.2
- Apply weight decay separately to the parameters, helping prevent the decay rate from adapting with the learning rates. section 3.4.1
PyTorch configurations
A ViT classifier uses AdamW with a learning rate of 0.01 and weight decay of 1e-2. A diffusion denoising U-Net uses learning rate 0.0001, betas of (0.95, 0.999), weight decay of 0.00001, and eps of 1e-8. section 3.4.1 section 7.3.2
from diffusers.optimization import get_scheduler
optimizer=torch.optim.AdamW(model.parameters(),lr=0.0001,
betas=(0.95,0.999),weight_decay=0.00001,eps=1e-8) #1
lr_scheduler=get_scheduler( #2
"cosine",
optimizer=optimizer,
num_warmup_steps=300,
num_training_steps=(len(train_dataloader) * 100))
Inspect a complete AdamW configuration alongside a cosine learning-rate scheduler with warmup for diffusion training. section 7.3.2
