Implement SGD with Momentum
Problem Statement
Extend SGD with momentum, tracking velocity buffers in the optimizer state.
Background
Momentum SGD maintains a velocity: v = momentum * v + grad, then param -= lr * v. This smooths updates and helps escape local minima. The velocity buffer must persist across optimizer steps.
Your Task
The starter code defines a MomentumSGD class with init already implemented. Implement the step method that maintains a velocity buffer for each parameter and applies the momentum update rule.
The training loop and verification are pre-filled.
Output Format
Returns a dictionary with "initial_weight", "final_weight", "has_velocity", and "weight_changed".
Example:
None
{'initial_weight': [0.7645, -0.8049, 0.2343], 'final_weight': [1.0577, 0.3795, 2.5983], 'has_velocity': True, 'weight_changed': True}- We initialize the
MomentumSGDoptimizer with a learning rate of 0.1 and momentum of 0.9, and create ann.Linear(3, 1)model. - The initial weights of the model are recorded as
[0.7645, -0.8049, 0.2343], which will be used to calculate theinitial_weightoutput. - We run 5 training steps with MomentumSGD, updating the model parameters using the momentum update rule: v=0.9⋅v+grad, then param−=0.1⋅v, which changes the weights to
[1.0577, 0.3795, 2.5983]. - The optimizer state is checked for velocity buffers, and since
MomentumSGDstores velocity inself.state[p],has_velocityis True, and since the weights have changed,weight_changedis also True.
Constraints:
- Store velocity in self.state[p]
- Update: v = momentum * v + grad, param -= lr * v
- Use @torch.no_grad() on step
Background Knowledge
The problem involves implementing a custom optimizer in PyTorch, specifically Stochastic Gradient Descent (SGD) with Momentum. SGD is a fundamental optimization algorithm used in machine learning to minimize the loss function of a model. It iteratively updates the model's parameters in the direction of the negative gradient of the loss function. However, SGD can suffer from oscillations and slow convergence, especially in the presence of local minima. To address this, Momentum SGD introduces a momentum term that helps smooth the updates and escape local minima. The momentum term is calculated as a fraction of the previous update, which is added to the current update.
The Momentum SGD update rule can be written as: v = momentum * v + grad, where v is the velocity, momentum is the momentum coefficient, and grad is the gradient of the loss function. The parameter update is then param -= lr * v, where lr is the learning rate. This update rule helps the optimizer to "remember" the previous updates and make more informed decisions about the direction of the next update. The momentum term can be thought of as a "memory" that helps the optimizer to avoid oscillations and converge faster.
In the context of PyTorch, the torch.optim.Optimizer class provides a base class for implementing custom optimizers. To implement Momentum SGD, we need to define a custom optimizer class that inherits from torch.optim.Optimizer and overrides the init and step methods. The init method is used to initialize the optimizer's parameters, such as the learning rate and momentum coefficient. The step method is used to perform a single update step, which involves calculating the velocity and updating the parameters.
Algorithm/Approach
The general approach to solving this problem involves the following steps:
- Define a custom optimizer class that inherits from torch.optim.Optimizer
- Initialize the optimizer's parameters, such as the learning rate and momentum coefficient
- Implement the step method to perform a single update step, which involves calculating the velocity and updating the parameters
- Use the custom optimizer to train a model and verify that it works as expected
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.