Quickstart

This page gives you a practical, minimal path to get WeightsLab running. If you prefer to start from examples, see Examples right after this setup.

Prerequisites

  • Python v3.10+ installed

  • A virtual environment tool like venv or Conda (optional).

  • Your training project available locally.

Install WeightsLab

Create and activate a virtual environment and install WeightsLab.

python -m pip install weightslab

Tip

For reproducible experiments, you can install in a virtual environment with the following command:

# From the repository root
python -m venv .venv

# Windows PowerShell
.\.venv\Scripts\Activate.ps1
# Linux/macOS
# source .venv/bin/activate

Try the bundled example

To see WeightsLab working end to end without writing any code, start one of the bundled examples (–cls, –seg, –det, –2d_det, –3d_det). It run a small bundled experiment:

weightslab start example --cls

Then, in another terminal, launch the UI and open the URL printed by the command:

weightslab start

Local integration in your own Python script (MNIST)

The same MNIST CNN, three ways, each shown as a diff: - lines go away, + lines are what WeightsLab adds. The first tab starts from a plain PyTorch script, the other two from one already wired to an experiment tracker. The Copy button on a diff block drops the - lines and the + markers, so what lands in your clipboard is the runnable WeightsLab version.

Migrating a real codebase? The full guides are From TensorBoard and From Weights & Biases.

Starting from a plain PyTorch loop, with no tracker of any kind. The + lines are everything WeightsLab needs; the - lines are what it replaces.

 1import torch
 2import torch.nn as nn
 3import torch.optim as optim
 4from torchvision import datasets, transforms
 5+
 6+  import weightslab as wl
 7
 8
 9class CNN(nn.Module):
10    def __init__(self):
11        super().__init__()
12        self.net = nn.Sequential(
13            nn.Conv2d(1, 32, 3, padding=1),
14            nn.ReLU(),
15            nn.MaxPool2d(2),
16            nn.Conv2d(32, 64, 3, padding=1),
17            nn.ReLU(),
18            nn.MaxPool2d(2),
19            nn.Flatten(),
20            nn.Linear(64 * 7 * 7, 10),
21        )
22
23    def forward(self, x):
24        return self.net(x)
25
26
27cfg = {
28    "device": "auto",
29    "data_root": "./data",
30    "data": {
31        "train_loader": {
32            "batch_size": 64,
33        }
34    },
35    "optimizer": {
36        "lr": 1e-3,
37    },
38}
39device = "cuda" if torch.cuda.is_available() and cfg["device"] in ["auto", "cuda"] else "cpu"
40
41train_ds = datasets.MNIST(cfg["data_root"], train=True, download=True, transform=transforms.ToTensor())
42-  train_loader = torch.utils.data.DataLoader(train_ds, batch_size=cfg["data"]["train_loader"]["batch_size"], shuffle=True)
43
44model = CNN().to(device)
45optimizer = optim.Adam(model.parameters(), lr=cfg.get("optimizer", {}).get("lr", 1e-3))
46-  loss = nn.CrossEntropyLoss()
47+  loss = nn.CrossEntropyLoss(reduction="none")
48+
49+ # Wrap your objects with WeightsLab to watch and edit them in real time.
50+ ## Wrap the hyperparameters first: the studio edits this dict in place.
51+  hp = wl.watch_or_edit(cfg, flag="hyperparameters")
52+
53+ ## Wrap the model and the optimizer next.
54+  model = wl.watch_or_edit(model, flag="model", device=device)
55+  optimizer = wl.watch_or_edit(optimizer, flag="optimizer")
56+
57+ ## Then the loss. The reduction="none" above is what makes it one value
58+ ## *per sample*, which is what lets the studio sort the grid by loss and
59+ ## take you from a spike in the curve to the images that caused it.
60+  loss = wl.watch_or_edit(
61+      loss,
62+      flag="loss",
63+      signal_name="train/loss",
64+      log=True,
65+  )
66+
67+ ## And the dataset, which comes back as a tracked dataloader.
68+  train_loader = wl.watch_or_edit(
69+      train_ds,
70+      flag="data",
71+      loader_name="train_loader",
72+      batch_size=hp["data"]["train_loader"]["batch_size"],
73+      shuffle=True,
74+      is_training=True,
75+  )
76+
77+ # Finally start the WeightsLab backend and keep it running while you train.
78+  wl.serve(serving_grpc=True, serving_cli=True)
79
80step = 0
81while True:
82+      ## guard_training_context is how pause/resume and the train/test
83+      ## split work -- without it, Play/Pause and the stats misbehave.
84+      with wl.guard_training_context:
85-      inputs, labels = next(iter(train_loader))
86+          inputs, uids, labels = next(train_loader)
87+          inputs, labels = inputs.to(device), labels.to(device)
88+          optimizer.zero_grad()
89+          logits = model(inputs)
90-          loss_per_sample = loss(logits, labels)
91+          loss_per_sample = loss(logits, labels, batch_ids=uids, preds=logits.argmax(1, keepdim=True))
92+          loss_per_sample.mean().backward()
93+          optimizer.step()
94
95    if step % 20 == 0:
96        print(f"Loss: {loss_per_sample.mean().item():.4f}")
97    step += 1

Three things to notice. The tracked train_loader yields (inputs, uids, labels) , those uids are what tie a loss value back to the sample that produced it, which is why they are handed to the loss as batch_ids. The loop is open-ended: you stop it from the studio, not with a step budget. And it does not start on its own , run weightslab start in another terminal and press Play.

SummaryWriter is a file handle you push numbers into. WeightsLab has no equivalent object: you wrap the thing that produces the number once, and it reports itself from then on, so the add_scalar call, and the global step bookkeeping it needs, both leave the loop.

 1import torch
 2import torch.nn as nn
 3import torch.optim as optim
 4-  from torch.utils.tensorboard import SummaryWriter
 5from torchvision import datasets, transforms
 6+  import weightslab as wl
 7
 8
 9class CNN(nn.Module):
10    def __init__(self):
11        super().__init__()
12        self.net = nn.Sequential(
13            nn.Conv2d(1, 32, 3, padding=1),
14            nn.ReLU(),
15            nn.MaxPool2d(2),
16            nn.Conv2d(32, 64, 3, padding=1),
17            nn.ReLU(),
18            nn.MaxPool2d(2),
19            nn.Flatten(),
20            nn.Linear(64 * 7 * 7, 10),
21        )
22
23    def forward(self, x):
24        return self.net(x)
25
26
27cfg = {
28    "device": "auto",
29    "data_root": "./data",
30    "data": {
31        "train_loader": {
32            "batch_size": 64,
33        }
34    },
35    "optimizer": {
36        "lr": 1e-3,
37    },
38}
39device = "cuda" if torch.cuda.is_available() and cfg["device"] in ["auto", "cuda"] else "cpu"
40
41train_ds = datasets.MNIST(cfg["data_root"], train=True, download=True, transform=transforms.ToTensor())
42-  train_loader = torch.utils.data.DataLoader(train_ds, batch_size=cfg["data"]["train_loader"]["batch_size"], shuffle=True)
43
44model = CNN().to(device)
45optimizer = optim.Adam(model.parameters(), lr=cfg.get("optimizer", {}).get("lr", 1e-3))
46-  loss = nn.CrossEntropyLoss()
47+  loss = nn.CrossEntropyLoss(reduction="none")   # one value per sample, not per batch
48-  writer = SummaryWriter(log_dir="./runs/mnist_baseline")
49+
50+ # Wrap your objects with WeightsLab to watch and edit them in real time.
51+ ## Wrap the hyperparameters first
52+  hp = wl.watch_or_edit(cfg, flag="hyperparameters")
53+
54+ ## Wrap the model and optimizer next
55+  model = wl.watch_or_edit(model, flag="model", device=device)
56+  optimizer = wl.watch_or_edit(
57+      optimizer,
58+      flag="optimizer",
59+  )
60+
61+ ## Then wrap the loss and the dataset
62+  loss = wl.watch_or_edit(
63+      loss,
64+      flag="loss",
65+      signal_name="train/loss",
66+      log=True,
67+  )
68+  train_loader = wl.watch_or_edit(
69+      train_ds,
70+      flag="data",
71+      loader_name="train_loader",
72+      batch_size=hp["data"]["train_loader"]["batch_size"],
73+      shuffle=True,
74+      is_training=True,
75+  )
76+
77+ # Finally start the WeightsLab backend and keep it running while you train.
78+  wl.serve(serving_grpc=True, serving_cli=True)
79
80step = 0
81while True:
82+      with wl.guard_training_context:
83-      inputs, labels = next(iter(train_loader))
84+          inputs, uids, labels = next(train_loader)
85+          inputs, labels = inputs.to(device), labels.to(device)
86+          optimizer.zero_grad()
87+          logits = model(inputs)
88-          loss_per_sample = loss(logits, labels)
89+          loss_per_sample = loss(logits, labels, batch_ids=uids, preds=logits.argmax(1, keepdim=True))
90+          loss_per_sample.mean().backward()
91+          optimizer.step()
92-          writer.add_scalar("train/loss", loss_per_sample.mean().item(), step)
93
94    if step % 20 == 0:
95        print(f"Loss: {loss_per_sample.mean().item():.4f}")
96    step += 1
97
98-  writer.close()

The loop body has no reporting code left in it at all. What it gained instead is guard_training_context (pause/resume and the train/test split) and batch_ids=uids (which sample each loss value belongs to).

wandb.log() is a call you make at every point you want a number recorded; a watched object records itself. The other change worth noticing is cfg: wandb.config is frozen once init() returns, whereas watched hyperparameters stay editable from the studio while the run goes on, so read them out of the dict each step rather than caching them in locals.

  1import torch
  2import torch.nn as nn
  3import torch.optim as optim
  4-  import wandb
  5from torchvision import datasets, transforms
  6+  import weightslab as wl
  7
  8
  9class CNN(nn.Module):
 10    def __init__(self):
 11        super().__init__()
 12        self.net = nn.Sequential(
 13            nn.Conv2d(1, 32, 3, padding=1),
 14            nn.ReLU(),
 15            nn.MaxPool2d(2),
 16            nn.Conv2d(32, 64, 3, padding=1),
 17            nn.ReLU(),
 18            nn.MaxPool2d(2),
 19            nn.Flatten(),
 20            nn.Linear(64 * 7 * 7, 10),
 21        )
 22
 23    def forward(self, x):
 24        return self.net(x)
 25
 26
 27cfg = {
 28    "device": "auto",
 29    "data_root": "./data",
 30    "data": {
 31        "train_loader": {
 32            "batch_size": 64,
 33        }
 34    },
 35    "optimizer": {
 36        "lr": 1e-3,
 37    },
 38}
 39-  wandb.init(project="mnist", config=cfg)
 40-  cfg = dict(wandb.config)   # a record of the run, frozen from here on
 41device = "cuda" if torch.cuda.is_available() and cfg["device"] in ["auto", "cuda"] else "cpu"
 42
 43train_ds = datasets.MNIST(cfg["data_root"], train=True, download=True, transform=transforms.ToTensor())
 44-  train_loader = torch.utils.data.DataLoader(train_ds, batch_size=cfg["data"]["train_loader"]["batch_size"], shuffle=True)
 45
 46model = CNN().to(device)
 47-  wandb.watch(model, log="all")
 48optimizer = optim.Adam(model.parameters(), lr=cfg.get("optimizer", {}).get("lr", 1e-3))
 49-  loss = nn.CrossEntropyLoss()
 50+  loss = nn.CrossEntropyLoss(reduction="none")   # one value per sample, not per batch
 51+
 52+ # Wrap your objects with WeightsLab to watch and edit them in real time.
 53+ ## wandb.config becomes a live dict -- edit it from the studio mid-run
 54+  hp = wl.watch_or_edit(cfg, flag="hyperparameters")
 55+
 56+ ## wandb.watch(model) becomes a watched model, plus the optimizer W&B never saw
 57+  model = wl.watch_or_edit(model, flag="model", device=device)
 58+  optimizer = wl.watch_or_edit(
 59+      optimizer,
 60+      flag="optimizer",
 61+  )
 62+
 63+ ## wandb.log({"train/loss": ...}) becomes a watched loss that logs itself
 64+  loss = wl.watch_or_edit(
 65+      loss,
 66+      flag="loss",
 67+      signal_name="train/loss",
 68+      log=True,
 69+  )
 70+
 71+ ## and the dataset becomes the tracked table -- no wandb.Table to build
 72+  train_loader = wl.watch_or_edit(
 73+      train_ds,
 74+      flag="data",
 75+      loader_name="train_loader",
 76+      batch_size=hp["data"]["train_loader"]["batch_size"],
 77+      shuffle=True,
 78+      is_training=True,
 79+  )
 80+
 81+ # wandb.init() becomes: start the backend, keep it up while you train.
 82+  wl.serve(serving_grpc=True, serving_cli=True)
 83
 84step = 0
 85while True:
 86+      with wl.guard_training_context:
 87-      inputs, labels = next(iter(train_loader))
 88+          inputs, uids, labels = next(train_loader)
 89+          inputs, labels = inputs.to(device), labels.to(device)
 90+          optimizer.zero_grad()
 91+          logits = model(inputs)
 92-          loss_per_sample = loss(logits, labels)
 93+          loss_per_sample = loss(logits, labels, batch_ids=uids, preds=logits.argmax(1, keepdim=True))
 94+          loss_per_sample.mean().backward()
 95+          optimizer.step()
 96-          wandb.log({"train/loss": loss_per_sample.mean().item()}, step=step)
 97
 98    if step % 20 == 0:
 99        print(f"Loss: {loss_per_sample.mean().item():.4f}")
100    step += 1
101
102-  wandb.finish()

One thing has no equivalent, and is not meant to: there are no sweeps. WeightsLab is built around staying inside one run and steering it, raise the learning rate when the curve flattens, discard the samples poisoning it, keep going. If you need a sweep, keep the tool you sweep with.

Notebook Code with Google Colab

Start by opening this notebook:

Use Weightslab Studio (UI)

For a full visual experiment monitoring workflow (agent, samples, tags, discard/restore, plots), deploy the Weights Studio web app with the bundled CLI.

Without certificates the UI runs unsecured (HTTP, no gRPC auth). Once you have generated them with weightslab se, weightslab start finds them in $WEIGHTSLAB_CERTS_DIR (else ~/.weightslab-certs) and serves HTTPS + gRPC auth automatically, the same rule the training backend applies, so both sides agree:

weightslab se                 # once: generate TLS certificates + a gRPC auth token
weightslab start              # HTTPS + gRPC auth when certs exist, HTTP otherwise
weightslab start --no-certs   # force plain HTTP even when certs exist

Important

When using certs, it is prefered to set manually the WEIGHTSLAB_CERTS_DIR environment variable so the training backend and any new terminal use the same certificates, it is the single source of truth for TLS/auth. Please note that this step has to be done before starting the experiment.

Run weightslab, weightslab help, or weightslab -h to see the banner and the full command reference (se, start, start example ...).

To stop the UI, press Ctrl+C in the terminal running weightslab start.

Prefer a terminal over a browser? weightslab cli opens an interactive console connected to the running experiment (pause/resume, status, evaluate, tag/discard samples, query the agent, …), no UI container required:

weightslab cli

Full reference for both, every weightslab subcommand and every console command, with all flags and defaults, lives in User Commands Reference.

Tip

Let an AI agent integrate WeightsLab for you.

The repository ships with AGENTS.md, a compact context file that gives any AI coding assistant (Claude, Copilot, Cursor, …) a complete picture of the WeightsLab API. Open your training script, attach AGENTS.md as context, and ask:

"Using the context in AGENTS.md, integrate WeightsLab into this training script."

The agent will wire up your model, data loader, loss, and hyperparameters in a few edits, no manual API lookup needed. Otherwise use the Agent Quickstart to connect the integrated OpenCode agent to a running experiment and have it generate code for you from the UI.