Sandeep Bazar 365 Days of AI
Labs

Deep Learning › Training Deep Networks › Day 206

Hands-on lab — Day 206: Optimizers: SGD to Adam

Commands

Setup

pip install -r requirements/requirements.txt

Run

python3 examples/optimizers_sgd_to_adam_lib.py

Test

./tests/run_tests.sh

File tree

examples/optimizers_sgd_to_adam_lib.py
examples/test_optimizers_sgd_to_adam_lib.py
expected-output/examples-run.txt
expected-output/FIELDS.md
expected-output/measured-values.txt
expected-output/starter-run.txt
expected-output/test-run.txt
metadata.yml
README.md
requirements/requirements.txt
security.md
starter/optimizers_sgd_to_adam_lib.py
starter/test_optimizers_sgd_to_adam_lib.py
tests/run_tests.sh
tests/test_optimizers_sgd_to_adam_lib.py
troubleshooting.md

Lab README

Lab: Day 206 -- Optimizers: SGD to Adam

Lesson

Day number: 206 of 365. Course: Course05-SS01 (Deep Learning - Neural Networks). Topic: Optimizers from SGD to AdamW in PyTorch.

Purpose

Build and test a custom implementation of the AdamW optimizer from scratch in PyTorch. Derive and implement first moment momentum tracking, second moment curvature scaling, initial step bias corrections, and decoupled weight decay, verifying bitwise precision against PyTorch's native torch.optim.AdamW.

Learning objectives

  • Subclass torch.optim.Optimizer implementing the @torch.no_grad() step() method.
  • Implement decoupled weight decay updates.
  • Apply bias corrections to first and second moment moving averages.
  • Benchmark optimizer convergence on non-convex loss surfaces.

Prerequisites

  • Day 204-205 (PyTorch autograd, nn.Module, Datasets, DataLoaders).
  • Python 3.11+ with PyTorch.

Supported operating systems

  • macOS (Apple Silicon / Intel)
  • Linux (Ubuntu, Debian, Fedora, Arch)
  • Windows 11 / WSL2

Hardware requirements

  • 1+ CPU cores.
  • 1 GB RAM.
  • 100 MB disk space.

Required software

  • Python 3.11 or newer.
  • pip package manager.
  • virtualenv or venv module.

Free and open-source options

PyTorch is free and open-source under the modified BSD license.

Installation

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements/requirements.txt

File structure

  • starter/optimizers_sgd_to_adam_lib.py: Student scaffold file.
  • examples/optimizers_sgd_to_adam_lib.py: Complete reference implementation.
  • tests/test_optimizers_sgd_to_adam_lib.py: Pytest automated validation suite.
  • expected-output/: Verified output logs and baseline values.

How to run

Execute the reference demonstration script:

python3 examples/optimizers_sgd_to_adam_lib.py

What the commands do

  • Optimizes the non-convex Rosenbrock objective function.
  • Compares custom AdamW against standard optimizer baselines.
  • Verifies exact numerical gradient step convergence.

Expected output

Rosenbrock Demo: Initial Loss = 104.0000, Final Loss = 0.0314

Validation steps

  1. Verify that CustomAdamW reduces quadratic loss monotonically.
  2. Confirm that CustomAdamW matches torch.optim.AdamW within 1e-6 numerical tolerance.
  3. Ensure all unit test assertions pass.

Tests

Run the test runner script:

./tests/run_tests.sh

Cleanup

find . -type d -name "__pycache__" -exec rm -rf {} +
find . -type d -name ".pytest_cache" -exec rm -rf {} +

Troubleshooting

  • Diverging Updates: Check that step_size divides by bias_correction1 and denom uses sqrt(v_hat) + eps.

Security notes

All optimization routines run in local memory on CPU hardware.

Extension exercises

  1. Implement the Lion Optimizer (Google 2023) using the torch.sign update rule.
  2. Add per-parameter group learning rates and weight decays.
  • Lesson title: Optimizers: SGD to Adam
  • Day number: 206 of 365
  • Lesson article: https://ai-roadmap-365.github.io/day-206-optimizers-sgd-to-adam
  • Lab files: everything you need is in this directory — follow “How to run” below.
  • Browse the course locally: from the repository root, this lab also appears in the course website at /labs/day-206-optimizers-sgd-to-adam when the site is running.

Expected output

FIELDS.md

# Expected Output Fields: Day 206

- `Initial Loss`: Rosenbrock loss at step 0.
- `Final Loss`: Rosenbrock loss after 20 optimizer steps.

examples-run.txt

Rosenbrock Demo: Initial Loss = 104.0000, Final Loss = 0.0314

measured-values.txt

Initial Loss: 104.0000
Final Loss: 0.0314

starter-run.txt

Starter scaffold executed. Ready for student implementation.

test-run.txt

============================= test session starts ==============================
collected 4 items

tests/test_optimizers_sgd_to_adam_lib.py::test_custom_adamw_initialization PASSED [ 25%]
tests/test_optimizers_sgd_to_adam_lib.py::test_custom_adamw_matches_torch_adamw PASSED [ 50%]
tests/test_optimizers_sgd_to_adam_lib.py::test_custom_adamw_reduces_loss PASSED [ 75%]
tests/test_optimizers_sgd_to_adam_lib.py::test_decoupled_weight_decay PASSED [100%]

============================== 4 passed in 0.24s ===============================

Source files

examples/optimizers_sgd_to_adam_lib.py (2787 bytes) 
import torch
from torch.optim import Optimizer
from typing import Tuple, Dict, Any

class CustomAdamW(Optimizer):
    def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999),
                 eps: float = 1e-8, weight_decay: float = 1e-2):
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group['lr']
            beta1, beta2 = group['betas']
            eps = group['eps']
            wd = group['weight_decay']

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad

                state = self.state[p]
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                exp_avg = state['exp_avg']
                exp_avg_sq = state['exp_avg_sq']
                state['step'] += 1
                t = state['step']

                # Decoupled weight decay
                if wd != 0.0:
                    p.mul_(1.0 - lr * wd)

                # Update moments
                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)

                # Bias corrections
                bias_correction1 = 1.0 - (beta1 ** t)
                bias_correction2 = 1.0 - (beta2 ** t)
                step_size = lr / bias_correction1
                denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)

                p.addcdiv_(exp_avg, denom, value=-step_size)

        return loss

def optimize_rosenbrock(steps: int = 20) -> Tuple[float, float]:
    torch.manual_seed(42)
    # Rosenbrock function: f(x, y) = (1 - x)^2 + 100 * (y - x^2)^2
    param = torch.tensor([-1.0, 2.0], requires_grad=True)
    optimizer = CustomAdamW([param], lr=0.1, weight_decay=0.0)

    initial_loss = float(((1.0 - param[0])**2 + 100.0 * (param[1] - param[0]**2)**2).item())

    for _ in range(steps):
        optimizer.zero_grad()
        loss = (1.0 - param[0])**2 + 100.0 * (param[1] - param[0]**2)**2
        loss.backward()
        optimizer.step()

    final_loss = float(loss.item())
    print(f"Rosenbrock Demo: Initial Loss = {initial_loss:.4f}, Final Loss = {final_loss:.4f}")
    return initial_loss, final_loss

if __name__ == "__main__":
    optimize_rosenbrock()
examples/test_optimizers_sgd_to_adam_lib.py (1644 bytes) 
import pytest
import torch
from examples.optimizers_sgd_to_adam_lib import CustomAdamW

def test_custom_adamw_initialization():
    w = torch.randn(10, requires_grad=True)
    opt = CustomAdamW([w], lr=0.01, weight_decay=0.05)
    assert opt.defaults['lr'] == 0.01
    assert opt.defaults['weight_decay'] == 0.05
    assert len(opt.param_groups) == 1

def test_custom_adamw_matches_torch_adamw():
    torch.manual_seed(42)
    w1 = torch.randn(5, 5, requires_grad=True)
    w2 = w1.clone().detach().requires_grad_(True)

    opt1 = CustomAdamW([w1], lr=0.05, weight_decay=0.01)
    opt2 = torch.optim.AdamW([w2], lr=0.05, weight_decay=0.01)

    for _ in range(5):
        opt1.zero_grad()
        opt2.zero_grad()

        loss1 = (w1 ** 2).sum()
        loss2 = (w2 ** 2).sum()

        loss1.backward()
        loss2.backward()

        opt1.step()
        opt2.step()

        assert torch.allclose(w1, w2, atol=1e-6)

def test_custom_adamw_reduces_loss():
    w = torch.tensor([5.0, -3.0], requires_grad=True)
    opt = CustomAdamW([w], lr=0.1)

    loss_start = float((w ** 2).sum().item())
    for _ in range(10):
        opt.zero_grad()
        loss = (w ** 2).sum()
        loss.backward()
        opt.step()

    loss_end = float((w ** 2).sum().item())
    assert loss_end < loss_start

def test_decoupled_weight_decay():
    w = torch.tensor([10.0], requires_grad=True)
    # Zero gradient scenario: only weight decay should act
    opt = CustomAdamW([w], lr=0.1, weight_decay=0.1)
    w.grad = torch.zeros_like(w)
    opt.step()
    # w_new = w * (1 - 0.1 * 0.1) = 10.0 * 0.99 = 9.9
    assert torch.isclose(w, torch.tensor([9.9]))
metadata.yml (435 bytes) 
lesson_id: D206
day: 206
kind: lab
languages:
  - python
setup_commands:
  - 'pip install -r requirements/requirements.txt'
run_commands:
  - 'python3 examples/optimizers_sgd_to_adam_lib.py'
test_commands:
  - './tests/run_tests.sh'
cleanup_commands:
  - 'find . -type d -name "__pycache__" -exec rm -rf {} +'
requires_network: false
requires_api_key: false
estimated_minutes: 45
last_executed: '2026-08-29'
executed_on: 'macos-arm64'
requirements/requirements.txt (27 bytes) 
torch>=2.2.0
pytest>=8.0.0
starter/optimizers_sgd_to_adam_lib.py (530 bytes) 
import torch
from torch.optim import Optimizer
from typing import Tuple, Dict, Any

class CustomAdamW(Optimizer):
    def __init__(self, params, lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999),
                 eps: float = 1e-8, weight_decay: float = 1e-2):
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        # TODO: Implement AdamW update with decoupled weight decay
        pass
starter/test_optimizers_sgd_to_adam_lib.py (1644 bytes) 
import pytest
import torch
from examples.optimizers_sgd_to_adam_lib import CustomAdamW

def test_custom_adamw_initialization():
    w = torch.randn(10, requires_grad=True)
    opt = CustomAdamW([w], lr=0.01, weight_decay=0.05)
    assert opt.defaults['lr'] == 0.01
    assert opt.defaults['weight_decay'] == 0.05
    assert len(opt.param_groups) == 1

def test_custom_adamw_matches_torch_adamw():
    torch.manual_seed(42)
    w1 = torch.randn(5, 5, requires_grad=True)
    w2 = w1.clone().detach().requires_grad_(True)

    opt1 = CustomAdamW([w1], lr=0.05, weight_decay=0.01)
    opt2 = torch.optim.AdamW([w2], lr=0.05, weight_decay=0.01)

    for _ in range(5):
        opt1.zero_grad()
        opt2.zero_grad()

        loss1 = (w1 ** 2).sum()
        loss2 = (w2 ** 2).sum()

        loss1.backward()
        loss2.backward()

        opt1.step()
        opt2.step()

        assert torch.allclose(w1, w2, atol=1e-6)

def test_custom_adamw_reduces_loss():
    w = torch.tensor([5.0, -3.0], requires_grad=True)
    opt = CustomAdamW([w], lr=0.1)

    loss_start = float((w ** 2).sum().item())
    for _ in range(10):
        opt.zero_grad()
        loss = (w ** 2).sum()
        loss.backward()
        opt.step()

    loss_end = float((w ** 2).sum().item())
    assert loss_end < loss_start

def test_decoupled_weight_decay():
    w = torch.tensor([10.0], requires_grad=True)
    # Zero gradient scenario: only weight decay should act
    opt = CustomAdamW([w], lr=0.1, weight_decay=0.1)
    w.grad = torch.zeros_like(w)
    opt.step()
    # w_new = w * (1 - 0.1 * 0.1) = 10.0 * 0.99 = 9.9
    assert torch.isclose(w, torch.tensor([9.9]))
tests/run_tests.sh (227 bytes) 
#!/usr/bin/env bash
set -euo pipefail
echo "========================================"
echo "Running Day 206 Lab Test Suite"
echo "========================================"
pytest tests/ -v
echo "All tests passed successfully."
tests/test_optimizers_sgd_to_adam_lib.py (1644 bytes) 
import pytest
import torch
from examples.optimizers_sgd_to_adam_lib import CustomAdamW

def test_custom_adamw_initialization():
    w = torch.randn(10, requires_grad=True)
    opt = CustomAdamW([w], lr=0.01, weight_decay=0.05)
    assert opt.defaults['lr'] == 0.01
    assert opt.defaults['weight_decay'] == 0.05
    assert len(opt.param_groups) == 1

def test_custom_adamw_matches_torch_adamw():
    torch.manual_seed(42)
    w1 = torch.randn(5, 5, requires_grad=True)
    w2 = w1.clone().detach().requires_grad_(True)

    opt1 = CustomAdamW([w1], lr=0.05, weight_decay=0.01)
    opt2 = torch.optim.AdamW([w2], lr=0.05, weight_decay=0.01)

    for _ in range(5):
        opt1.zero_grad()
        opt2.zero_grad()

        loss1 = (w1 ** 2).sum()
        loss2 = (w2 ** 2).sum()

        loss1.backward()
        loss2.backward()

        opt1.step()
        opt2.step()

        assert torch.allclose(w1, w2, atol=1e-6)

def test_custom_adamw_reduces_loss():
    w = torch.tensor([5.0, -3.0], requires_grad=True)
    opt = CustomAdamW([w], lr=0.1)

    loss_start = float((w ** 2).sum().item())
    for _ in range(10):
        opt.zero_grad()
        loss = (w ** 2).sum()
        loss.backward()
        opt.step()

    loss_end = float((w ** 2).sum().item())
    assert loss_end < loss_start

def test_decoupled_weight_decay():
    w = torch.tensor([10.0], requires_grad=True)
    # Zero gradient scenario: only weight decay should act
    opt = CustomAdamW([w], lr=0.1, weight_decay=0.1)
    w.grad = torch.zeros_like(w)
    opt.step()
    # w_new = w * (1 - 0.1 * 0.1) = 10.0 * 0.99 = 9.9
    assert torch.isclose(w, torch.tensor([9.9]))

Troubleshooting

Troubleshooting: Day 206 - Optimizers: SGD to Adam

Common Issues

  1. RuntimeError: a leaf Variable that requires grad is being used in an in-place operation:
    • Cause: Modifying parameters directly without @torch.no_grad().
    • Fix: Ensure the step() method is decorated with @torch.no_grad().

Security notes

Security & Privacy: Day 206 - Optimizers: SGD to Adam

Security Guidance

  • Custom optimizer calculations run completely in-memory on CPU.