Built for AI Engineers working with messy real-world data

Pause training, mine live loss signals to surface mislabels, class imbalance & outliers,
then curate your image, video & LiDAR data, without restarting.

Website · Docs · Demo (Images) · Demo (VLA) · Demo (LiDAR)

Overview

WeightsLab hooks into your existing PyTorch training loop and exposes a live UI where you can inspect per-sample signals, edit the dataset, and steer training. Without restarting.

Weightslab in Motion

Quickstart

1. Install & Launch

pip install weightslab

weightslab start # launch the UI

2. Start in the cloud

OR Wrap your training script locally

# wrap the objects in your training scriptimportweightslabaswl
...
model=wl.watch_or_edit(model, flag='model')
optim=wl.watch_or_edit(optim, flag='opt')
loss=wl.watch_or_edit(loss, flag='signal', name="loss", per_sample=True, log=True)
loader=wl.watch_or_edit(dataset, flag='data', loader_name="train")
...
wl.serve(serving_grpc=True, serving_cli=False)
...

Tip

Quick examples to get started

weightslab start example # classification (default)
weightslab start example --cls # classification
weightslab start example --seg # segmentation
weightslab start example --det # detection
weightslab start example --clus # clustering

Explore our sandbox. For a detailed installation guide and advanced configuration: Documentation.

How can you use it ?

1. Find bad data fast

Pause training mid-run, sort samples by loss, spot mislabels and outliers before they impact your model.

2. Fix and resume without re-starting

Relabel or drop samples live, then continue training from the same checkpoint.

3. Catch model regressions early on

Analyze per-sample loss trajectories to see exactly where the model is struggling.

Resources & Community

Training script with Weightslab - Step-by-Step Integration

1. Add the import at the top of your script:

importweightslabaswl

2. Wrap your parameters, model, optimizer, signals, and dataset:

parameters=wl.watch_or_edit(parameters, flag='hp', ...) # ← WeightsLab monitors your parameters and lets you update them from the UImodel=wl.watch_or_edit(model, flag='model', ...) # ← WeightsLab monitors your model stateoptimizer=wl.watch_or_edit(optim.Adam(...), flag='opt', ...) # ← Tracks optimizer state and lets you update the learning rate from the UItrain_criterion=wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), flag='signal', name="train_loss/sample", per_sample=True, log=True) # ← Wrap and plot your signals on the UItest_criterion=wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), flag='signal', name="test_loss/sample", per_sample=True, log=False) # ← Per-sample only, plot disabledtrain_loader=wl.watch_or_edit(train_dataset, flag='data', loader_name="train_loader", ...) # ← Track your training datasetval_loader=wl.watch_or_edit(val_dataset, flag='data', loader_name="val_loader", ...) # ← Track your validation dataset

3. Run your script, then launch the UI in a separate terminal:

 python train.py
weightslab start

4. Open your browser at the URL printed by weightslab start and inspect your training in real time.

Training script with Weightslab - Full Example

#!/usr/bin/env python3"""Basic PyTorch training script with WeightsLab integration"""importtorchimporttorch.nnasnnimporttorch.optimasoptimfromtorch.utils.dataimportDataLoader, TensorDatasetimportweightslabaswlclassSimpleModel(nn.Module):
def__init__(self, input_shape=10, output_shape=1):
super().__init__()
self.linear=nn.Linear(input_shape, output_shape)
defforward(self, x):
returnself.linear(x)
defcreate_data(n_samples=1000):
X=torch.randn(n_samples, 10)
y=X.sum(dim=1, keepdim=True) +0.1*torch.randn(n_samples, 1)
returnTensorDataset(X, y)
defmain():
parameters=wl.watch_or_edit({}, flag="hyperparameters", poll_interval=1.0) or {}
model=wl.watch_or_edit(SimpleModel(), flag='model')
optimizer=wl.watch_or_edit(optim.Adam(model.parameters(), lr=0.01), flag='optimizer')
criterion=wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), flag="loss", signal_name="train-loss-CE", log=True)
loader=wl.watch_or_edit(create_data(), flag="data", loader_name="loader", batch_size=8, is_training=True)
forepochinrange(parameters.get('n_epochs', 5)):
total_loss=0forbatch_X, batch_yinloader:
predictions=model(batch_X)
loss=criterion(predictions, batch_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss+=loss.item()
# Write the history of these samples every x stepsifmodel.get_age() %100==0:
print(f'Dump signals history and dataframe at age {model.get_age()}')
wl.write_history(
# path=None, # Use root_log_dir by default, filename generated from parameters md5 hashtype_of_history="all",
graph_name=[
'train/clsf_instance',
'val/clsf_instance'
],
# experiment_hash=None, Default is 'last', i.e., current experiment hashsample_id=['11', '29', '28', '27', '22'],
instance_id=[1, 2, 3]
)
# Dump the sample dataframe: all signals plus the loss_shape categorical tag,wl.write_dataframe(
columns=["signals", "tag:loss_shape"],
format='csv'# sample_id=['0', '28']# instance_id=[1, 2],
)
avg_loss=total_loss/len(loader)
print(f"Epoch {epoch+1}/5 - Loss: {avg_loss:.4f}")
print("✅ Training complete!")
if__name__=="__main__":
main()

Migrating from Weights & Biases?

WeightsLab vs Weights & Biases

Weights & Biases (wandb) tracks experiments. WeightsLab connects training signals back to the exact samples causing them — so you can fix your data, not just log it.

--- train_baseline.py+++ train_wl.py@@ -1,11 +1,12 @@
import argparse
import torch
import torch.nn as nn
-from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
from torchmetrics.classification import MulticlassAccuracy
-import wandb+import weightslab as wl+from weightslab.components.global_monitoring import (+ guard_training_context, guard_testing_context)++@wl.signal(name="byte_adjusted_loss", subscribe_to="loss/CE")+def byte_adjusted_loss(ctx): return ctx.subscribed_value / ctx.image_bytes+
def main():
@@ -15,29 +16,38 @@
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
parameters = {"batch_size": 128, "lr": 1e-3}
- wandb.init(project="cifar10")-
transform = transforms.Compose([...])
train_set = datasets.CIFAR10("./data", train=True, download=True, transform=transform)
test_set = datasets.CIFAR10("./data", train=False, download=True, transform=transform)
- train_loader = DataLoader(train_set, batch_size=parameters["batch_size"], shuffle=True, num_workers=2)- test_loader = DataLoader(test_set, batch_size=256, num_workers=2)+ wl.watch_or_edit(parameters, flag="hyperparameters") # live-editable in UI++ train_loader = wl.watch_or_edit(+ train_set, flag="data", loader_name="train_loader",+ batch_size=parameters["batch_size"], shuffle=True, is_training=True)+ test_loader = wl.watch_or_edit(+ test_set, flag="data", loader_name="test_loader",+ batch_size=256, shuffle=False, is_training=False)
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, 10)
optimizer = torch.optim.Adam(model.parameters(), lr=parameters["lr"])
- criterion = nn.CrossEntropyLoss()- accuracy = MulticlassAccuracy(num_classes=10).to(device)+ criterion = wl.watch_or_edit(nn.CrossEntropyLoss(), flag="loss", signal_name="loss/CE")+ accuracy = wl.watch_or_edit(MulticlassAccuracy(num_classes=10).to(device), flag="metric", signal_name="acc")++ wl.serve(serving_grpc=True)
for epoch in range(1, args.epochs + 1):
model.train()
for x, y in train_loader:
+ with guard_training_context:
logits = model(x.to(device))
loss = criterion(logits, y.to(device))
optimizer.zero_grad(); loss.backward(); optimizer.step()
accuracy.update(logits, y)
- wandb.log({"train/loss": loss.item()})- wandb.log({"train/acc": accuracy.compute().item(), "epoch": epoch})+ wl.save_signals(preds_raw=logits, targets=y,+ signals={"metric/accuracy": accuracy.compute().item()})
model.eval()
with torch.no_grad():
for x, y in test_loader:
+ with guard_testing_context:
accuracy.update(model(x.to(device)), y)
- wandb.log({"test/acc": accuracy.compute().item(), "epoch": epoch})+ wl.save_signals(preds_raw=logits, targets=y,+ signals={"metric/accuracy": accuracy.compute().item()})- wandb.finish()+ wl.keep_serving()

Documentation (API + SDK)

Find our documentation online.

Contributing & Onboarding

New here (human or AI coding agent)? Start with AGENTS.md — it captures the cross-repo architecture (weightslab backend ↔ weights_studio frontend via the shared proto), the module maps, the integration pattern, where tests live, and the gotchas that aren't obvious from any single file. It's the fastest way to orient before a first change and contribution.

Community

We're building a community of ML engineers around data-centric training tooling. Interested in contributing or just want to say hi? → hello [at] graybx [dot] com