Gradient Accumulation Over Steps

mediumPROGradient ManipulationPyTorch Advanced

Problem Statement

Implement gradient accumulation, where you accumulate gradients over multiple mini-batches before taking an optimizer step.

Background

When GPU memory is limited, you can simulate a larger batch size by accumulating gradients over N forward/backward passes before calling optimizer.step(). Only zero gradients every N steps.

Your Task

The starter code provides 4 mini-batches, a model, and an optimizer. Accumulate gradients over all 4 mini-batches before taking a single optimizer step. This simulates a larger effective batch size of 4x.

Output Format

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

Example

Input

None

Output

{'initial_weight': [0.7645, -0.8049], 'final_weight': [19.8107, 22.0265], 'weight_changed': True}

Reasoning

We start by initializing the model with torch.manual_seed(42) and creating an nn.Linear(2, 1) model, resulting in initial weights of approximately $[0.7645, -0.8049]$ .
We then accumulate gradients over 4 mini-batches, computing the MSE loss for each batch and calling backward() without zeroing the gradients between batches. The losses are calculated as $((prediction - target)^2)$ for each batch.
After all 4 batches, we call optimizer.step() to update the model weights, and then optimizer.zero_grad() to reset the gradients. This results in final weights of approximately $[19.8107, 22.0265]$ .
Since the final weights are different from the initial weights, we return a dictionary with "weight_changed" set to True, along with the initial and final weight values.

Constraints

Accumulate over 4 batches before stepping
Use SGD with lr=0.1
Use MSE loss

Learn this in

PyTorch FundamentalsFoundations PretrainingLLM Study Plan

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy

Concept primerPREVIEW

Background Knowledge

Gradient Accumulation is a technique used in deep learning to simulate a larger batch size when the available GPU memory is limited. This is achieved by accumulating gradients over multiple mini-batches before taking an optimizer step. The key concept here is that the gradients are not zeroed out after each mini-batch, but rather accumulated and then used to update the model parameters after a specified number of mini-batches.

In PyTorch, this can be implemented using the optimizer.step() and optimizer.zero_grad() functions. The optimizer.step() function updates the model parameters based on the accumulated gradients, while the optimizer.zero_grad() function resets the gradients to zero. By controlling when these functions are called, we can accumulate gradients over multiple mini-batches and simulate a larger batch size.

The Mean Squared Error (MSE) loss function is commonly used in regression problems, and it measures the average squared difference between the predicted and actual values. In PyTorch, the MSE loss function can be computed using the nn.MSELoss() function. The backward() function is then used to compute the gradients of the loss with respect to the model parameters.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Initialize the model, optimizer, and loss function
Accumulate gradients over multiple mini-batches by calling backward() after each mini-batch
Update the model parameters using optimizer.step() after a specified number of mini-batches
Reset the gradients to zero using optimizer.zero_grad() after updating the model parameters

Step-by-Step Strategy

To implement the solution, follow these steps:

Initialize the model, optimizer, and loss function
Record the initial model weights
For each mini-batch:

Compute the output of the model
Compute the MSE loss
Call backward() to compute the gradients

After all mini-batches:

Call optimizer.step() to update the model parameters
Call optimizer.zero_grad() to reset the gradients to zero

Record the final model weights
Return the initial and final weights, and a boolean indicating whether the weights changed

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