Visualization toolkit for neural networks in PyTorch

visualization, machine-learning, deep-learning, cnn, pytorch, neural-networks, interpretability, explainability
pip install flashtorch==0.1.3



PyPI - Python Version Build Status PyPI PyPI - License DOI Say Thanks!

A Python visualization toolkit, built with PyTorch, for neural networks in PyTorch.

Neural networks are often described as "black box". The lack of understanding on how neural networks make predictions enables unpredictable/biased models, causing real harm to society and a loss of trust in AI-assisted systems.

Feature visualization is an area of research, which aims to understand how neural networks perceive images. However, implementing such techniques is often complicated.

FlashTorch was created to solve this problem!

You can apply feature visualization techniques (such as saliency maps and activation maximization) on your model, with as little as a few lines of code.

It is compatible with pre-trained models that come with torchvision, and seamlessly integrates with other custom models built in PyTorch.


Take a look at the quick 3min intro/demo to FlashTorch below!

FlashTorch demo

Want to try?

Head over to example notebooks on Colab!

  • Saliency maps: Saliency map demo

  • Activation maximization: Activation maximization demo



If you are installing FlashTorch for the first time:

$ pip install flashtorch

Or if you are upgrading it:

$ pip install flashtorch -U

API guide

These are currently available modules.

  • flashtorch.utils: some useful utility functions for data handling & transformation
  • flashtorch.utils.imagenet: ImageNetIndex class for easy-ish retrieval of class index
  • flashtorch.saliency.backprop: Backprop class for calculating gradients
  • flashtorch.activmax.gradient_ascent: GradientAscent class for activation maximization

You can inspect each module with Python built-in function help. The output of that is available on Quick API Guide for your convenience.


Use FlashTorch

Below, you can find simple demos to get you started, as well as links to some handy notebooks showing additional examples of using FlashTorch.

Image handling (flashtorch.utils)

Saliency maps (flashtorch.saliency)

Saliency in human visual perception is a subjective quality that makes certain things within the field of view stand out from the rest and grabs our attention.

Saliency maps in computer vision provide indications of the most salient regions within images. By creating a saliency map for neural networks, we can gain some intuition on "where the network is paying the most attention to" in an input image.

Using flashtorch.saliency module, let's visualize image-specific class saliency maps of AlexNet pre-trained on ImageNet classification tasks.

Saliency map of great grey owl in Alexnet

The network is focusing on the sunken eyes and the round head for this owl.

Activation maximization (flashtorch.activmax)

Activation maximization is one form of feature visualization that allows us to visualize what CNN filters are "looking for", by applying each filter to an input image and updating the input image so as to maximize the activation of the filter of interest (i.e. treating it as a gradient ascent task with filter activation values as the loss).

Using flashtorch.activmax module, let's visualize images optimized with filters from VGG16 pre-trained on ImageNet classification tasks.

VGG16 conv5_1 filters

Concepts such as 'eyes' (filter 45) and 'entrances (?)' (filter 271) seem to appear in the conv5_1 layer of VGG16.

Visit the notebook above to see what earlier layers do!

Develop FlashTorch

Here is how to setup a dev environment for FlashTorch.

From the project root:

  1. Create a conda environment.

    $ conda env create -f environment.yml
  2. Activate the environment.

    $ conda activate flashtorch
  3. Install FlashTorch in a development mode.

    $ pip install -e .
  4. Run the linter & test suit.

    $ flake8 flashtorch tests && pytest
  5. Add a kernel to Jupyter notebook.

    $ python -m ipykernel install --user --name flashtorch \
      --display-name <e.g. flashtorch-dev>
  6. Launch Jupyter notebook

    $ jupyter notebook
  7. Open a notebook in the ./examples directory.

  8. From the top menu, Kernel -> Change kernel -> flashtorch-dev

  9. From the top menu, Cell -> Run All

If the test suit runs and all the cells in the notebook execute - congratulations, you're good to go!

If you encounter any problem setting up the dev environment, please open an issue.

How to contribute

Thanks for your interest in contributing!

Please first head over to the Code of Conduct, which helps set the ground rules for participation in communities and helps build a culture of respect.

Next, please make sure that you have a dev environment set up (see the Develop FlashTorch section above).

Still here? Great! There are many ways to contribute to this project. Get started here.


Talks & blog posts




Misa Ogura, & Ravi Jain. (2020, January 2).
MisaOgura/flashtorch: 0.1.2 (Version v0.1.2).


Misa Ogura

Medium | twitter | LinkedIn

R&D Software Engineer @ BBC

Co-founder of Women Driven Development