wandb_plot.py

import sys
import pandas as pd
import numpy as np
import os
import wandb
import matplotlib.pyplot as plt
import json
import seaborn as sns
import json
import re
import matplotlib.pyplot as plt
plt.style.use('https://github.com/dhaitz/matplotlib-stylesheets/raw/master/pacoty.mplstyle')


def smooth(scalars, weight):
    last = scalars[0]  # First value in the plot (first timestep)
    smoothed = list()
    for point in scalars:
        smoothed_val = last * weight + (1 - weight) * point  # Calculate smoothed value
        smoothed.append(smoothed_val)                        # Save it
        last = smoothed_val                                  # Anchor the last smoothed value
    return np.array(smoothed)


def getting_metrics(url, metrics):
    # get results from wandb
    wandb_run = api.run(url)
    return wandb_run.history(keys=[metrics]).to_numpy()


api = wandb.Api(timeout=120)
runs = []

settings = json.load(open(sys.argv[1]))
experiments_id = settings['id']

print("Getting all runs.")

regex = []
for k, v in settings['experiments'].items():
    regex.append(v['regex'])

wandb_runs = api.runs(path=experiments_id, filters={"display_name": {"$regex": "|".join(regex)}})

for run in wandb_runs:
  runs.append((run.name, run.id))

print(f'Retrieved {len(runs)} runs.')

EXPERIMENTS = {}
min_steps = 2**32
for k, v in settings['experiments'].items():
    group = []
    if type(v) == dict:
        print(f'Getting runs for experiment {k}')
        for run in runs:
            name, id = run
            if re.match(v['regex'], name):
                url = experiments_id + f"/{id}"
                arr = getting_metrics(url, settings['metric'])
                group.append(arr)
                min_steps = min(min_steps, arr[:, 0].max().item())
    else:
        ids = v
    EXPERIMENTS[k] = group

fig, ax = plt.subplots(figsize=(8, 8))
for k, v in EXPERIMENTS.items():
    df = pd.concat([pd.DataFrame(run[:, 1], index=run[:, 0]) for run in EXPERIMENTS[k]], axis=1)
    df.sort_index(inplace=True)

    steps = df.index.values
    mean = df.mean(axis=1).backfill().interpolate().to_numpy()
    stderr = df.interpolate().sem(axis=1).backfill().interpolate().to_numpy()

    mean = smooth(mean, 0.99)
    stderr = smooth(stderr, 0.99)

    ax.plot(steps, mean, label=k, linewidth=4, color=settings['experiments'][k]['color'])
    ax.fill_between(steps, mean-stderr, mean+stderr, alpha=0.3, color=settings['experiments'][k]['color'])

ax.set(xlabel='steps', ylabel=settings['metric'])
ax.legend()
fig.savefig(settings['output'])