SGD with Momentum - TorchCode

Problem Description

Implement SGD with momentum from scratch. Momentum accumulates a velocity — an exponentially-weighted average of past gradients — so consistent directions build up speed while oscillations across a ravine cancel out.

Signature

class MySGD:
    def __init__(self, params, lr=0.1, momentum=0.9): ...
    def step(self): ...
    def zero_grad(self): ...

Algorithm (per parameter)

v = momentum * v + grad p -= lr * v

Template

Implement the class below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

class MySGD:
    def __init__(self, params, lr=0.1, momentum=0.9):
        pass  # store params, lr, momentum; init velocity buffers to zeros

    def step(self):
        pass  # v = momentum*v + grad ; p -= lr*v

    def zero_grad(self):
        pass  # zero all gradients

Test Your Implementation

Use this code to debug before submitting.

# 🧪 Debug
torch.manual_seed(0)
w = torch.randn(4, 3, requires_grad=True)
opt = MySGD([w], lr=0.1, momentum=0.9)
for i in range(5):
    loss = (w ** 2).sum()
    loss.backward()
    opt.step()
    opt.zero_grad()
    print(f'Step {i}: loss={loss.item():.4f}')

Reference Solution

Try solving it yourself first! Click below to reveal the solution.

# ✅ SOLUTION

class MySGD:
    def __init__(self, params, lr=0.1, momentum=0.9):
        self.params = list(params)
        self.lr = lr
        self.momentum = momentum
        self.v = [torch.zeros_like(p) for p in self.params]

    def step(self):
        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                self.v[i] = self.momentum * self.v[i] + p.grad
                p -= self.lr * self.v[i]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

Tips

Run Locally

For interactive practice with auto-grading, run TorchCode locally:
pip install torch-judge then use check("sgd_momentum")

Key Concepts

Velocity buffer (EWA of gradients), damps oscillations in steep ravines. With a constant gradient g the velocity settles at g/(1−momentum) — a speed-up along steady directions. Covered in DL Module 7 §4.