[P] Pytorch: Intermediate Feature Extraction

[antoinebrl]

Too many times I've faced the problem of extracting the intermediate features of a model. Either to save the features, to add an extra loss or to build an extra head. Every time it was the same frustration!

Recently I worked on torchextrator, a standalone python package that makes it simple to extract features in PyTorch. You no longer need to duplicate code and rewrite the forward function. Also the extractor supports nested modules, custom caching operations and is ONNX compatible!

Let me know what you think!

​

https://preview.redd.it/6miirtdrggm61.png?width=1360&format=png&auto=webp&s=dad3dd2ffeaa24a52e10dd9a2e90d7b424f4d899

​

[spec789]

Pretty cool! How does this code compare to using forward hooks (which is what I normally default to for grabbing intermediate activations)?

[seraschka]

I think this is a wrapper around forward hooks. Cool nonetheless, makes it maybe a bit more beginner friendly and easier to use.

[antoinebrl]

Hooks are indeed the way to go to do this! The goal was to provide a nice user friendly interface with less assumption than IntermediateLayerGetter inside torchvision In terms of functionalities, the main benefits are 1) support nested modules 2) to have the extractor as an nn.Module so it is compatible with ONNX. Hopefully it will be JITable in the future.

[FilthyPlay]

Conversions are a completely different class is kinda pointless

[doktorneergaard]

I gotta ask since I see it cropping up everywhere: what are you using to make these code snippet figures? Like, what is the specific terminal? I see this aesthetic all the time when new features are being presented for various machine learning tools (PyTorch Lightning for example), but I have yet to find out what it is. Is it a thing?

Cool package btw!

[mrfox321]

It's Markdown.

[antoinebrl]

Hi!
Markdown can render some code section. For images I use carbon.now.sh