Understanding Adapted Learning Rates in Adam with PyTorch

• Internal Calculation: The adapted rate is an internal variable used by the Adam algorithm. It's not meant to be directly accessed or modified by the user.

There are alternative approaches to monitor the learning rate's behavior:

import torch
from torch import nn
from torch.optim import Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau

# Define model and loss function (replace with your specific model and loss)
model = nn.Linear(10, 1)
loss_fn = nn.MSELoss()

# Sample data (modify based on your data)
x = torch.randn(100, 10)
y = torch.randn(100, 1)

# Set initial learning rate and other Adam parameters
learning_rate = 0.01
beta1 = 0.9
beta2 = 0.999
epsilon = 1e-8

# Create optimizer with Adam and learning rate scheduler
optimizer = Adam(model.parameters(), lr=learning_rate, betas=(beta1, beta2), eps=epsilon)
scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=5)  # Reduce LR on plateau

# Training loop (adjust for your training needs)
for epoch in range(10):
  for i in range(len(x)):
    # Forward pass, calculate loss
    y_pred = model(x[i])
    loss = loss_fn(y_pred, y[i])

    # Backward pass and update weights
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

  # Get current learning rates after update step
  for param_group in optimizer.param_groups:
    current_lr = param_group['lr']
    print(f"Epoch: {epoch+1}, Current Learning Rate: {current_lr}")

  # Scheduler step (update learning rate based on validation loss)
  # Replace with your validation logic
  # scheduler.step(validation_loss)

In this example, the ReduceLROnPlateau scheduler is used. Within the training loop, after the optimizer update step, the code iterates through optimizer parameter groups and retrieves the current learning rate using param_group['lr']. This reflects the adjusted learning rate by Adam based on the gradients encountered during training.

1. Custom Callback:

   • Create a custom callback class that gets called at specific points during training (e.g., after each epoch).
   • Inside the callback function, access the optimizer's parameter groups and extract the learning rate similar to the scheduler example.
   • This approach offers more flexibility to log or visualize the learning rate alongside other training metrics.

2. TensorBoard Integration:

   • If you're using TensorBoard for visualization, leverage its functionalities to track the learning rate.
   • During training, manually log the learning rate (e.g., using writer.add_scalar) at desired intervals.
   • This allows visualizing the learning rate alongside loss and other metrics within TensorBoard.

3. Monitoring Gradients:

   • While not a direct reflection of the learning rate, monitoring the gradients during training can be informative.
   • Large gradients might indicate the need for a smaller learning rate to prevent oscillations, while very small gradients might suggest a stagnant learning process.

Here's a brief code example for a custom callback:

class LearningRateMonitor(object):
  def __init__(self, writer):
    self.writer = writer
    self.epoch = 0

  def on_epoch_end(self, trainer):
    optimizer = trainer.optimizer
    for param_group in optimizer.param_groups:
      current_lr = param_group['lr']
      self.writer.add_scalar('Learning Rate', current_lr, self.epoch)
    self.epoch += 1

This callback logs the learning rate to TensorBoard after each epoch. Remember to adapt it to your specific training loop and logging setup.