TorchHook 0.1.0

TouchHook: A PyTorch hook management library

TorchHook

TorchHook is a library for managing PyTorch model hooks, providing convenient interfaces to capture feature maps and debug models.

Installation

pip install torchhook

Usage Example

import torch
import torch.nn as nn
from torchhook import HookManager

# Define a simple model
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(16 * 30 * 30, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Initialize model and HookManager
model = MyModel()
hook_manager = HookManager(model)

# Register hooks using layer_name (recommended for simplicity)
hook_manager.register_forward_hook(layer_name="conv1")

# Register hooks using layer object (automatically named as: ClassName+Index)
hook_manager.register_forward_hook(layer=model.relu)

# Register hooks with a custom name (useful for distinguishing hooks when debugging)
hook_manager.register_forward_hook('CustomName', layer=model.fc)

# Run the model
for _ in range(5):
    # Generate random input data
    input_tensor = torch.randn(2, 3, 32, 32)
    output = model(input_tensor)

# Print HookManager information
print(hook_manager)
print("Current keys:", hook_manager.get_keys())  # Get all registered hook names

# Get intermediate results (feature maps)
print("\nconv1:", hook_manager.get_features('conv1')[0].shape)  # Feature map of conv1
print("   fc:", hook_manager.get_features('CustomName')[0].shape)  # Feature map of fc

# Get all feature maps
all_features = hook_manager.get_all()

# Concatenate feature maps for each layer (may cause memory overflow if data is too large)
concatenated_features = {key: torch.cat(features, dim=0) for key, features in all_features.items()}

# Compute mean and standard deviation
stats = {key: (torch.mean(value), torch.std(value)) for key, value in concatenated_features.items()}

# Print results
print("\nMean and Std of features:")
for key, (mean, std) in stats.items():
    print(f"Layer: {key}, Mean: {mean.item():.4f}, Std: {std.item():.4f}")

# Clear hooks and features
hook_manager.clear_hooks()
hook_manager.clear_features()

Example Output:

Model: MyModel
Layer Name                    Feature Count       Feature Shape                 
--------------------------------------------------------------------------------
conv1                         5                   (2, 16, 30, 30)               
ReLU_0                        5                   (2, 16, 30, 30)               
CustomName                    5                   (2, 10)                       
--------------------------------------------------------------------------------
Current keys: ['conv1', 'ReLU_0', 'CustomName']

conv1: torch.Size([2, 16, 30, 30])
   fc: torch.Size([2, 10])

Mean and Std of features:
Layer: conv1, Mean: -0.0060, Std: 0.5978
Layer: ReLU_0, Mean: 0.2344, Std: 0.3463
Layer: CustomName, Mean: 0.0245, Std: 0.2332

License

MIT License