Awesome
ProtoFlow: An Invertible Prototypical Neural Network
This repository contains code for the paper, "This Probably Looks Exactly Like That: An Invertible Prototypical Network," which was accepted to ECCV 2024. The proposed architecture, ProtoFlow, represents prototypical distributions as Gaussians in the latent space of a normalizing flow. The approach enables rich interpretation, effective uncertainty estimation, and a promising research path forward for intrinsically interpretable neural networks.
Abstract
We combine concept-based neural networks with generative, flow-based classifiers into a novel, intrinsically explainable, exactly invertible approach to supervised learning. Prototypical neural networks, a type of concept-based neural network, represent an exciting way forward in realizing human-comprehensible machine learning without concept annotations, but a human-machine semantic gap continues to haunt current approaches. We find that reliance on indirect interpretation functions for prototypical explanations imposes a severe limit on prototypes' informative power. From this, we posit that invertibly learning prototypes as distributions over the latent space provides more robust, expressive, and interpretable modeling. We propose one such model, called ProtoFlow, by composing a normalizing flow with Gaussian mixture models. ProtoFlow (1) sets a new state-of-the-art in joint generative and predictive modeling and (2) achieves predictive performance comparable to existing prototypical neural networks while enabling richer interpretation.
Installation
Install the requirements in requirements.txt
as follows:
pip install requirements.txt
Alternatively, the exact environment that was used in this research can be reproduced using conda. After installing
conda, create a new environment using the provided environment.yml
:
conda env create -f environment.yml
Training
To train an instance of ProtoFlow, the train.py
script should be used. Run python train.py --help
for usage details.
You will want the DenseFlow pretrained checkpoints to initialize the model, which can be downloaded following the
instructions here.
For example, ProtoFlow can be trained on CIFAR-10 as follows:
# Optionally enter your dataset root here
#export DATASET_ROOT='/mnt/data/ml_datasets/'
python train.py --flow_ckpt checkpoints/denseflow/imn32/imagenet32/ \
--img_size 32 \
--dataset cifar10 \
--extra my_test_run \
-e 10 \
--batch_steps 32 \
--batch_size 256 \
--trainable all \
--lr 2e-4 \
--gmm_lr 2e-3 \
--consistency_loss \
--protos_per_class 5 \
--elbo_loss2
To run using PyTorch DDP (distributed/parallel training), you can use the following:
torchrun --nproc_per_node=2 train.py ...
Testing
To train an instance of ProtoFlow, the test.py
script should be used. Run python test.py --help
for usage details.
A pre-trained model can be downloaded (see proceeding section) and be evaluated using this script. For example:
# Optionally enter your dataset root here
#export DATASET_ROOT='/mnt/data/ml_datasets/'
python test.py --resume checkpoints/cifar10/checkpoint.pt \
--tta \
--tta_num 5 \
--num_samples 5 \
--proto_scores
If you run out of GPU VRAM, adjust the --batch_size
.
Pre-trained Models
All checkpoints and configurations (including hyperparameters) for trained models are available here.
License
This repository is distributed under the GNU GPL v2.0 License.
This repository contains code from the following projects:
kmeans_pytorch
(MIT License):./protoflow/kmeans/
DenseFlow
(GNU GPL v2.0):./denseflow/
and./experiments/
gmm-torch
(MIT License):./protoflow/gmm.py
Citation
@inproceedings{protoflowECCV2024,
author = {Carmichael, Zachariah and
Redgrave, Timothy and
Gonzalez Cedre, Daniel and
Scheirer, Walter J.},
title = {This Probably Looks Exactly Like That: An Invertible Prototypical Network},
booktitle = {European Conference on Computer Vision},
year = {2024},
publisher = {Springer Nature},
}