Feature Extraction with Hooks
Problem Statement
Extract features from multiple layers simultaneously using hooks.
Background
In many applications (e.g., FPN, style transfer), you need features from multiple layers. Hooks let you capture all of them in one forward pass.
Your Task
The starter code creates a 5-layer model with named modules. Register hooks on fc1, fc2, and fc3 to capture their outputs in a single forward pass. Use a pattern that correctly captures the layer name in the hook closure (a common gotcha with closures in loops).
Output Format
Returns a dictionary with "fc1_shape", "fc2_shape", "fc3_shape", and "num_features_captured".
Example
Input
None
Output
{'fc1_shape': [1, 8], 'fc2_shape': [1, 4], 'fc3_shape': [1, 2], 'num_features_captured': 3}Reasoning
- The function
feature_extraction_test()starts by seeding the random number generator withtorch.manual_seed(42)to ensure reproducibility. - A model is created with the specified layers:
"fc1","relu1","fc2","relu2", and"fc3", and forward hooks are registered on"fc1","fc2", and"fc3"to capture their outputs. - The input
[[1.0, 2.0, 3.0, 4.0]]is passed through the model, and the outputs of the hooked layers are captured:"fc1"output shape is because the input has 1 sample and"fc1"has 8 output neurons"fc2"output shape is because the input to"fc2"has 1 sample and"fc2"has 4 output neurons"fc3"output shape is because the input to"fc3"has 1 sample and"fc3"has 2 output neurons
- The function returns a dictionary with the shapes of the captured outputs and the number of layers captured, which is 3.
Constraints
- Register hooks on 3 separate layers
- Capture outputs in a single forward pass
- Remove all hooks after