A library to inspect itermediate layers of PyTorch models.

Last update: Dec 28, 2022

Overview

A library to inspect itermediate layers of PyTorch models.

Why?

It's often the case that we want to inspect intermediate layers of a model without modifying the code e.g. visualize attention matrices of language models, get values from an intermediate layer to feed to another layer, or applying a loss function to intermediate layers.

Install

$ pip install surgeon-pytorch

Usage

Inspect

Given a PyTorch model we can display all layers using get_layers:

import torch
import torch.nn as nn

from surgeon_pytorch import Inspect, get_layers

class SomeModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(5, 3)
        self.layer2 = nn.Linear(3, 2)
        self.layer3 = nn.Linear(2, 1)

    def forward(self, x):
        x1 = self.layer1(x)
        x2 = self.layer2(x1)
        y = self.layer3(x2)
        return y


model = SomeModel()
print(get_layers(model)) # ['layer1', 'layer2', 'layer3']

Then we can wrap our model to be inspected using Inspect and in every forward call the new model we will also output the provided layer outputs (in second return value):

model_wrapped = Inspect(model, layer='layer2')
x = torch.rand(1, 5)
y, x2 = model_wrapped(x)
print(x2) # tensor([[-0.2726,  0.0910]], grad_fn=<AddmmBackward0>)

We can also provide a list of layers:

model_wrapped = Inspect(model, layer=['layer1', 'layer2'])
x = torch.rand(1, 5)
y, [x1, x2] = model_wrapped(x)
print(x1) # tensor([[ 0.1739,  0.3844, -0.4724]], grad_fn=<AddmmBackward0>)
print(x2) # tensor([[-0.2238,  0.0107]], grad_fn=<AddmmBackward0>)

Or a dictionary to get named outputs:

model_wrapped = Inspect(model, layer={'x1': 'layer1', 'x2': 'layer2'})
x = torch.rand(1, 5)
y, layers = model_wrapped(x)
print(layers)
"""
{
    'x1': tensor([[ 0.3707,  0.6584, -0.2970]], grad_fn=<AddmmBackward0>),
    'x2': tensor([[-0.1953, -0.3408]], grad_fn=<AddmmBackward0>)
}
"""

TODO

add extract function to get intermediate block

Comments

Use one backbone with different heads

Is it possible to save the results from the backbone and apply them on the heads of the all the other models. My goal was to try to save time by avoiding repeating the backbone part. Instead of running the 3 complete models (left), only run the backbone 1 time and switch only the heads for the 3 models (right), therefore not repeating executing the backbone every time in yolov5 model.

Thank you for the help!
question

opened by brunopatricio2012 4
Support for DataParallel?

Hi, I noticed that the current version does not support parallel models (at least those created using torch.nn.DataParallel) since the forward hook does not differentiate between the different copies of the model and a model wrapped with Inspect will just return the intermediate features of the last copy of the parallelized model to run.

Are you planning on fixing this issue/supporting this use case?

opened by zimmerrol 1

A library to inspect itermediate layers of PyTorch models.

Related tags

Overview

Why?

Install

Usage

Inspect

TODO

You might also like...

pyhsmm - library for approximate unsupervised inference in Bayesian Hidden Markov Models (HMMs) and explicit-duration Hidden semi-Markov Models (HSMMs), focusing on the Bayesian Nonparametric extensions, the HDP-HMM and HDP-HSMM, mostly with weak-limit approximations.

PyTorch implementation and pretrained models for XCiT models. See XCiT: Cross-Covariance Image Transformer

TorchGeo is a PyTorch domain library, similar to torchvision, that provides datasets, transforms, samplers, and pre-trained models specific to geospatial data.

Pytorch library for end-to-end transformer models training and serving

This repository provides an efficient PyTorch-based library for training deep models.

TorchMultimodal is a PyTorch library for training state-of-the-art multimodal multi-task models at scale.

Objective of the repository is to learn and build machine learning models using Pytorch. 30DaysofML Using Pytorch

Pretrained SOTA Deep Learning models, callbacks and more for research and production with PyTorch Lightning and PyTorch

Comments

Use one backbone with different heads

Support for DataParallel?

Releases(0.0.4)

0.0.4(May 3, 2022)

0.0.3(May 3, 2022)

0.0.2(Apr 21, 2022)

0.0.1(Apr 21, 2022)

Owner

archinet.ai

Unsupervised Learning of Probably Symmetric Deformable 3D Objects from Images in the Wild

Implementation of fast algorithms for Maximum Spanning Tree (MST) parsing that includes fast ArcMax+Reweighting+Tarjan algorithm for single-root dependency parsing.

(CVPR 2022) Pytorch implementation of "Self-supervised transformers for unsupervised object discovery using normalized cut"

The (Official) PyTorch Implementation of the paper "Deep Extraction of Manga Structural Lines"

Manipulation OpenAI Gym environments to simulate robots at the STARS lab

Categorical Depth Distribution Network for Monocular 3D Object Detection

Code for the paper "Regularizing Variational Autoencoder with Diversity and Uncertainty Awareness"

Implementation of "Meta-rPPG: Remote Heart Rate Estimation Using a Transductive Meta-Learner"

PyTorch implementation of Value Iteration Networks (VIN): Clean, Simple and Modular. Visualization in Visdom.

Official implementation of our neural-network-based fast diffuse room impulse response generator (FAST-RIR)

Real-Time-Student-Attendence-System - Real Time Student Attendence System

[AAAI 2022] Sparse Structure Learning via Graph Neural Networks for Inductive Document Classification

Interpretation of T cell states using reference single-cell atlases

Artificial Intelligence search algorithm base on Pacman

This is the repo for the paper "Improving the Accuracy-Memory Trade-Off of Random Forests Via Leaf-Refinement".

Embeds a story into a music playlist by sorting the playlist so that the order of the music follows a narrative arc.

Breast Cancer Detection 🔬 ITI "AI_Pro" Graduation Project

PointCNN: Convolution On X-Transformed Points (NeurIPS 2018)

Official code release for "GRAF: Generative Radiance Fields for 3D-Aware Image Synthesis"

Propagate Yourself: Exploring Pixel-Level Consistency for Unsupervised Visual Representation Learning, CVPR 2021