|
| 1 | +""" |
| 2 | +Momentum SGD Optimizer |
| 3 | +
|
| 4 | +Implements SGD with momentum for neural network training using NumPy. |
| 5 | +Momentum helps accelerate gradients in the relevant direction and dampens oscillations. |
| 6 | +
|
| 7 | +Reference: https://en.wikipedia.org/wiki/Stochastic_gradient_descent#Momentum |
| 8 | +Author: Adhithya Laxman Ravi Shankar Geetha |
| 9 | +Github: https://github.com/Adhithya-Laxman |
| 10 | +Date: 2025.10.22 |
| 11 | +""" |
| 12 | + |
| 13 | +import numpy as np |
| 14 | + |
| 15 | + |
| 16 | +class MomentumSGD: |
| 17 | + """ |
| 18 | + SGD with momentum optimizer. |
| 19 | +
|
| 20 | + Updates parameters using momentum: |
| 21 | + velocity = momentum * velocity - learning_rate * gradient |
| 22 | + param = param + velocity |
| 23 | + """ |
| 24 | + |
| 25 | + def __init__(self, learning_rate: float = 0.01, momentum: float = 0.9) -> None: |
| 26 | + """ |
| 27 | + Initialize Momentum SGD optimizer. |
| 28 | +
|
| 29 | + Args: |
| 30 | + learning_rate (float): Learning rate for weight updates. |
| 31 | + momentum (float): Momentum factor. |
| 32 | +
|
| 33 | + >>> optimizer = MomentumSGD(learning_rate=0.01, momentum=0.9) |
| 34 | + >>> optimizer.momentum |
| 35 | + 0.9 |
| 36 | + """ |
| 37 | + self.learning_rate = learning_rate |
| 38 | + self.momentum = momentum |
| 39 | + self.velocity: dict[int, np.ndarray] = {} |
| 40 | + |
| 41 | + def update( |
| 42 | + self, param_id: int, params: np.ndarray, gradients: np.ndarray |
| 43 | + ) -> np.ndarray: |
| 44 | + """ |
| 45 | + Update parameters using momentum. |
| 46 | +
|
| 47 | + Args: |
| 48 | + param_id (int): Unique identifier for parameter group. |
| 49 | + params (np.ndarray): Current parameters. |
| 50 | + gradients (np.ndarray): Gradients of parameters. |
| 51 | +
|
| 52 | + Returns: |
| 53 | + np.ndarray: Updated parameters. |
| 54 | +
|
| 55 | + >>> optimizer = MomentumSGD(learning_rate=0.1, momentum=0.9) |
| 56 | + >>> params = np.array([1.0, 2.0]) |
| 57 | + >>> grads = np.array([0.1, 0.2]) |
| 58 | + >>> updated = optimizer.update(0, params, grads) |
| 59 | + >>> updated.shape |
| 60 | + (2,) |
| 61 | + """ |
| 62 | + if param_id not in self.velocity: |
| 63 | + self.velocity[param_id] = np.zeros_like(params) |
| 64 | + |
| 65 | + self.velocity[param_id] = ( |
| 66 | + self.momentum * self.velocity[param_id] - self.learning_rate * gradients |
| 67 | + ) |
| 68 | + return params + self.velocity[param_id] |
| 69 | + |
| 70 | + |
| 71 | +# Usage example |
| 72 | +if __name__ == "__main__": |
| 73 | + import doctest |
| 74 | + |
| 75 | + doctest.testmod() |
| 76 | + |
| 77 | + print("Momentum SGD Example: Minimizing f(x) = x^2") |
| 78 | + |
| 79 | + optimizer = MomentumSGD(learning_rate=0.1, momentum=0.9) |
| 80 | + x = np.array([5.0]) |
| 81 | + |
| 82 | + for step in range(20): |
| 83 | + gradient = 2 * x |
| 84 | + x = optimizer.update(0, x, gradient) |
| 85 | + if step % 5 == 0: |
| 86 | + print(f"Step {step}: x = {x[0]:.4f}, f(x) = {x[0] ** 2:.4f}") |
| 87 | + |
| 88 | + print(f"Final: x = {x[0]:.4f}, f(x) = {x[0] ** 2:.4f}") |
0 commit comments