Spaces:
Sleeping
Sleeping
File size: 2,540 Bytes
adcc0ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from torch import nn
class ActivationExtractor(nn.Module):
"""
A utility class for extracting activations from intermediate layers of a PyTorch model.
This class uses forward hooks to capture the output (activations) of a specified layer during the forward pass of the model.
It is designed to work as a context manager for safe and automatic hook management.
Args:
model: The PyTorch model from which activations are to be extracted.
target_layer: The specific layer iin the model whose activation needs to be captured.
Attributes:
activation: Captures the output (activations) of the target layer during the forward pass.
Usage:
>>> model = torchvision.models.resnet50(pretrained=True)
>>> target_layer = model.layer4[0].conv1
>>> with ActivationExtractor(model, target_layer) as extractor:
... output = model(input_tensor)
... activations = extractor.activation
"""
def __init__(self, model, target_layer):
"""
Initialize the ActivationExtractor with the model and target layer.
Args:
model: the PyTorch model
target_layer: the layer from which activations are to be extracted.
"""
super().__init__()
self.model = model
self.target_layer = target_layer
self.activation = None # Stores activation from target layer
def hook_fn(self, module, input, output):
"""
The forward hook function to capture the output of the target layer.
Args:
module: The target layer (passed automatically by PyTorch)
input: Input to the target layer (not used here but required by PyTorch)
output: Output from the target layer (stored as th activation).
"""
self.activation = output # Save the activation output during the forward pass.
def __enter__(self):
"""
Enter the context and register the forward hook on the target layer.
Returns:
self: The instance of the ActivationExtractor with the hook registered.
"""
# Register the forward hook to capture activations
self.hook = self.target_layer.register_forward_hook(self.hook_fn)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""
Exit the context and remove the forward hook to clean up resources.
Args:
exc_type: Exception type (if any) raised inside the context.
exc_val: Exception value (if any) raised inside the context.
exc_tb: Traceback (if any) raised inside the context.
"""
# Remove the forward hook to avoid memory leaks.
self.hook.remove()
|