Implement Custom SGD Optimizer

hardPROCustom OptimizersPyTorch Advanced

Problem Statement

Implement a basic SGD optimizer from scratch by subclassing torch.optim.Optimizer.

Background

Building a custom optimizer teaches the fundamentals: iterate parameter groups, access .grad, and update .data in-place.

Your Task

The starter code defines a SimpleSGD class with init already implemented. Implement the step method that applies the basic SGD update rule to each parameter that has a gradient.

The training loop and verification are pre-filled.

Output Format

Returns a dictionary with "initial_weight", "final_weight", "weight_changed", and "is_optimizer".

Example

Input

None

Output

{'initial_weight': [0.7645, -0.8049, 0.2343], 'final_weight': [0.8373, -0.3944, 0.9289], 'weight_changed': True, 'is_optimizer': True}

Reasoning

We start by defining the SimpleSGD class, which inherits from torch.optim.Optimizer, and create an instance of it with a learning rate of 0.1.
We then create a nn.Linear(3, 1) model, record its initial weights, and run 5 training steps using the SimpleSGD optimizer with mean squared error (MSE) loss: $L = \frac{1}{n} \sum_{i=1}^{n} (y_i - \hat{y_i})^2$ , where $y_i$ is the target and $\hat{y_i}$ is the predicted output.
In each training step, the weights are updated according to the SGD update rule: $param.data -= lr \cdot param.grad$ , where $lr$ is the learning rate and $param.grad$ is the gradient of the loss with respect to the parameter.
After the training steps, we record the final weights and compare them to the initial weights to determine if the weights have changed, and check if SimpleSGD is an instance of torch.optim.Optimizer to determine the value of "is_optimizer".

Constraints

Subclass torch.optim.Optimizer
Implement step() with basic SGD update rule
Use @torch.no_grad() for step method

Learn this in

PretrainingLLM Study Plan

More from PyTorch Advanced

Implement SGD with Momentumhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy Implement Bilinear Form with Einsumhard

Concept primerPREVIEW

Background Knowledge

The problem involves implementing a basic Stochastic Gradient Descent (SGD) optimizer from scratch using PyTorch. SGD is a widely used optimization algorithm in machine learning for minimizing the loss function of a model. It works by iteratively updating the model's parameters in the direction of the negative gradient of the loss function. In the context of PyTorch, optimizers like SGD are implemented as classes that inherit from torch.optim.Optimizer. This base class provides a framework for defining custom optimizers by implementing the init and step methods.

To understand this problem, it's essential to have a grasp of the backpropagation algorithm, which is used to compute the gradients of the loss function with respect to the model's parameters. Additionally, familiarity with PyTorch's tensor operations and the nn.Module API is necessary. The problem also involves defining a custom optimizer class, SimpleSGD, which will inherit from torch.optim.Optimizer. This requires understanding of object-oriented programming principles and how to work with PyTorch's optimizer API.

The problem statement involves training a simple neural network model using the custom SGD optimizer. This requires understanding how to define a model, compute the loss function, and perform backpropagation to update the model's parameters. The problem also involves recording the initial and final weights of the model, which requires understanding how to access and manipulate the model's parameters.

Algorithm/Approach

The general approach to solving this problem involves:

Defining a custom optimizer class, SimpleSGD, that inherits from torch.optim.Optimizer.
Implementing the init method to initialize the optimizer with the model's parameters and learning rate.
Implementing the step method to update the model's parameters using the SGD algorithm.
Creating a simple neural network model and defining a loss function.
Training the model using the custom optimizer and recording the initial and final weights.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

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.

solution.py

Python 3.13

Test results