Loss Rider

Finally, a Python plotting library that can (only) output Line Rider maps!

ML practitioners can experience gradient descent like never before!

With support for all important features of a line graph.

And don't forget interactive plotting in Jupyter Notebooks!

The above plots all use data from the Unit-Scaled Maximal Update Parameterization paper which proposes a more usable version of μP.

Installation

pip install lossrider

Usage

import pandas as pd
from lossrider import lossrider

# Load a csv that contains columns named "Validation Loss", "Run Count" and "model_type"
data = pd.read_csv("./_data/sweep_df.csv")

# Plot it!
lossrider(
    data,
    x="Run Count", 
    y="Validation Loss",
    hue="model_type",
    xlim=(0.6, 340),
    ylim=(3.2, 3.8),
    xticks=(1, 10, 100), 
    yticks=[x/10 for x in range(32, 39)],
    width=1000, height=500, fontsize=30,
    logx=True, grid=False,
    legend=True, legend_loc=(.65, 1),
    outfile='maps/sweep_strategies',
)

The above produces the below plot

Related Work

The team at wandb claim to have had this idea already.

Ge. et al. have definitely already published in this area. We note the unique contributions of our work are:

• successfully scaling beyond toy datasets such as MNIST, using models with up to 1B parameters
• evaluation on a variety of down-stream tasks
• cool gifs?