torchextractor: PyTorch Intermediate Feature Extraction
Introduction
Too many times some model definitions get remorselessly copy-pasted just because the forward function does not return what the person expects. You provide module names and torchextractor takes care of the extraction for you.It's never been easier to extract feature, add an extra loss or plug another head to a network. Ler us know what amazing things you build with torchextractor!
Installation
pip install torchextractor # stable
pip install git+https://github.com/antoinebrl/torchextractor.git # latest
Requirements:
- Python >= 3.6+
- torch >= 1.4.0
Usage
import torch
import torchvision
import torchextractor as tx
model = torchvision.models.resnet18(pretrained=True)
model = tx.Extractor(model, ["layer1", "layer2", "layer3", "layer4"])
dummy_input = torch.rand(7, 3, 224, 224)
model_output, features = model(dummy_input)
feature_shapes = {name: f.shape for name, f in features.items()}
print(feature_shapes)
# {
# 'layer1': torch.Size([1, 64, 56, 56]),
# 'layer2': torch.Size([1, 128, 28, 28]),
# 'layer3': torch.Size([1, 256, 14, 14]),
# 'layer4': torch.Size([1, 512, 7, 7]),
# }
FAQ
• How do I know the names of the modules?
You can print all module names like this:
tx.list_module_names(model)
# OR
for name, module in model.named_modules():
print(name)
• Why do some operations not get listed?
It is not possible to add hooks if operations are not defined as modules. Therefore, F.relu cannot be captured but nn.Relu() can.
• How can I avoid listing all relevant modules?
You can specify a custom filtering function to hook the relevant modules:
# Hook everything !
module_filter_fn = lambda module, name: True
# Capture of all modules inside first layer
module_filter_fn = lambda module, name: name.startswith("layer1")
# Focus on all convolutions
module_filter_fn = lambda module, name: isinstance(module, torch.nn.Conv2d)
model = tx.Extractor(model, module_filter_fn=module_filter_fn)
• Is it compatible with ONNX?
tx.Extractor is compatible with ONNX! This means you can also access intermediate features maps after the export.
Pro-tip: name the output nodes by using output_names when calling torch.onnx.export.
• Is it compatible with TorchScript?
Not yet, but we are working on it. Compiling registered hook of a module was just recently added in PyTorch v1.8.0.
• "One more thing!"
By default we capture the latest output of the relevant modules, but you can specify your own custom operations.
For example, to accumulate features over 10 forward passes you can do the following:
import torch
import torchvision
import torchextractor as tx
model = torchvision.models.resnet18(pretrained=True)
def capture_fn(module, input, output, module_name, feature_maps):
if module_name not in feature_maps:
feature_maps[module_name] = []
feature_maps[module_name].append(output)
extractor = tx.Extractor(model, ["layer3", "layer4"], capture_fn=capture_fn)
for i in range(20):
for i in range(10):
x = torch.rand(7, 3, 224, 224)
model(x)
feature_maps = extractor.collect()
# Do your stuffs here
# Discard collected elements
extractor.clear_placeholder()
Contributing
All feedbacks and contributions are welcomed. Feel free to report an issue or to create a pull request!
If you want to get hands-on:
- (Fork and) clone the repo.
- Create a virtual environment:
virtualenv -p python3 .venv && source .venv/bin/activate - Install dependencies:
pip install -r requirements.txt && pip install -r requirements-dev.txt - Hook auto-formatting tools:
pre-commit install - Hack as much as you want!
- Run tests:
python -m unittest discover -vs ./tests/ - Share your work and create a pull request.
To Build documentation:
cd docs
pip install requirements.txt
make html
139 Jan 1, 2023
49 Nov 21, 2022
146 Nov 29, 2022
62 Dec 7, 2022
186 Jan 4, 2023
3.4k Jan 4, 2023
69 Nov 15, 2022
45 Dec 21, 2022
13 Jan 6, 2023
41 Dec 26, 2022
594 Dec 31, 2022
1k Dec 30, 2022
42 Nov 25, 2022
13k Jan 09, 2023
1 Dec 18, 2021
180 Jan 07, 2023
48 Sep 30, 2022
3 Jan 18, 2022
5 Dec 13, 2021
29 Nov 25, 2022
64 Nov 11, 2022
2 Jul 12, 2022
15 Dec 02, 2022
374 Dec 23, 2022
3 Jan 08, 2022
28 Jan 04, 2023
633 Jan 04, 2023
14 Oct 21, 2022
464 Dec 27, 2022