loss-landscape-anim

To project the high-dimensional optimization path into a 2D plane, you can specify a reduction_method.

• "pca" (Default): Uses Principal Component Analysis to find the two orthogonal directions with the most variance. This is recommended for showing the most motion in 2D.
• "random": Picks two random directions. Note that in high-dimensional space, these are likely to be orthogonal and may result in a visualization that looks like a random walk.
• "custom": Allows you to provide two specific directions via the custom_directions parameter.

# Using random directions
loss_landscape_anim(n_epochs=300, load_model=False, reduction_method="random")

# Using custom directions
import numpy as np

n_params = ... # number of parameters your model has
u_gen = np.random.normal(size=n_params)
u = u_gen / np.linalg.norm(u_gen)
v_gen = np.random.normal(size=n_params)
v = v_gen / np.linalg.norm(v_gen)

loss_landscape_anim(
    n_epochs=300, 
    load_model=False, 
    reduction_method="custom", 
    custom_directions=(u, v)
)

To use your own PyTorch model for loss landscape analysis, inherit from GenericModel. Your implementation should follow the PyTorch Lightning style. GenericModel provides essential utilities for flattening and unflattening model parameters, which is required for tracking the optimization trajectory.

Key requirements for your subclass:

1. Implement forward(self, x) for the model's forward pass.
2. Implement loss_fn(self, y_pred, y) to define the loss calculation.
3. Implement training_step(self, batch, batch_idx) to handle the training logic. Crucially, this method must call self.get_flat_params() and include the resulting tensor in the dictionary returned by the method to ensure the optimization path is recorded.
4. Implement on_train_epoch_end(self) to append the epoch's results to self.optim_path.

from loss_landscape_anim.model import GenericModel
import torch
import torch.nn as nn

class MyCustomModel(GenericModel):
    def __init__(self, learning_rate, num_classes, ...):
        super().__init__(optimizer="adam", learning_rate=learning_rate, num_classes=num_classes)
        self.layer = nn.Linear(10, num_classes)

    def forward(self, x):
        return self.layer(x)

    def loss_fn(self, y_pred, y):
        return torch.nn.functional.cross_entropy(y_pred, y)

    def training_step(self, batch, batch_idx):
        X, y = batch
        y_pred = self(X)
        flat_w = self.get_flat_params()
        loss = self.loss_fn(y_pred, y)
        # ... log metrics ...
        return {"loss": loss, "flat_w": flat_w}

    def on_train_epoch_end(self):
        self.optim_path.append(self.training_step_outputs[-1])

To visualize your own models and data, pass your model and a PyTorch Lightning DataModule to loss_landscape_anim.

Key parameters include:

• n_epochs: Number of training epochs.
• model: Your model instance.
• datamodule: Your PyTorch Lightning DataModule instance.
• optimizer: The optimizer to use (e.g., "adam", "sgd", "adagrad", "rmsprop").
• giffps: Frames per second for the output GIF.
• seed: Seed for reproducible results.
• load_model: Boolean indicating whether to load a pre-trained model.
• output_to_file: Boolean to save the animation to a file.
• return_data: Boolean to return optim_path, loss_steps, and accu_steps.
• gpus: Number of GPUs to use for training.

bs = 16
lr = 1e-3
datamodule = MNISTDataModule(batch_size=bs, n_examples=3000)
model = LeNet(learning_rate=lr)

optim_path, loss_steps, accu_steps = loss_landscape_anim(
    n_epochs=10,
    model=model,
    datamodule=datamodule,
    optimizer="adam",
    giffps=15,
    seed=SEED,
    load_model=False,
    output_to_file=True,
    return_data=True,
    gpus=1
)

The loss_landscape_anim function is the primary entrypoint for generating an animation of a neural network's optimization trajectory within its loss landscape. It automates the entire pipeline: training a model (or loading an existing one), recording parameter changes, performing dimensionality reduction (PCA, random, or custom), calculating a 2D loss grid, and rendering a GIF.

Workflow

1. Data & Model: Provide a PyTorch Lightning datamodule and a model. If omitted, it defaults to a SpiralsDataModule and a simple MLP.
2. Training: If load_model=False, the function trains the model for n_epochs and saves the checkpoint.
3. Reduction: Projects high-dimensional parameters into 2D space using the specified reduction_method.
4. Visualization: Generates a contour plot animation showing the trajectory moving through the loss landscape.

Key Arguments

• n_epochs: Number of training epochs.
• datamodule: A PyTorch Lightning datamodule instance.
• model: A PyTorch model instance.
• optimizer: String name of the optimizer (e.g., `

The internal function _animate_decision_area animates how a model's decision boundaries change over training epochs. It updates the model parameters using model.init_from_flat_params(W) at each step and redraws the contour plot and data points.

Arguments:

• model: The neural network model. Must implement init_from_flat_params(W).
• X_train: Training features (2D array).
• y_train: Training labels.
• steps: A list of tuples where each tuple is (W, loss, acc). W is the flattened parameter vector.
• giffps: Frames per second for the output.
• write2gif: (Optional) If True, saves the animation as a GIF. Defaults to False.
• file: (Optional) Base filename for the GIF. Defaults to "nn_decision".

# steps should be a list of (flattened_weights, loss, accuracy)
_animate_decision_area(
    model=my_model,
    X_train=X_train,
    y_train=y_train,
    steps=training_history,
    giffps=20,
    write2gif=True,
    file="my_decision_boundary_anim"
)

The LossGrid class computes the loss values for a 2D grid of weight vectors centered around an optimization endpoint. This is used to generate contour plots of the loss landscape.

Key Arguments:

• optim_path: Full-dimensional flattened parameters from training.
• model: The model instance. Note: The model must implement init_from_flat_params(weights) and loss_fn(y_pred, y).
• data: A tuple (X, y) used for loss evaluation.
• path_2d: The 2D coordinates of the optimization path.
• directions: A tuple of the two 2D axes (e.g., from DimReduction.reduce()).
• res: Grid resolution (defaults to 50).
• filepath: Path to save/load the computed grid via pickle.

Workflow:

1. Initialize LossGrid.
2. If load_grid=False (default), it automatically calls compute_loss_2d to evaluate the model at every grid point.
3. If save_grid=True (default), it pickles the results to filepath for future use.

from loss_landscape_anim.loss_landscape import LossGrid

# Assuming model and data are already defined
grid = LossGrid(
    optim_path=params_list, 
    model=my_model, 
    data=(X_train, y_train), 
    path_2d=reduced_result["path_2d"], 
    directions=(reduced_result["reduced_dirs"][0], reduced_result["reduced_dirs"][1]),
    res=30,
    filepath="./my_loss_grid.p"
)