Register Forward Hook to Capture Activations

Problem Statement

Use a forward hook to capture intermediate activations from a hidden layer.

Background

module.register_forward_hook(hook_fn) lets you inspect outputs of any layer without modifying the model. The hook function receives (module, input, output).

Your Task

The starter code creates a 3-layer model and an activations dict. Use a forward hook to capture the output of the first layer during a forward pass. Store the captured activation in the provided dict and remember to clean up the hook afterward.

Output Format

Returns a dictionary with "activation_shape", "activation_values", and "final_output".

Example

Input

None

Output

{'activation_shape': [1, 4], 'activation_values': [1.4206, 0.7051, 2.2023, 0.8228], 'final_output': [-0.2042, -0.7735]}

Reasoning

Input: none. The seed fixes every weight of the $3 \to 4 \to 2$ network, and the forward pass uses the single row $x = [1, 2, 3]$ , a $1 \times 3$ batch of one sample.
First layer: $h = x W_1^\top + b_1$ , a $1 \times 4$ result. Its first entry is $0.4414 \cdot 1 + 0.4792 \cdot 2 - 0.1353 \cdot 3 + 0.4266 \approx 1.4206$ ; the other three give $0.7051$ , $2.2023$ , $0.8228$ . That row is the captured activation, shape [1, 4].
Rest of the network: ReLU zeroes negatives, but all four entries are positive, so they pass unchanged. The second layer maps them to $y = \mathrm{ReLU}(h) W_2^\top + b_2 = [-0.2042, -0.7735]$ .
Output: the shape, the four rounded activations, and the two final values.
Notice that the activation is taken straight out of the first layer, before ReLU: had any $h_i$ been negative, the captured value would still be negative even though the next layer sees $0$ .

Constraints

Use register_forward_hook
Hook captures output of first Linear layer
Do not modify the model architecture

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

PyTorch provides a powerful feature called hooks, which allow you to inspect and modify the behavior of modules (e.g., layers) in a neural network. There are two primary types of hooks: forward hooks and backward hooks. Forward hooks are used to inspect the output of a module, while backward hooks are used to inspect the gradients of a module.

Understanding Forward Hooks

A forward hook is a function that is called after the forward method of a module has been executed. The hook function receives three arguments: (module, input, output), where module is the module that the hook is registered on, input is the input to the module, and output is the output of the module. This allows you to capture intermediate activations or inspect the output of a specific layer.

Registering Forward Hooks

To register a forward hook on a module, you can use the register_forward_hook method, which takes a hook function as an argument. The hook function will be called after the forward method of the module has been executed, allowing you to inspect the output of the module. This is useful for debugging, visualization, or capturing intermediate activations.

Algorithm/Approach

The general approach to solving this problem involves:

Creating a PyTorch model with multiple layers
Registering a forward hook on a specific layer to capture its output
Running input through the model to trigger the hook function
Storing and returning the captured activation and final output

Step-by-Step Strategy

To solve this problem, follow these steps:

Seed the random number generator: Use torch.manual_seed to ensure reproducibility.
Create a PyTorch model: Define a model with multiple layers using nn.Sequential.
Register a forward hook: Use register_forward_hook to register a hook function on the first Linear layer.
Define the hook function: Create a function that captures the output of the module and stores it in a variable.
Run input through the model: Use the forward method to run input through the model and trigger the hook function.
Return the captured activation and final output: Return a dictionary with the shape and values of the captured activation, as well as the final output of the model.

Common Pitfalls

Forgetting to seed the random number generator, which can lead to non-reproducible results.
Registering the hook on the wrong layer or module.
Failing to store the captured activation and final output correctly.

Time & Space Complexity

The time complexity of this solution is O(1), since we are only running a single input through the model. The space complexity is also O(1), since we are only storing a single activation and final output. However, the space complexity of the model itself will depend on the number of layers and parameters.

solution.py

Python 3.13

Test results