Register Backward Hook to Capture Gradients

Problem Statement

Use a backward hook to capture gradients flowing through a specific layer.

Background

module.register_full_backward_hook(hook_fn) captures gradients during the backward pass. The hook receives (module, grad_input, grad_output).

Your Task

The starter code creates a model and grad_data dict. Use a backward hook to capture the gradients flowing through the second Linear layer during backpropagation. Store both grad_input and grad_output, and clean up the hook afterward.

Output Format

Returns a dictionary with "grad_output_shape", "grad_input_shapes", and "num_grad_inputs".

Example

Input

None

Output

{'grad_output_shape': [1, 1], 'grad_input_shapes': [[1, 4], [1, 1], [1]], 'num_grad_inputs': 3}

Reasoning

The backward_hook_test function creates a neural network with two linear layers and a ReLU activation function, then registers a full backward hook on the second linear layer.
When the input [[1.0, 2.0, 3.0]] is passed through the network, the output is computed and the mean squared error (MSE) loss is calculated with respect to the target [[1.0]].
During the backward pass, the hook is triggered, capturing the gradients flowing through the second linear layer: grad_input contains the gradients of the loss with respect to the layer's input and parameters, while grad_output contains the gradients of the loss with respect to the layer's output.
The hook stores the shapes of grad_output and grad_input, which are then returned as a dictionary: grad_output has shape $[1, 1]$ , grad_input has shapes $[1, 4]$ , $[1, 1]$ , and $[1]$ (corresponding to the layer's input, weight, and bias, respectively), and there are $3$ non-None entries in grad_input.

Constraints

Use register_full_backward_hook
Capture grad_input and grad_output
Hook on the second Linear layer

Learn this in

PyTorch FundamentalsFoundations

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 primer

Background Knowledge

Introduction to PyTorch Hooks

In PyTorch, hooks are a powerful tool for modifying or inspecting the behavior of modules (e.g., neural network layers) during the forward or backward pass. A hook is essentially a function that is called at a specific point in the execution of a module. There are two types of hooks: forward hooks and backward hooks. Forward hooks are called after the forward pass, while backward hooks are called after the backward pass.

Understanding Backward Hooks

A backward hook is used to capture gradients flowing through a module during the backward pass. The register_full_backward_hook method is used to register a backward hook on a module. The hook function receives three arguments: module, grad_input, and grad_output. module is the module that the hook is registered on, grad_input is the gradient of the loss with respect to the input of the module, and grad_output is the gradient of the loss with respect to the output of the module.

Gradients and Backpropagation

In the context of neural networks, gradients represent the rate of change of the loss with respect to the model's parameters. During the backward pass, PyTorch computes these gradients using the chain rule of calculus. The gradients are used to update the model's parameters during training. By capturing gradients using a backward hook, you can inspect or modify the gradients at a specific point in the network, which can be useful for debugging, visualization, or implementing custom training algorithms.

Algorithm/Approach

The general approach to solving this problem involves:

Creating a PyTorch module (in this case, a sequential neural network)
Registering a backward hook on a specific layer of the module
Running a forward pass and computing the loss
Calling the backward pass to compute gradients
Inspecting the gradients captured by the hook

Step-by-Step Strategy

To implement the solution:

Create a PyTorch module using nn.Sequential.
Register a full backward hook on the second Linear layer using register_full_backward_hook.
Define the hook function to capture and store grad_input and grad_output.
Run a forward pass with a sample input and compute the loss using a loss function (e.g., MSE).
Call the backward pass using loss.backward().
Inspect the gradients captured by the hook and compute the required metrics (e.g., shape of grad_output, shapes of grad_input tensors).

Common Pitfalls

Forgetting to call loss.backward() to trigger the backward pass.
Not handling None entries in grad_input when computing shapes.
Incorrectly indexing the module or layer when registering the hook.

Time & Space Complexity

The time complexity of this solution is dominated by the forward and backward passes, which are $O(n)$ , where $n$ is the number of parameters in the network. The space complexity is $O(n)$ as well, since we need to store the gradients and other intermediate results. However, the exact complexity may vary depending on the specific implementation and the size of the input data.

solution.py

Python 3.13

Test results