Visualizing Torch Landscape

---

Overview

Visualizing Torch Landscape is a Python library for visualizing the loss landscapes of PyTorch models, with a focus on neural networks and GNNs. It provides tools to:

• Compute and visualize 2D/3D loss surfaces around trained model parameters.
• Project training trajectories onto these landscapes.
• Analyze directions in parameter space using random, PCA, SVD, and learnable methods.
• Support both linear and nonlinear subspace visualizations.

Features

• Flexible Direction Calculation: Random, PCA, SVD, and learnable (autoencoder-based) directions.
• Trajectory Projection: Visualize optimizer paths in the loss landscape.
• Parallelized Computation: Fast surface calculation with progress bars.
• Interactive Visualizations: 2D and 3D plots using Plotly.
• Extensible API: Modular design for custom directions, subspaces, and visualizations.

Installation

# 1. Create a new environment (recommended)
conda create -n torch_landscape python=3.10
conda activate torch_landscape

# 2. Install PyTorch and CUDA (if needed)
conda install pytorch cudatoolkit=11.3 -c pytorch

# 3. Clone the repository
git clone https://github.com/SamirMoustafa/GNNs-loss-landscape
cd GNNs-loss-landscape

# 4. Install the package
pip install -e .

Repository Structure

torch_landscape/
    directions.py            # Methods for direction finding (Random, PCA, SVD, Learnable, etc.)
    landscape_linear.py      # Linear loss surface calculation
    landscape_subspace.py    # Subspace-based loss surface calculation
    trajectory.py            # Project and analyze optimizer trajectories
    visualize.py             # 2D/3D visualization with Plotly
    visualize_options.py     # Visualization configuration options
    utils.py                 # Utility functions (seeding, parameter cloning, etc.)
    subspace.py              # Linear and nonlinear subspace abstractions
    parallel_progress_bar.py # Parallelization with progress bars
    names.py                 # Axis/component naming constants
test/
    unit/                   # Unit tests for all major modules
requirements.txt
requirements_dev.txt
setup.py
README.md

Quick Start

from torch_landscape.directions import PcaDirections
from torch_landscape.landscape_linear import LinearLandscapeCalculator
from torch_landscape.trajectory import TrajectoryCalculator
from torch_landscape.utils import clone_parameters
from torch_landscape.visualize import VisualizationData, Plotly2dVisualization
from torch_landscape.visualize_options import VisualizationOptions

# 1. Collect intermediate model parameters during training
intermediate_results = []
for ... in training_loop:
    intermediate_results.append(clone_parameters(model.parameters()))

# 2. Define an evaluation function
def evaluate_func(model, test_loader):
    device = next(model.parameters()).device
    model.eval()
    test_loss = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += torch.nn.functional.nll_loss(output, target, reduction="sum").item()
    return test_loss / len(test_loader.dataset)

# 3. Calculate directions (e.g., PCA)
directions = PcaDirections(optimized_parameters, intermediate_results).calculate_directions()

# 4. Project trajectory
trajectory = TrajectoryCalculator(optimized_parameters, directions).project_disregard_z(intermediate_results)

# 5. Set visualization options
options = VisualizationOptions(num_points=10)
trajectory.set_range_to_fit_trajectory(options)

# 6. Calculate loss surface
landscape_calculator = LinearLandscapeCalculator(optimized_parameters, directions, options=options)
landscape = landscape_calculator.calculate_loss_surface_data_model(model, lambda: evaluate_func(model, data))

# 7. Visualize
Plotly2dVisualization(options).plot(VisualizationData(landscape, trajectory), "my_landscape_surface", file_extension="pdf")

Testing

To run the unit tests:

pip install -r requirements_dev.txt
pytest

Contributing

Contributions are welcome! Please open issues or pull requests for bug fixes, new features, or documentation improvements.

License

GPL-2.0

---

For more details, see the docstrings in each module and the unit tests in test/unit/.

To cite this project

If you use this codebase in your research, please cite it as follows:

@inproceedings{Moustafa2025visualization,
  title         =   {Visualization and Analysis of the Loss Landscape in Graph Neural Networks},
  author        =   {Moustafa, Samir and Kummer, Lorenz and Fetzel, Simon and Kriege, Nils and Gansterer, Wilfried},
  booktitle     =   {Neural Networks for Graphs and Beyond Special Session of the 34th International Conference on Artificial Neural Networks (ICANN 2025)},
  year          =   {2025},
  organization  =   {Springer}
}