Skip to content

Quickstart

This page walks through the flagship example: run a parameter sweep and get the results back as a labeled xarray.Dataset.

The decorator flow below runs as examples/parallel_sweep.py in the repository; examples/sweep_to_dataset.py is the same flow written with the class API (MultiRunMetricsWorkflow), which this page introduces further down.

Decorate a function with @mushin.sweep, sweep it over a grid, and get results back as a labeled xarray.Dataset — no subclassing, no callbacks. Whatever the function returns as a dict becomes data variables in the output dataset.

import mushin


@mushin.sweep
def experiment(lr, seed):
    ...  # train, evaluate
    return dict(accuracy=acc)  # returned dict -> dataset variables


ds = experiment.run(
    lr=mushin.multirun([0.01, 0.1, 1.0]),
    seed=mushin.multirun([0, 1, 2]),
)

Prefer pandas? You never have to touch xarray: one call turns the sweep into a tidy long-form table (one row per cell, params and metrics as columns):

df = experiment.workflow.to_dataframe()  # or ds.to_dataframe().reset_index()
#      lr  seed  accuracy
# 0  0.01     0     0.914
# 1  0.01     1     0.907
# ...
df.groupby("lr")["accuracy"].mean()  # plain pandas from here on

The labeled dataset is still the richer object (.sel, .mean("seed"), provenance in attrs, netCDF round-trip) — but it's an option, not a prerequisite.

Need the full tool — .failures, .plot(), provenance, custom to_xarray? Drop to experiment.workflow (the last-run instance), or use the MultiRunMetricsWorkflow class directly (shown next).

Going deeper: the workflow class

For advanced control (custom pre_task, jobs_post_process, subclassing like RobustnessCurve), subclass MultiRunMetricsWorkflow and implement a static task method — this is exactly what @mushin.sweep builds for you:

import torch as tr
from mushin import multirun
from mushin.workflows import MultiRunMetricsWorkflow

LEARNING_RATES = [0.01, 0.1, 1.0]
SEEDS = [0, 1, 2]
POINTS_PER_CLASS = 256


def _make_data(seed: int, n: int = POINTS_PER_CLASS) -> tuple[tr.Tensor, tr.Tensor]:
    g = tr.Generator().manual_seed(seed)
    x0 = tr.randn(n, 2, generator=g) + tr.tensor([2.0, 2.0])
    x1 = tr.randn(n, 2, generator=g) + tr.tensor([-2.0, -2.0])
    x = tr.cat([x0, x1])
    y = tr.cat([tr.zeros(n), tr.ones(n)])
    return x, y


class LRSweep(MultiRunMetricsWorkflow):
    @staticmethod
    def task(lr: float, seed: int) -> dict:
        tr.manual_seed(seed)
        x, y = _make_data(seed)
        model = tr.nn.Linear(2, 1)
        opt = tr.optim.SGD(model.parameters(), lr=lr)
        for _ in range(100):
            opt.zero_grad()
            logits = model(x).squeeze(1)
            loss = tr.nn.functional.binary_cross_entropy_with_logits(logits, y)
            loss.backward()
            opt.step()
        with tr.no_grad():
            preds = (model(x).squeeze(1) > 0).float()
            acc = (preds == y).float().mean().item()
        # returning the dict is what populates the dataset; saving is optional
        result = dict(accuracy=acc)
        tr.save(result, "metrics.pt")
        return result

Run the sweep

Call wf.run(...) with multirun(...) wrapped arguments. Hydra launches one job per combination — 3 learning rates × 3 seeds = 9 runs total.

wf = LRSweep()
wf.run(
    lr=multirun(LEARNING_RATES),
    seed=multirun(SEEDS),
)

Heads up: your task() runs in a per-job directory. If it reads or writes files by relative path, wrap them with mushin.original_cwd() / "..." — see Concepts.

Get results as a dataset

ds = wf.to_xarray()
print(ds)

Expected output:

<xarray.Dataset>
Dimensions:   (lr: 3, seed: 3)
Coordinates:
  * lr        (lr) float64 0.01 0.1 1.0
  * seed      (seed) int64 0 1 2
Data variables:
    accuracy  (lr, seed) float64 ...

From there, standard xarray/pandas operations apply:

# average accuracy across seeds, per learning rate
mean_acc = ds["accuracy"].mean("seed")
print(mean_acc)

# plot
import matplotlib.pyplot as plt

mean_acc.plot.line(x="lr", marker="o")
plt.xscale("log")
plt.savefig("sweep_accuracy.png", dpi=120, bbox_inches="tight")

Run the full example

uv run python examples/sweep_to_dataset.py

Next steps