' +
+ '' +
+ ''
+ : '';
+ // The wrapper carries the filter data so the whole card (badge included)
+ // shows/hides together.
return (
- '' +
- '' + esc(ex.badge) + '' +
- '' + esc(ex.title) + '
' + - '' + esc(ex.desc) + '
' + - '' + - '' + '' + esc(ex.title) + '
' + + '' + esc(ex.desc) + '
' + + '' + + '' + + colabHtml + + '
+
+
+
+
+
+
+
+
+
+
+
+ ![\"WeightsLab](\"https://raw.githubusercontent.com/GrayboxTech/.github/main/profile/weightslab-banner-dark.png\")
\n",
+ "\n",
+ " ![\"License\"](\"https://img.shields.io/badge/License-Apache%202.0-blue.svg\"){kind=icon}
\n",
+ " ![\"Stars\"](\"https://img.shields.io/github/stars/GrayboxTech/weightslab?style=flat&color=5865F2\")
\n",
+ " ![\"Version\"](\"https://img.shields.io/pypi/v/weightslab?style=flat&color=5865F2&logo=pypi&logoColor=white\")
\n",
+ " ![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
\n",
+ "\n",
+ " Welcome to the WeightsLab image-classification notebook! WeightsLab is an open-source PyTorch tool for dataset debugging, mislabel detection, and mid-training data curation. Browse the Docs for details, and raise an issue on GitHub for support.![\"PyPI](\"https://img.shields.io/pypi/v/weightslab?color=5865F2&logo=pypi&logoColor=white\")
\n",
+ "![\"PyPI](\"https://img.shields.io/pypi/dm/weightslab?color=5865F2\")
\n",
+ "![\"PyPI](\"https://img.shields.io/pypi/pyversions/weightslab?color=5865F2&logo=python&logoColor=white\")
"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%pip install -q weightslab\n",
+ "# Override these for Colab compatibility.\n",
+ "!pip install --upgrade \"protobuf>=6.31.1\"\n",
+ "%pip install \"torchvision>=0.16\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 1. Imports\n",
+ "\n",
+ "`weightslab` is imported as `wl`. The two `guard_*_context` managers scope a block as training vs. evaluation so signals are attributed to the right phase."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "import tempfile\n",
+ "import logging\n",
+ "\n",
+ "import torch\n",
+ "import torch.nn as nn\n",
+ "import torch.optim as optim\n",
+ "from torch.utils.data import Dataset\n",
+ "from torchvision import datasets, transforms\n",
+ "from torchmetrics.classification import Accuracy\n",
+ "from tqdm.auto import tqdm\n",
+ "\n",
+ "import weightslab as wl\n",
+ "from weightslab.examples.utils.baseline_models.pytorch.models import FashionCNN as CNN\n",
+ "from weightslab.components.global_monitoring import (\n",
+ " guard_training_context,\n",
+ " guard_testing_context,\n",
+ ")\n",
+ "\n",
+ "logging.basicConfig(level=logging.ERROR)\n",
+ "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+ "print(f\"Using device: {device}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 2. A dataset that carries sample identity\n",
+ "\n",
+ "WeightsLab attributes every signal to a **sample id**. This thin wrapper around torchvision's MNIST returns `(image, id, label)` and attaches a virtual filepath per sample, so signals can be traced back to individual images in the UI."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "class MNISTCustomDataset(Dataset):\n",
+ " \"\"\"MNIST that returns (image, index, label) and tracks a virtual filepath.\"\"\"\n",
+ "\n",
+ " def __init__(self, root, train=True, download=False, transform=None, max_samples=None):\n",
+ " self.mnist = datasets.MNIST(root=root, train=train, download=download, transform=None)\n",
+ " self.transform = transform\n",
+ " self.train = train\n",
+ " self.max_samples = max_samples\n",
+ " split = \"train\" if train else \"test\"\n",
+ " self.filepaths = {}\n",
+ " for idx in range(len(self.mnist)):\n",
+ " if max_samples is not None and idx >= max_samples:\n",
+ " break\n",
+ " label = self.mnist.targets[idx].item()\n",
+ " self.filepaths[idx] = os.path.join(\n",
+ " \"MNIST\", \"processed\", split, f\"class_{label}\", f\"sample_{idx:05d}.pt\"\n",
+ " )\n",
+ "\n",
+ " def __len__(self):\n",
+ " if self.max_samples is not None:\n",
+ " return min(len(self.mnist), self.max_samples)\n",
+ " return len(self.mnist)\n",
+ "\n",
+ " def __getitem__(self, idx):\n",
+ " image, label = self.mnist[idx]\n",
+ " if self.transform:\n",
+ " image = self.transform(image)\n",
+ " return image, idx, label"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 3. Configuration\n",
+ "\n",
+ "Every tunable lives here, in one dict, like a `config.yaml` with comments. Wrapping it with `flag=\"hyperparameters\"` lets the Studio UI read (and live-edit) these values while training."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "log_dir = tempfile.mkdtemp(prefix=\"weightslab_mnist_\")\n",
+ "\n",
+ "config = {\n",
+ " # -- Experiment -------------------------------------------------------\n",
+ " \"experiment_name\": \"mnist_classification\", # name shown in Weights Studio\n",
+ " \"device\": str(device), # \"cuda\" if a GPU is available, else \"cpu\"\n",
+ " \"root_log_dir\": log_dir, # where signal history / dataframes are written\n",
+ " \"num_classes\": 10, # MNIST digits 0-9\n",
+ "\n",
+ " # -- Training schedule ------------------------------------------------\n",
+ " \"training_steps_to_do\": 2000, # total optimizer steps (raise for a longer live run)\n",
+ " \"eval_full_to_train_steps_ratio\": 100, # run a full eval every N steps\n",
+ " \"write_export_ratio\": 100, # export signal history + dataframe every N steps\n",
+ "\n",
+ " # -- Optimizer --------------------------------------------------------\n",
+ " \"learning_rate\": 0.01, # Adam learning rate\n",
+ "\n",
+ " # -- Data loaders -----------------------------------------------------\n",
+ " \"train_batch_size\": 16, # training batch size\n",
+ " \"test_batch_size\": 128, # evaluation batch size\n",
+ " \"max_train_samples\": None, # cap train set (None = full 60k)\n",
+ " \"max_test_samples\": None, # cap test set (None = full 10k)\n",
+ "\n",
+ " # -- Services ---------------------------------------------------------\n",
+ " \"serving_grpc\": True, # expose the gRPC backend for Weights Studio\n",
+ "\n",
+ " # -- Live UI tunnel (see the Expose section) -------------------------\n",
+ " \"expose_ui\": True, # open a bore tunnel so a local Studio can connect\n",
+ " \"backend_port\": 50051, # gRPC port the tunnel forwards\n",
+ "}\n",
+ "\n",
+ "wl.watch_or_edit(config, flag=\"hyperparameters\", poll_interval=1.0)\n",
+ "print(\"Experiment logs ->\", log_dir)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 4. Wrap the training objects\n",
+ "\n",
+ "This is the heart of WeightsLab. Each object is passed through `wl.watch_or_edit(...)` with a `flag` describing its role. The returned objects behave exactly like the originals, but now report their state and per-sample signals to WeightsLab."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_root = os.path.join(log_dir, \"data\")\n",
+ "os.makedirs(data_root, exist_ok=True)\n",
+ "\n",
+ "train_ds = MNISTCustomDataset(root=data_root, train=True, download=True,\n",
+ " transform=transforms.ToTensor(),\n",
+ " max_samples=config[\"max_train_samples\"])\n",
+ "test_ds = MNISTCustomDataset(root=data_root, train=False, download=True,\n",
+ " transform=transforms.ToTensor(),\n",
+ " max_samples=config[\"max_test_samples\"])\n",
+ "\n",
+ "# Model + optimizer\n",
+ "model = wl.watch_or_edit(CNN().to(device), flag=\"model\", device=device)\n",
+ "optimizer = wl.watch_or_edit(\n",
+ " optim.Adam(model.parameters(), lr=config[\"learning_rate\"]), flag=\"optimizer\")\n",
+ "\n",
+ "# Tracked dataloaders\n",
+ "train_loader = wl.watch_or_edit(\n",
+ " train_ds, flag=\"data\", loader_name=\"train_loader\",\n",
+ " batch_size=config[\"train_batch_size\"], shuffle=True, is_training=True,\n",
+ " compute_hash=False, preload_labels=True,\n",
+ " preload_metadata=False, enable_h5_persistence=True,\n",
+ ")\n",
+ "test_loader = wl.watch_or_edit(\n",
+ " test_ds, flag=\"data\", loader_name=\"test_loader\",\n",
+ " batch_size=config[\"test_batch_size\"], shuffle=False, is_training=False,\n",
+ " compute_hash=False, preload_labels=True,\n",
+ " preload_metadata=False, enable_h5_persistence=True,\n",
+ ")\n",
+ "\n",
+ "# Watched losses + metric (they log themselves per sample)\n",
+ "train_criterion = wl.watch_or_edit(nn.CrossEntropyLoss(reduction=\"none\"),\n",
+ " flag=\"loss\", signal_name=\"train-loss-CE\", log=True)\n",
+ "test_criterion = wl.watch_or_edit(nn.CrossEntropyLoss(reduction=\"none\"),\n",
+ " flag=\"loss\", signal_name=\"test-loss-CE\", log=True)\n",
+ "metric = wl.watch_or_edit(\n",
+ " Accuracy(task=\"multiclass\", num_classes=config[\"num_classes\"]).to(device),\n",
+ " flag=\"metric\", signal_name=\"metric-ACC\", log=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 5. Train and evaluate steps\n",
+ "\n",
+ "The `guard_training_context` / `guard_testing_context` blocks tell WeightsLab which phase it's in. `criterion(..., batch_ids=ids, preds=preds)` passes the sample ids so the loss is stored **per sample**, and `wl.save_signals(...)` logs a custom per-sample accuracy signal during evaluation."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def train(loader, model, optimizer, criterion, device):\n",
+ " with guard_training_context:\n",
+ " inputs, ids, labels = next(loader)\n",
+ " inputs, labels = inputs.to(device), labels.to(device)\n",
+ "\n",
+ " optimizer.zero_grad()\n",
+ " logits = model(inputs)\n",
+ " preds = logits.argmax(dim=1, keepdim=True)\n",
+ "\n",
+ " loss = criterion(logits.float(), labels.long(), batch_ids=ids, preds=preds)\n",
+ " total_loss = loss.mean()\n",
+ " total_loss.backward()\n",
+ " optimizer.step()\n",
+ " return total_loss.detach().cpu().item()\n",
+ "\n",
+ "\n",
+ "def test(loader, model, criterion, metric, device, n_batches):\n",
+ " losses = torch.tensor(0.0, device=device)\n",
+ " for inputs, ids, labels in loader:\n",
+ " with guard_testing_context:\n",
+ " inputs, labels = inputs.to(device), labels.to(device)\n",
+ " logits = model(inputs)\n",
+ " preds = logits.argmax(dim=1, keepdim=True)\n",
+ "\n",
+ " loss = criterion(logits, labels, batch_ids=ids, preds=preds)\n",
+ " losses += loss.mean()\n",
+ " metric.update(logits, labels)\n",
+ "\n",
+ " correct = (preds.view(-1) == labels.view(-1)).float()\n",
+ " wl.save_signals(\n",
+ " preds_raw=logits, targets=labels, batch_ids=ids, preds=preds,\n",
+ " signals={\n",
+ " \"test_metric/Accuracy_per_sample\": correct,\n",
+ " \"test_metric/Inverse_Accuracy_per_sample\": 1.0 - correct,\n",
+ " },\n",
+ " )\n",
+ " return (losses / n_batches).item(), (metric.compute() * 100).item()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 6. (Optional) Expose the backend for the live UI\n",
+ "\n",
+ "Skip this if you only want the inline results below. To watch training **live in Weights Studio on your own machine**, keep `config[\"expose_ui\"] = True`.\n",
+ "\n",
+ "This opens a **raw-TCP** tunnel over [`bore`](https://github.com/ekzhang/bore) and its free public relay (`bore.pub`) - **no signup, no account, no credit card** - and prints the exact `weightslab tunnel ...` command to run locally.\n",
+ "\n",
+ "> Note: raw TCP is required so gRPC's HTTP/2 frames pass through untouched (this is also why we avoid `ngrok`, which now needs a credit card for TCP endpoints).\n",
+ ">\n",
+ "> Note: `bore.pub` is a shared public relay; the random port is the only thing protecting your endpoint. Fine for a demo, not for sensitive data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import re, tarfile, threading, urllib.request, subprocess\n",
+ "\n",
+ "endpoint = None\n",
+ "if config[\"expose_ui\"]:\n",
+ " BORE = \"v0.6.0\"\n",
+ " if not os.path.exists(\"bore\"):\n",
+ " urllib.request.urlretrieve(\n",
+ " f\"https://github.com/ekzhang/bore/releases/download/{BORE}/bore-{BORE}-x86_64-unknown-linux-musl.tar.gz\",\n",
+ " \"bore.tar.gz\")\n",
+ " with tarfile.open(\"bore.tar.gz\") as t:\n",
+ " t.extractall()\n",
+ " os.chmod(\"bore\", 0o755)\n",
+ "\n",
+ " proc = subprocess.Popen(\n",
+ " [\"./bore\", \"local\", str(config[\"backend_port\"]), \"--to\", \"bore.pub\"],\n",
+ " stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)\n",
+ " for line in proc.stdout:\n",
+ " m = re.search(r\"bore\\.pub:(\\d+)\", line)\n",
+ " if m:\n",
+ " endpoint = f\"bore.pub:{m.group(1)}\"\n",
+ " break\n",
+ " threading.Thread(target=lambda: [None for _ in proc.stdout], daemon=True).start()\n",
+ "\n",
+ " print(\"=\" * 60)\n",
+ " print(\"Backend exposed at:\", endpoint)\n",
+ " print(\"On your OWN machine (Docker running), run these two commands:\")\n",
+ " print(\" weightslab ui launch\")\n",
+ " print(f\" weightslab tunnel {endpoint}\")\n",
+ " print(\"=\" * 60)\n",
+ "else:\n",
+ " print(\"expose_ui is False - inline results only, no tunnel.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 7. Serve and train\n",
+ "\n",
+ "`wl.serve(serving_grpc=True)` starts the background gRPC server (non-blocking) that Weights Studio connects to. `wl.start_training(...)` flips the experiment into the *training* state, then we run the loop, periodically evaluating and exporting signals.\n",
+ "\n",
+ "Leave this cell **running** while you watch it stream in the UI."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "wl.serve(serving_grpc=config[\"serving_grpc\"])\n",
+ "wl.start_training(timeout=3)\n",
+ "\n",
+ "steps = config[\"training_steps_to_do\"]\n",
+ "eval_ratio = config[\"eval_full_to_train_steps_ratio\"]\n",
+ "export_ratio = config[\"write_export_ratio\"]\n",
+ "n_test_batches = len(test_loader)\n",
+ "\n",
+ "test_loss = test_acc = None\n",
+ "pbar = tqdm(range(steps), desc=\"Training\")\n",
+ "for step in pbar:\n",
+ " age = model.get_age() if hasattr(model, \"get_age\") else step\n",
+ "\n",
+ " train_loss = train(train_loader, model, optimizer, train_criterion, device)\n",
+ "\n",
+ " if age > 0 and age % eval_ratio == 0:\n",
+ " test_loss, test_acc = test(test_loader, model, test_criterion, metric, device, n_test_batches)\n",
+ "\n",
+ " if age > 0 and age % export_ratio == 0:\n",
+ " wl.write_history()\n",
+ " wl.write_dataframe()\n",
+ "\n",
+ " postfix = {\"loss\": f\"{train_loss:.3f}\"}\n",
+ " if test_acc is not None:\n",
+ " postfix[\"test_acc\"] = f\"{test_acc:.1f}%\"\n",
+ " pbar.set_postfix(postfix)\n",
+ "\n",
+ "wl.write_history()\n",
+ "wl.write_dataframe()\n",
+ "print(\"Training complete. Logs at:\", log_dir)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 8. Which samples hurt the most?\n",
+ "\n",
+ "WeightsLab exports a per-sample dataframe to `root_log_dir`. Ranking by loss surfaces the samples the model struggles with - the usual suspects for **mislabels and outliers**. This is the same signal the Studio UI visualizes; here we peek at it inline."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import glob\n",
+ "import pandas as pd\n",
+ "\n",
+ "paths = sorted(glob.glob(os.path.join(log_dir, \"**\", \"*.json\"), recursive=True),\n",
+ " key=os.path.getmtime)\n",
+ "if paths:\n",
+ " df = pd.read_json(paths[-1])\n",
+ " print(f\"Loaded {paths[-1]} (shape={df.shape})\")\n",
+ " loss_cols = [c for c in df.columns if \"loss\" in c.lower()]\n",
+ " if loss_cols:\n",
+ " display(df.sort_values(loss_cols[-1], ascending=False).head(50))\n",
+ " else:\n",
+ " display(df.head())\n",
+ "else:\n",
+ " print(\"No export found - run the training cell first.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 9. See it live in Weights Studio\n",
+ "\n",
+ "Everything above ran headless. The real payoff is the **Weights Studio** UI, where you browse the highest-loss images, filter by class, and curate the dataset mid-training.\n",
+ "\n",
+ "Studio runs as a local Docker stack (a static frontend + an Envoy gRPC-Web proxy). **Colab has no Docker daemon**, so you don't run the UI *inside* Colab - you run Studio on your own machine and point it at this notebook's backend using the `bore.pub:\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " Welcome to the WeightsLab Face Recognition (triplet loss) notebook! WeightsLab traces training signals back to the exact samples producing them. Browse the Docs for details.\n",
+ "\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Clone the repo (for the example's utils/ + data) and install WeightsLab.\n",
+ "import os\n",
+ "if not os.path.isdir(\"weightslab\"):\n",
+ " !git clone --branch torch_collab_examples https://github.com/GrayboxTech/weightslab.git\n",
+ "%pip install -q ./weightslab\n",
+ "!pip install --upgrade \"protobuf>=6.31.1\"\n",
+ "%pip install \"torchvision>=0.16\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Move into the example directory so its `utils/` and config.yaml resolve.\n",
+ "%cd weightslab/weightslab/examples/PyTorch/ws-clustering\n",
+ "# Install the example's own extra requirements, if any.\n",
+ "import os\n",
+ "if os.path.exists(\"requirements.txt\"):\n",
+ " !pip install -q -r requirements.txt"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Configuration\n",
+ "\n",
+ "Every tunable lives in **`config.yaml`** (already commented). We load it into a `config` dict\n",
+ "so you can tweak values here; the changes are written back before training."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import yaml\n",
+ "\n",
+ "with open(\"config.yaml\") as f:\n",
+ " config = yaml.safe_load(f)\n",
+ "\n",
+ "# --- tweak any parameter here, e.g. ---\n",
+ "# config[\"eval_full_to_train_steps_ratio\"] = 5\n",
+ "config[\"serving_grpc\"] = True # expose the gRPC backend for Weights Studio\n",
+ "\n",
+ "with open(\"config.yaml\", \"w\") as f:\n",
+ " yaml.safe_dump(config, f, sort_keys=False)\n",
+ "\n",
+ "print(yaml.safe_dump(config, sort_keys=False))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## (Optional) Expose the backend for the live UI\n",
+ "\n",
+ "To watch training **live in Weights Studio on your own machine**, run this cell first (it opens a\n",
+ "raw-TCP [`bore`](https://github.com/ekzhang/bore) tunnel over the free public relay `bore.pub` -\n",
+ "**no signup, no card**), then start training below.\n",
+ "\n",
+ "> `bore.pub` is a shared public relay; the random port is the only thing protecting your endpoint.\n",
+ "> Fine for a demo, not for sensitive data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os, re, tarfile, threading, urllib.request, subprocess\n",
+ "\n",
+ "EXPOSE_UI = True # set False for a headless run (no tunnel)\n",
+ "BACKEND_PORT = 50051 # gRPC port to forward\n",
+ "\n",
+ "endpoint = None\n",
+ "if EXPOSE_UI:\n",
+ " bore = os.path.join(os.getcwd(), \"bore\")\n",
+ " if not os.path.exists(bore):\n",
+ " BORE = \"v0.6.0\"\n",
+ " urllib.request.urlretrieve(\n",
+ " f\"https://github.com/ekzhang/bore/releases/download/{BORE}/bore-{BORE}-x86_64-unknown-linux-musl.tar.gz\",\n",
+ " \"bore.tar.gz\")\n",
+ " with tarfile.open(\"bore.tar.gz\") as t:\n",
+ " t.extractall()\n",
+ " os.chmod(bore, 0o755)\n",
+ "\n",
+ " proc = subprocess.Popen(\n",
+ " [bore, \"local\", str(BACKEND_PORT), \"--to\", \"bore.pub\"],\n",
+ " stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)\n",
+ " for line in proc.stdout:\n",
+ " m = re.search(r\"bore\\.pub:(\\d+)\", line)\n",
+ " if m:\n",
+ " endpoint = f\"bore.pub:{m.group(1)}\"\n",
+ " break\n",
+ " threading.Thread(target=lambda: [None for _ in proc.stdout], daemon=True).start()\n",
+ "\n",
+ " print(\"=\" * 60)\n",
+ " print(\"Backend exposed at:\", endpoint)\n",
+ " print(\"On your OWN machine (Docker running), run:\")\n",
+ " print(\" weightslab ui launch\")\n",
+ " print(f\" weightslab tunnel {endpoint}\")\n",
+ " print(\"=\" * 60)\n",
+ "else:\n",
+ " print(\"EXPOSE_UI is False - no tunnel.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Train\n",
+ "\n",
+ "This runs the example's `main.py`, which wraps the model / optimizer / loaders / losses with\n",
+ "`wl.watch_or_edit(...)`, starts the gRPC backend, and trains while streaming per-sample signals.\n",
+ "\n",
+ "> This example trains **open-ended** (until you interrupt the cell). Interrupt the cell (stop button) to stop training; the exported history and\n",
+ "> dataframe are written to a temp `root_log_dir` (printed at startup)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "!python main.py"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## See it live in Weights Studio\n",
+ "\n",
+ "**Colab has no Docker daemon**, so run Studio on your own machine and point it at this notebook's\n",
+ "backend using the `bore.pub:\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " Welcome to the WeightsLab Object Detection (Penn-Fudan) notebook! WeightsLab traces training signals back to the exact samples producing them. Browse the Docs for details.\n",
+ "\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Clone the repo (for the example's utils/ + data) and install WeightsLab.\n",
+ "import os\n",
+ "if not os.path.isdir(\"weightslab\"):\n",
+ " !git clone --branch torch_collab_examples https://github.com/GrayboxTech/weightslab.git\n",
+ "%pip install -q ./weightslab\n",
+ "!pip install --upgrade \"protobuf>=6.31.1\"\n",
+ "%pip install \"torchvision>=0.16\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Move into the example directory so its `utils/` and config.yaml resolve.\n",
+ "%cd weightslab/weightslab/examples/PyTorch/ws-detection\n",
+ "# Install the example's own extra requirements, if any.\n",
+ "import os\n",
+ "if os.path.exists(\"requirements.txt\"):\n",
+ " !pip install -q -r requirements.txt"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Configuration\n",
+ "\n",
+ "Every tunable lives in **`config.yaml`** (already commented). We load it into a `config` dict\n",
+ "so you can tweak values here; the changes are written back before training."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import yaml\n",
+ "\n",
+ "with open(\"config.yaml\") as f:\n",
+ " config = yaml.safe_load(f)\n",
+ "\n",
+ "# --- tweak any parameter here, e.g. ---\n",
+ "# config[\"eval_full_to_train_steps_ratio\"] = 5\n",
+ "config[\"serving_grpc\"] = True # expose the gRPC backend for Weights Studio\n",
+ "\n",
+ "with open(\"config.yaml\", \"w\") as f:\n",
+ " yaml.safe_dump(config, f, sort_keys=False)\n",
+ "\n",
+ "print(yaml.safe_dump(config, sort_keys=False))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## (Optional) Expose the backend for the live UI\n",
+ "\n",
+ "To watch training **live in Weights Studio on your own machine**, run this cell first (it opens a\n",
+ "raw-TCP [`bore`](https://github.com/ekzhang/bore) tunnel over the free public relay `bore.pub` -\n",
+ "**no signup, no card**), then start training below.\n",
+ "\n",
+ "> `bore.pub` is a shared public relay; the random port is the only thing protecting your endpoint.\n",
+ "> Fine for a demo, not for sensitive data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os, re, tarfile, threading, urllib.request, subprocess\n",
+ "\n",
+ "EXPOSE_UI = True # set False for a headless run (no tunnel)\n",
+ "BACKEND_PORT = 50051 # gRPC port to forward\n",
+ "\n",
+ "endpoint = None\n",
+ "if EXPOSE_UI:\n",
+ " bore = os.path.join(os.getcwd(), \"bore\")\n",
+ " if not os.path.exists(bore):\n",
+ " BORE = \"v0.6.0\"\n",
+ " urllib.request.urlretrieve(\n",
+ " f\"https://github.com/ekzhang/bore/releases/download/{BORE}/bore-{BORE}-x86_64-unknown-linux-musl.tar.gz\",\n",
+ " \"bore.tar.gz\")\n",
+ " with tarfile.open(\"bore.tar.gz\") as t:\n",
+ " t.extractall()\n",
+ " os.chmod(bore, 0o755)\n",
+ "\n",
+ " proc = subprocess.Popen(\n",
+ " [bore, \"local\", str(BACKEND_PORT), \"--to\", \"bore.pub\"],\n",
+ " stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)\n",
+ " for line in proc.stdout:\n",
+ " m = re.search(r\"bore\\.pub:(\\d+)\", line)\n",
+ " if m:\n",
+ " endpoint = f\"bore.pub:{m.group(1)}\"\n",
+ " break\n",
+ " threading.Thread(target=lambda: [None for _ in proc.stdout], daemon=True).start()\n",
+ "\n",
+ " print(\"=\" * 60)\n",
+ " print(\"Backend exposed at:\", endpoint)\n",
+ " print(\"On your OWN machine (Docker running), run:\")\n",
+ " print(\" weightslab ui launch\")\n",
+ " print(f\" weightslab tunnel {endpoint}\")\n",
+ " print(\"=\" * 60)\n",
+ "else:\n",
+ " print(\"EXPOSE_UI is False - no tunnel.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Train\n",
+ "\n",
+ "This runs the example's `main.py`, which wraps the model / optimizer / loaders / losses with\n",
+ "`wl.watch_or_edit(...)`, starts the gRPC backend, and trains while streaming per-sample signals.\n",
+ "\n",
+ "> This example trains **open-ended** (until you interrupt the cell). Interrupt the cell (stop button) to stop training; the exported history and\n",
+ "> dataframe are written to a temp `root_log_dir` (printed at startup)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "!python main.py"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## See it live in Weights Studio\n",
+ "\n",
+ "**Colab has no Docker daemon**, so run Studio on your own machine and point it at this notebook's\n",
+ "backend using the `bore.pub:\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " Welcome to the WeightsLab Image Generation (VAD U-Net) notebook! WeightsLab traces training signals back to the exact samples producing them. Browse the Docs for details.\n",
+ "\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Clone the repo (for the example's utils/ + data) and install WeightsLab.\n",
+ "import os\n",
+ "if not os.path.isdir(\"weightslab\"):\n",
+ " !git clone --branch torch_collab_examples https://github.com/GrayboxTech/weightslab.git\n",
+ "%pip install -q ./weightslab\n",
+ "!pip install --upgrade \"protobuf>=6.31.1\"\n",
+ "%pip install \"torchvision>=0.16\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Move into the example directory so its `utils/` and config.yaml resolve.\n",
+ "%cd weightslab/weightslab/examples/PyTorch/ws-generation\n",
+ "# Install the example's own extra requirements, if any.\n",
+ "import os\n",
+ "if os.path.exists(\"requirements.txt\"):\n",
+ " !pip install -q -r requirements.txt"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Configuration\n",
+ "\n",
+ "Every tunable lives in **`config.yaml`** (already commented). We load it into a `config` dict\n",
+ "so you can tweak values here; the changes are written back before training."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import yaml\n",
+ "\n",
+ "with open(\"config.yaml\") as f:\n",
+ " config = yaml.safe_load(f)\n",
+ "\n",
+ "# --- tweak any parameter here, e.g. ---\n",
+ "# config[\"eval_full_to_train_steps_ratio\"] = 5\n",
+ "config[\"serving_grpc\"] = True # expose the gRPC backend for Weights Studio\n",
+ "\n",
+ "with open(\"config.yaml\", \"w\") as f:\n",
+ " yaml.safe_dump(config, f, sort_keys=False)\n",
+ "\n",
+ "print(yaml.safe_dump(config, sort_keys=False))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## (Optional) Expose the backend for the live UI\n",
+ "\n",
+ "To watch training **live in Weights Studio on your own machine**, run this cell first (it opens a\n",
+ "raw-TCP [`bore`](https://github.com/ekzhang/bore) tunnel over the free public relay `bore.pub` -\n",
+ "**no signup, no card**), then start training below.\n",
+ "\n",
+ "> `bore.pub` is a shared public relay; the random port is the only thing protecting your endpoint.\n",
+ "> Fine for a demo, not for sensitive data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os, re, tarfile, threading, urllib.request, subprocess\n",
+ "\n",
+ "EXPOSE_UI = True # set False for a headless run (no tunnel)\n",
+ "BACKEND_PORT = 50051 # gRPC port to forward\n",
+ "\n",
+ "endpoint = None\n",
+ "if EXPOSE_UI:\n",
+ " bore = os.path.join(os.getcwd(), \"bore\")\n",
+ " if not os.path.exists(bore):\n",
+ " BORE = \"v0.6.0\"\n",
+ " urllib.request.urlretrieve(\n",
+ " f\"https://github.com/ekzhang/bore/releases/download/{BORE}/bore-{BORE}-x86_64-unknown-linux-musl.tar.gz\",\n",
+ " \"bore.tar.gz\")\n",
+ " with tarfile.open(\"bore.tar.gz\") as t:\n",
+ " t.extractall()\n",
+ " os.chmod(bore, 0o755)\n",
+ "\n",
+ " proc = subprocess.Popen(\n",
+ " [bore, \"local\", str(BACKEND_PORT), \"--to\", \"bore.pub\"],\n",
+ " stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)\n",
+ " for line in proc.stdout:\n",
+ " m = re.search(r\"bore\\.pub:(\\d+)\", line)\n",
+ " if m:\n",
+ " endpoint = f\"bore.pub:{m.group(1)}\"\n",
+ " break\n",
+ " threading.Thread(target=lambda: [None for _ in proc.stdout], daemon=True).start()\n",
+ "\n",
+ " print(\"=\" * 60)\n",
+ " print(\"Backend exposed at:\", endpoint)\n",
+ " print(\"On your OWN machine (Docker running), run:\")\n",
+ " print(\" weightslab ui launch\")\n",
+ " print(f\" weightslab tunnel {endpoint}\")\n",
+ " print(\"=\" * 60)\n",
+ "else:\n",
+ " print(\"EXPOSE_UI is False - no tunnel.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Train\n",
+ "\n",
+ "This runs the example's `main.py`, which wraps the model / optimizer / loaders / losses with\n",
+ "`wl.watch_or_edit(...)`, starts the gRPC backend, and trains while streaming per-sample signals.\n",
+ "\n",
+ "> Training runs for `training_steps` (default 200000); interrupt the cell to stop early. Interrupt the cell (stop button) to stop training; the exported history and\n",
+ "> dataframe are written to a temp `root_log_dir` (printed at startup)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "!python main.py"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## See it live in Weights Studio\n",
+ "\n",
+ "**Colab has no Docker daemon**, so run Studio on your own machine and point it at this notebook's\n",
+ "backend using the `bore.pub:\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " Welcome to the WeightsLab Semantic Segmentation (BDD100k) notebook! WeightsLab traces training signals back to the exact samples producing them. Browse the Docs for details.\n",
+ "\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Clone the repo (for the example's utils/ + data) and install WeightsLab.\n",
+ "import os\n",
+ "if not os.path.isdir(\"weightslab\"):\n",
+ " !git clone --branch torch_collab_examples https://github.com/GrayboxTech/weightslab.git\n",
+ "%pip install -q ./weightslab\n",
+ "!pip install --upgrade \"protobuf>=6.31.1\"\n",
+ "%pip install \"torchvision>=0.16\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Move into the example directory so its `utils/` and config.yaml resolve.\n",
+ "%cd weightslab/weightslab/examples/PyTorch/ws-segmentation\n",
+ "# Install the example's own extra requirements, if any.\n",
+ "import os\n",
+ "if os.path.exists(\"requirements.txt\"):\n",
+ " !pip install -q -r requirements.txt"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Configuration\n",
+ "\n",
+ "Every tunable lives in **`config.yaml`** (already commented). We load it into a `config` dict\n",
+ "so you can tweak values here; the changes are written back before training."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import yaml\n",
+ "\n",
+ "with open(\"config.yaml\") as f:\n",
+ " config = yaml.safe_load(f)\n",
+ "\n",
+ "# --- tweak any parameter here, e.g. ---\n",
+ "# config[\"eval_full_to_train_steps_ratio\"] = 5\n",
+ "config[\"serving_grpc\"] = True # expose the gRPC backend for Weights Studio\n",
+ "\n",
+ "with open(\"config.yaml\", \"w\") as f:\n",
+ " yaml.safe_dump(config, f, sort_keys=False)\n",
+ "\n",
+ "print(yaml.safe_dump(config, sort_keys=False))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## (Optional) Expose the backend for the live UI\n",
+ "\n",
+ "To watch training **live in Weights Studio on your own machine**, run this cell first (it opens a\n",
+ "raw-TCP [`bore`](https://github.com/ekzhang/bore) tunnel over the free public relay `bore.pub` -\n",
+ "**no signup, no card**), then start training below.\n",
+ "\n",
+ "> `bore.pub` is a shared public relay; the random port is the only thing protecting your endpoint.\n",
+ "> Fine for a demo, not for sensitive data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os, re, tarfile, threading, urllib.request, subprocess\n",
+ "\n",
+ "EXPOSE_UI = True # set False for a headless run (no tunnel)\n",
+ "BACKEND_PORT = 50051 # gRPC port to forward\n",
+ "\n",
+ "endpoint = None\n",
+ "if EXPOSE_UI:\n",
+ " bore = os.path.join(os.getcwd(), \"bore\")\n",
+ " if not os.path.exists(bore):\n",
+ " BORE = \"v0.6.0\"\n",
+ " urllib.request.urlretrieve(\n",
+ " f\"https://github.com/ekzhang/bore/releases/download/{BORE}/bore-{BORE}-x86_64-unknown-linux-musl.tar.gz\",\n",
+ " \"bore.tar.gz\")\n",
+ " with tarfile.open(\"bore.tar.gz\") as t:\n",
+ " t.extractall()\n",
+ " os.chmod(bore, 0o755)\n",
+ "\n",
+ " proc = subprocess.Popen(\n",
+ " [bore, \"local\", str(BACKEND_PORT), \"--to\", \"bore.pub\"],\n",
+ " stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)\n",
+ " for line in proc.stdout:\n",
+ " m = re.search(r\"bore\\.pub:(\\d+)\", line)\n",
+ " if m:\n",
+ " endpoint = f\"bore.pub:{m.group(1)}\"\n",
+ " break\n",
+ " threading.Thread(target=lambda: [None for _ in proc.stdout], daemon=True).start()\n",
+ "\n",
+ " print(\"=\" * 60)\n",
+ " print(\"Backend exposed at:\", endpoint)\n",
+ " print(\"On your OWN machine (Docker running), run:\")\n",
+ " print(\" weightslab ui launch\")\n",
+ " print(f\" weightslab tunnel {endpoint}\")\n",
+ " print(\"=\" * 60)\n",
+ "else:\n",
+ " print(\"EXPOSE_UI is False - no tunnel.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Train\n",
+ "\n",
+ "This runs the example's `main.py`, which wraps the model / optimizer / loaders / losses with\n",
+ "`wl.watch_or_edit(...)`, starts the gRPC backend, and trains while streaming per-sample signals.\n",
+ "\n",
+ "> This example trains **open-ended** (until you interrupt the cell). Interrupt the cell (stop button) to stop training; the exported history and\n",
+ "> dataframe are written to a temp `root_log_dir` (printed at startup)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "!python main.py"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## See it live in Weights Studio\n",
+ "\n",
+ "**Colab has no Docker daemon**, so run Studio on your own machine and point it at this notebook's\n",
+ "backend using the `bore.pub:\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " Welcome to the WeightsLab custom decorated-signal usecase! This notebook trains segmentation and adds a user-defined, decorated per-sample signal that classifies each sample's loss trajectory. See the Docs for details.\n",
+ "\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "if not os.path.isdir(\"weightslab\"):\n",
+ " !git clone --branch torch_collab_examples https://github.com/GrayboxTech/weightslab.git\n",
+ "%pip install -q ./weightslab\n",
+ "!pip install --upgrade \"protobuf>=6.31.1\"\n",
+ "%pip install \"torchvision>=0.16\"\n",
+ "%cd weightslab/weightslab/examples/PyTorch/ws-segmentation"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 1. Imports\n",
+ "\n",
+ "We reuse the segmentation example's helper modules (`utils/`) for the dataset, model, and the\n",
+ "per-sample / per-instance Dice + BCE criterions."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import torch\n",
+ "from torch import optim\n",
+ "from tqdm.auto import tqdm\n",
+ "\n",
+ "import weightslab as wl\n",
+ "from weightslab.components.global_monitoring import (\n",
+ " guard_training_context, guard_testing_context,\n",
+ ")\n",
+ "from utils.data import BDD100kSegDataset, seg_collate\n",
+ "from utils.model import SmallUNet\n",
+ "from utils.criterions import (\n",
+ " PerSampleDice, PerInstanceDice, PerSampleBCE, PerInstanceBCE,\n",
+ ")\n",
+ "\n",
+ "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+ "print(\"Using device:\", device)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 2. Configuration\n",
+ "\n",
+ "All tunables in one dict (like a `config.yaml`, with comments)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "config = {\n",
+ " # -- Experiment ----------------------------------------------------------\n",
+ " \"experiment_name\": \"seg_loss_shapes\", # name shown in Weights Studio\n",
+ " \"root_log_dir\": None, # None -> a temp dir is created\n",
+ "\n",
+ " # -- Data (bundled BDD_subset) -------------------------------------------\n",
+ " \"data_root\": \"BDD_subset\", # ships in the repo\n",
+ " \"num_classes\": 6, # BDD label set\n",
+ " \"class_names\": [\"Background\", \"Ego Road\", \"Driveable Area\",\n",
+ " \"Lane Line 1\", \"Lane Line 2\", \"Lane Line 3\"],\n",
+ " \"ignore_index\": 255, # void pixels\n",
+ " \"image_size\": 128, # images resized to image_size^2\n",
+ " \"max_train_samples\": 256, # cap for a quick demo\n",
+ " \"max_test_samples\": 100,\n",
+ " \"train_batch_size\": 2,\n",
+ " \"test_batch_size\": 2,\n",
+ "\n",
+ " # -- Optimizer -----------------------------------------------------------\n",
+ " \"learning_rate\": 1e-3,\n",
+ "\n",
+ " # -- Training schedule ---------------------------------------------------\n",
+ " \"training_steps_to_do\": 800, # > min_step so the classifier fires\n",
+ " \"eval_full_to_train_steps_ratio\": 100, # full eval every N steps\n",
+ " \"write_export_ratio\": 100, # export history + dataframe every N steps\n",
+ "\n",
+ " # -- Custom decorated signal --------------------------------------------\n",
+ " \"shape_subscribe_to\": \"train_bce/sample\", # per-sample loss the classifier watches\n",
+ " \"shape_compute_every_n_steps\": 25, # throttle the classifier\n",
+ " \"shape_min_step\": 505, # WARM-UP: don't classify before this step\n",
+ "\n",
+ " # -- Services / live UI --------------------------------------------------\n",
+ " \"serving_grpc\": True,\n",
+ " \"expose_ui\": True,\n",
+ " \"backend_port\": 50051,\n",
+ "}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 3. Build data, model, optimizer, and the tracked seg signals\n",
+ "\n",
+ "Same wrapping pattern as the segmentation example: dataloaders, model, optimizer, and the\n",
+ "per-sample / per-instance Dice (metric) + BCE (loss) signals are all passed through\n",
+ "`wl.watch_or_edit(...)`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "train_ds = BDD100kSegDataset(root=config[\"data_root\"], split=\"train\",\n",
+ " num_classes=config[\"num_classes\"], class_names=config[\"class_names\"],\n",
+ " ignore_index=config[\"ignore_index\"], image_size=config[\"image_size\"],\n",
+ " max_samples=config[\"max_train_samples\"])\n",
+ "val_ds = BDD100kSegDataset(root=config[\"data_root\"], split=\"val\",\n",
+ " num_classes=config[\"num_classes\"], class_names=config[\"class_names\"],\n",
+ " ignore_index=config[\"ignore_index\"], image_size=config[\"image_size\"],\n",
+ " max_samples=config[\"max_test_samples\"])\n",
+ "\n",
+ "wl.watch_or_edit(config, flag=\"hyperparameters\", name=config[\"experiment_name\"],\n",
+ " defaults=config, poll_interval=1.0)\n",
+ "\n",
+ "train_loader = wl.watch_or_edit(\n",
+ " train_ds, flag=\"data\", loader_name=\"train_loader\",\n",
+ " batch_size=config[\"train_batch_size\"], shuffle=True, is_training=True,\n",
+ " compute_hash=False, array_autoload_arrays=False, array_return_proxies=True,\n",
+ " array_use_cache=True, preload_labels=False, collate_fn=seg_collate)\n",
+ "test_loader = wl.watch_or_edit(\n",
+ " val_ds, flag=\"data\", loader_name=\"test_loader\",\n",
+ " batch_size=config[\"test_batch_size\"], shuffle=False, is_training=False,\n",
+ " compute_hash=False, array_autoload_arrays=False, array_return_proxies=True,\n",
+ " array_use_cache=True, preload_labels=True, collate_fn=seg_collate)\n",
+ "\n",
+ "model = wl.watch_or_edit(\n",
+ " SmallUNet(in_channels=3, num_classes=config[\"num_classes\"],\n",
+ " image_size=config[\"image_size\"]).to(device),\n",
+ " flag=\"model\", device=device, compute_dependencies=True)\n",
+ "optimizer = wl.watch_or_edit(\n",
+ " optim.Adam(model.parameters(), lr=config[\"learning_rate\"]), flag=\"optimizer\")\n",
+ "\n",
+ "\n",
+ "def make_seg_signals(split):\n",
+ " return {\n",
+ " \"dice_sample\": wl.watch_or_edit(PerSampleDice(), flag=\"metric\",\n",
+ " name=f\"{split}_dice/sample\", per_sample=True, log=True),\n",
+ " \"dice_instance\": wl.watch_or_edit(PerInstanceDice(), flag=\"metric\",\n",
+ " name=f\"{split}_dice/instance\", per_instance=True, log=True),\n",
+ " \"bce_sample\": wl.watch_or_edit(PerSampleBCE(), flag=\"loss\",\n",
+ " name=f\"{split}_bce/sample\", per_sample=True, log=True),\n",
+ " \"bce_instance\": wl.watch_or_edit(PerInstanceBCE(), flag=\"loss\",\n",
+ " name=f\"{split}_bce/instance\", per_instance=True, log=True),\n",
+ " }\n",
+ "\n",
+ "train_sig = make_seg_signals(\"train\")\n",
+ "test_sig = make_seg_signals(\"test\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 4. The custom, user-defined decorated signal\n",
+ "\n",
+ "This is the point of the usecase. `classify_loss_shape` is **your** function, decorated with\n",
+ "`@wl.signal`. Because it sets `subscribe_to=