Extracting and visualizing hidden activations and computational graphs of PyTorch models with TorchLens

Area of Science:

• Artificial Intelligence
• 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.
• Current methods lack comprehensive extraction and characterization of DNN internal operations.

Purpose of the Study:

• Introduce TorchLens, an open-source Python package for PyTorch.
• Enable exhaustive extraction and characterization of hidden-layer activations in DNNs.
• Facilitate deeper understanding of DNN internal computations and representations.

Main Methods:

• Developed TorchLens, a Python package for PyTorch models.
• Implemented exhaustive extraction of all intermediate operation results.
• Integrated intuitive visualization of computational graphs and metadata.
• Included algorithmic validation for activation accuracy.
• Ensured automatic applicability to diverse PyTorch model architectures.

Main Results:

• TorchLens exhaustively extracts all intermediate operations, providing a full computational graph record.
• The package visualizes the computational graph with metadata for analysis.
• Built-in validation ensures the accuracy of extracted activations.
• TorchLens seamlessly applies to various PyTorch models, including complex architectures.

Conclusions:

• TorchLens simplifies the analysis of DNN internal representations and operations.
• It serves as a valuable tool for AI and neuroscience researchers.
• The package facilitates model development, analysis, and education in deep learning.