TorchLens: A Python package for extracting and visualizing hidden activations of PyTorch models.

Area of Science:

• Artificial Intelligence (AI)
• Computational Neuroscience
• Machine Learning

Background:

• Deep neural networks (DNNs) are crucial for AI and serve as models for biological neural networks.
• Understanding DNN internal representations is vital for both AI and neuroscience research.
• Existing methods lack comprehensive extraction and characterization of DNN internal operations.

Approach:

• Introduced TorchLens, an open-source Python package for PyTorch models.
• TorchLens exhaustively extracts all intermediate operation results, not just module outputs.
• It visualizes the complete computational graph with metadata for analysis.

Key Points:

• TorchLens provides a built-in validation procedure to ensure activation accuracy.
• It automatically applies to any PyTorch model, including complex architectures (recurrent, conditional, branching).
• Minimal code integration is required, facilitating use in existing pipelines and education.

Conclusions:

• TorchLens enables easy and exhaustive extraction and characterization of DNN internal operations.
• This tool aids researchers in AI and neuroscience in understanding DNN internal representations.
• It supports the evaluation of DNNs as computational models of the brain.