Torch-Module-Observer

Torch-Module-Observer is a tool that uses the internal torch hook functionnalities to provide an API to observe on a model's nn.Modules.

Here is an example with a ConvObserver, but this tool can be used to observe any module, not only Convs.

Extract features :

First layer

-------------->

Last layer

Installation

pip install torchmo

This package also works with conda and poetry.

Use

from torchmo import ConvObserver

model = ... # nn.Module

observer = ConvObserver(model) # Observe model

features = observer(tensor_imgs) # Forward pass through the observer

# After the forward pass, the observer stored the feature info

out_p = ".cache"
out_p.mkdir(exist_ok=True)

# Save figures
observer.save_figs(cache_p)

# Iterate manually through conv features
for name, output in features:
    # output is a batch of
    # output.shape (B,N,H,W)

    print(output.shape)
    break

Create your own

All you need to do is to inherit the ModuleObserver class.

See ConvObserver to see how to inherit and use.

Contribute

The project is small and some peope may have other use cases.

Poetry is being used to manage the environment.

Feel free to open an issue and/or do a PR when you believe it would help the community.