- # ColabFold with W&B: Align, Fold, Log
+ This lesson uses [ColabFold](https://github.com/sokrypton/ColabFold) to
+ align a protein sequence and predict its structure with AlphaFold2. We then
+ log the experiment configuration, confidence metrics, plots, and a 3D
+ molecule to Weights & Biases. A prediction table keeps the sequence and results together
+ for comparison across runs.
- This notebook predicts a protein structure with
- [ColabFold](https://github.com/sokrypton/ColabFold), then logs the sequence,
- alignment summary, confidence plots, timings, and predicted molecule to
- Weights & Biases.
+ This version runs ColabFold from a marimo notebook and adds Weights & Biases
+ logging so you can compare prediction settings and inspect the saved outputs.
- It preserves the **align → fold → log** lesson from
- [DeepMind's original AlphaFold Colab](https://github.com/google-deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb)
- and the W&B adaptation described in
- [this Report](https://wandb.me/alphafold-short-report). ColabFold replaces
- the retired Colab setup with an MMseqs2 alignment service and local
- AlphaFold2 inference.
-
- No Docker image or local genetic database is required. The prediction runs
- inside the notebook; only the submitted sequence is sent to the public
- ColabFold MSA service.
- """)
- return
+ See the [AlphaFold-ed Proteins in Weights & Biases Tables report](https://forge.coreweave.com/wandb/wandb/examples/reports/AlphaFold-ed-Proteins-in-W-B-Tables--Vmlldzo4ODc0MDc) and view an example protein below.
+ ## Prepare the Runtime
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## 0. Runtime
+ Use a Linux-backed molab runtime, or open the notebook with its declared
+ dependencies on Linux:
- Open this notebook in molab. Its Python dependencies include ColabFold and
- the AlphaFold2 inference package.
+ ```bash
+ uvx marimo edit --sandbox alphafold_with_w_b_align_fold_log.py
+ ```
- The first prediction downloads the model parameters. To keep this interactive
- tutorial practical, the notebook uses a bounded MSA depth and automatically
- selects a compatible inference path for the attached runtime. ColabFold queries
- its shared public MSA service for this small demonstration, so submit one
- sequence at a time.
+ The next cell checks the runtime and chooses GPU or CPU inference. It
+ retains the CPU fallback for compute capability 12.0 devices used by this
+ example. You can inspect the selected backend before predicting.
- Opening the notebook does not query the service, download weights, run a
- prediction, or create W&B objects.
+ The prediction cell queries the public ColabFold MSA service and downloads
+ model parameters on first use. Submit one sequence at a time and follow
+ [ColabFold's shared-service guidance](https://github.com/sokrypton/ColabFold).
+ No local sequence database or Docker image is required for this workflow.
""")
return
-@app.cell(hide_code=True)
-def _(mo, subprocess, sys):
- _linux_available = sys.platform.startswith("linux")
+@app.cell
+def _():
+ runtime_ready = sys.platform.startswith("linux")
try:
- _gpu_probe = subprocess.run(
- [
- "nvidia-smi",
- "--query-gpu=name,compute_cap",
- "--format=csv,noheader",
- ],
+ gpu_probe = subprocess.run(
+ ["nvidia-smi", "--query-gpu=name,compute_cap", "--format=csv,noheader"],
check=True,
capture_output=True,
text=True,
)
- _gpu_rows = [
- _row.rsplit(",", maxsplit=1)
- for _row in _gpu_probe.stdout.splitlines()
- if _row.strip()
- ]
- _gpu_names = [_name.strip() for _name, _capability in _gpu_rows]
- _gpu_compute_caps = [
- _capability.strip() for _name, _capability in _gpu_rows
- ]
+ gpu_rows = [row.rsplit(",", maxsplit=1) for row in gpu_probe.stdout.splitlines() if row.strip()]
+ gpu_names = [name.strip() for name, capability in gpu_rows]
+ compute_capabilities = [capability.strip() for name, capability in gpu_rows]
except (FileNotFoundError, subprocess.CalledProcessError, ValueError):
- _gpu_names = []
- _gpu_compute_caps = []
+ gpu_names = []
+ compute_capabilities = []
- _blackwell_sm120 = any(
- _capability == "12.0" for _capability in _gpu_compute_caps
- )
- colabfold_inference_backend = (
- "cpu" if _blackwell_sm120 or not _gpu_names else "gpu"
- )
- colabfold_runtime_ready = _linux_available
-
- if colabfold_runtime_ready:
- _runtime_message = mo.callout(
- mo.md("Runtime check passed. ColabFold is ready."),
- kind="success",
- )
- else:
- _runtime_message = mo.callout(
- mo.md("Open this notebook in a Linux-backed molab runtime."),
- kind="warn",
- title="Molab runtime needed",
- )
- _runtime_message
- return colabfold_inference_backend, colabfold_runtime_ready
+ inference_backend = "cpu" if not gpu_names or "12.0" in compute_capabilities else "gpu"
+ print(f"Platform: {sys.platform}; inference backend: {inference_backend}")
+ if not runtime_ready:
+ print("Structure prediction requires a Linux-backed runtime for this notebook.")
+ return inference_backend, runtime_ready
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## Authentication
+ ## Authenticate with Weights & Biases
+
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- Enter your [W&B API key](https://wandb.ai/authorize), or leave it blank to
- use `WANDB_API_KEY` from the molab Secrets panel or credentials already
- stored in this runtime. A fresh molab session does not inherit credentials
- from your local computer.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- The entity is the team name in a W&B project URL:
- `wandb.ai/
- ## Quickstart
- Use [Weights & Biases](https://wandb.ai)
- for machine learning experiment tracking, dataset versioning, and project collaboration.
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/configs-in-w-b/configs_in_w_b.py/server)
- 
- """)
- return
+ ## Collect command-line settings with `argparse`
+ Training scripts often receive hyperparameters as command-line arguments.
+ Weights & Biases accepts the `argparse.Namespace` returned by `parse_args()` directly
+ in `run.config.update()`.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- You give us a (possibly nested) dictionary as your `config`, and we'll flatten the names using dots in our backend.
-
- > _Side Note_: We recommend that you avoid using dots in your config variable names, and use a dash or underscore instead. Once you've created your `config` dictionary, if your script accesses `wandb.config` keys below the root, use the dictionary access syntax, `["key"]["foo"]`, instead of the attribute access syntax, `config.key.foo`.
+ Here, `parse_args([])` selects the defaults without reading the notebook
+ process's arguments. Try `parse_args(["--batch_per_gpu", "16"])` to see
+ how an override changes the saved configuration. In a standalone training
+ script, use `parse_args()` to read arguments from the command line.
""")
return
-@app.cell(hide_code=True)
-def _():
- add_config_parameters = mo.ui.run_button(
- label="Add parameters to the W&B config"
- )
- mo.vstack(
- [
- mo.md(r"""
- ### Adding to the `config` by hand
- You can add more parameters to the `config` later if you want:
- """),
- add_config_parameters,
- ]
- )
- return (add_config_parameters,)
-
-
@app.cell
-def _(add_config_parameters, config_run):
- mo.stop(
- not add_config_parameters.value,
- mo.callout(
- mo.md(
- "Click the button above to add `epochs` and `batch_size` to "
- "this run's config."
- ),
- kind="info",
- ),
- )
-
- config_run.config.epochs = 4
- config_run.config["batch_size"] = 32
-
- manual_config_update = True
- dict(config_run.config)
- return
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- Now, your Config section on the dashboard has been updated:
- """)
- return
-
-
-@app.cell(hide_code=True)
def _():
- mo.md(r"""
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _():
- add_argparse_parameters = mo.ui.run_button(
- label="Add argparse parameters to the W&B config"
- )
- mo.vstack(
- [
- mo.md(r"""
- ### Adding to the `config` with `argparse`
-
- `config` is a dictionary-like object, and it can be built from lots of dictionary-like objects.
-
- For example, you can pass in the arguments object produced by `argparse`.
- [`argparse`](https://docs.python.org/3/library/argparse.html), short for `arg`ument `parse`r, is a standard library module in Python 3.2 and above that makes it easy to write scripts that take advantage of all the flexibility and power of command line arguments. And it's Pythonic!
-
- This is especially convenient for tracking results from scripts that are launched from the command line.
- """),
- add_argparse_parameters,
- ]
- )
- return (add_argparse_parameters,)
-
-
-@app.cell
-def _(add_argparse_parameters, config_run):
- mo.stop(
- not add_argparse_parameters.value,
- mo.callout(
- mo.md("Click the button above to add the parsed arguments to this run's config."),
- kind="info",
- ),
- )
-
parser = argparse.ArgumentParser()
- parser.add_argument('-b', '--batch_per_gpu', type=int, default=8,
- help='input batch size for training (default: 8)')
- parser.add_argument('-wd', '--weight_decay', type=float, default=0.1,
- help='weight decay (default: 0.1)')
-
- args = parser.parse_args(args=[])
- config_run.config.update(args)
- return
+ parser.add_argument("-b", "--batch_per_gpu", type=int, default=8)
+ parser.add_argument("-wd", "--weight_decay", type=float, default=0.1)
+ args = parser.parse_args([])
+ args
+ return (args,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- Here's the updated Config panel:
- """)
- return
+ ## Save and extend the run's configuration
+ `wandb.init(config=config)` saves the initial dictionary. While the run is
+ active, add new fields using attribute access or dictionary
+ access, then merge the parsed arguments with `update()`.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- 
+ The cell displays the combined configuration so you can check all six
+ fields. The context manager finishes and syncs the run before the next
+ example queries it. It also saves the actual entity, project, and run ID
+ in `run_path`, so no ID needs to be copied by hand.
""")
return
-@app.cell(hide_code=True)
-def _():
- update_config_with_api = mo.ui.run_button(
- label="Update the W&B config with the Public API"
- )
- mo.vstack(
- [
- mo.md(r"""
- ### Updating the `config` with the API
-
- What if your run has finished, but you realized you forgot to log something?
-
- Never fear, you can always use the
- [public API](https://docs.wandb.ai/ref/python/public-api)
- to update your `config`
- (or anything else about your run!)
- at any time. You just need to know the details of the `run` you want to update.
- """),
- update_config_with_api,
- ]
- )
- return (update_config_with_api,)
-
-
@app.cell
-def _(config_run, update_config_with_api):
- mo.stop(
- not update_config_with_api.value,
- mo.callout(
- mo.md("Click the button above to add `bar` through the W&B Public API."),
- kind="info",
- ),
- )
+def _(args, config, entity, project):
+ with wandb.init(entity=entity, project=project, config=config) as run:
+ run.config.epochs = 4
+ run.config["batch_size"] = 32
+ run.config.update(args)
+ run_path = f"{run.entity}/{run.project}/{run.id}"
+ saved_config = dict(run.config)
- api = wandb.Api()
-
- # pulling the relevant info automatically from the run object
- # this can also be found on the website
- username = config_run.entity
- project = config_run.project
- run_id = config_run.id
-
- api_run = api.run(f"{username}/{project}/{run_id}")
- api_run.config["bar"] = 32
- api_run.update()
- return
-
-
-@app.cell
-def _():
- mo.md(r"""
- Here's what the final Config panel looks like:
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## Using `config` for great good!
- The `config` parameters are useful for performing grouping, filtering, and aggregating on your experiments and their results.
-
- ### The examples below come from the project [here](https://wandb.ai/wandb/DistHyperOpt).
- """)
- return
+ saved_config
+ return (run_path,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Filtering Runs
- Filter tab allows you to display the runs that quality one or more conditions. These conditions can be formed by applying relational operators to any of the parameters logged in the `config` file.
-
- [Our example project](https://wandb.ai/wandb/DistHyperOpt) compares various hyper-parameter tuning methods and has more than 80 runs. Each run has a "Job Type" logged in the `config` which corresponds
-
- Let's say you want to visualize only the ones that are generated by a particular tuning algorithm, like Population Based Traing (`pbt`). You can do that by applying a filter on "Job Type".
+ ## Add metadata after a run has finished
- Run the cell below to see this in action!
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _():
- mo.Html(r"""
-
-
+ If you forgot to record a setting, use the
+ [Weights & Biases Public API](https://docs.coreweave.com/products/wandb/ref/python/public-api)
+ to retrieve the completed run at `entity/project/run_id`. This example
+ uses the `run_path` captured above and adds a `purpose` field to that run.
+ Calling `update()` saves the edited configuration to Weights & Biases.
""")
return
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ### Grouping Runs
- You can group your experiments in the dashboard of your project based on a particular column from `config`. A common use case for this would be grouping sub-experiments within a larger project.
-
- Our runs are grouped based on "Job Type". The Group tab is located next to the Filter Tab. You can group your runs by any parameter present in the config.
-
- 
- """)
+@app.cell
+def _(run_path):
+ api = wandb.Api()
+ completed_run = api.run(run_path)
+ completed_run.config["purpose"] = "configuration tutorial"
+ completed_run.update()
+ completed_run.config
return
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Parallel Coordinates Chart
-
- Often, the main thing we want to do with a group of Runs is make comparisons.
-
- The W&B Dashboard includes a Chart type for exactly this purpose:
- the Parallel Coordinates chart.
+ ## Verify and compare experiments
- A Parallel Coordinates chart represents each Run in the group as a line.
- This line passes through as many of the `config` values
- or logged metrics as you like,
- and is colored by its value on a single metric.
- This lets you take in, at a glance,
- which hyperparameter configurations were most and least successful.
- See the example below.
+ Open the run link printed by `wandb.init()` and find **Config** on the
+ run's **Overview** page. Confirm that it contains `dataset="CelebA"`,
+ `type="baseline"`, `epochs=4`, `batch_size=32`, `batch_per_gpu=8`, and
+ `weight_decay=0.1` (or your chosen overrides). Refresh after the Public API
+ step to see `purpose="configuration tutorial"` as well.
- Head to a [group of Runs in this project](https://wandb.ai/wandb/DistHyperOpt/groups/dcgan_train)
- and build a Parallel Coordinates chart like the one pictured below
- by
- 1. clicking the + sign in the top-right corner, aligned with "Charts",
- 2. selecting "Parallel Coordinates" from the available Charts, and
- 3. adding the columns in the image, in order.
+ Rerun the example with different settings, then open the project's runs
+ table to filter or group by a config field such as `type` or `batch_per_gpu`.
+ In a training experiment, log metrics with `run.log()` to compare outcomes
+ across configurations or build a parallel coordinates chart.
- 
+ Explore the [example project](https://wandb.ai/wandb/DistHyperOpt) for a
+ larger comparison, or read
+ [Configure experiments](https://docs.coreweave.com/products/wandb/track/config)
+ for more configuration patterns.
""")
return
diff --git a/marimo/convert/credit-scorecards-with-xgboost-and-w-b/credit_scorecards_with_xgboost_and_w_b.py b/marimo/convert/credit-scorecards-with-xgboost-and-w-b/credit_scorecards_with_xgboost_and_w_b.py
index 9d06fd09..dfb3e22d 100644
--- a/marimo/convert/credit-scorecards-with-xgboost-and-w-b/credit_scorecards_with_xgboost_and_w_b.py
+++ b/marimo/convert/credit-scorecards-with-xgboost-and-w-b/credit_scorecards_with_xgboost_and_w_b.py
@@ -1,7 +1,6 @@
# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
-# "dill>=0.3.8,<1",
# "marimo>=0.24.2",
# "numpy>=1.26,<3",
# "pandas>=2.1,<3",
@@ -10,151 +9,73 @@
# "wandb>=0.19.10,<1",
# "xgboost>=2.0,<4",
# ]
+#
+# [tool.marimo.runtime]
+# on_cell_change = "lazy"
# ///
-"""Build an interpretable vehicle-loan scorecard with XGBoost and W&B."""
+"""Track a vehicle-loan classification pipeline with XGBoost and Weights & Biases."""
import marimo
-__generated_with = "0.24.2"
+__generated_with = "0.25.1"
app = marimo.App(
width="medium",
- app_title="Vehicle loan scorecards with XGBoost and W&B",
+ app_title="Vehicle loan classification with XGBoost and Weights & Biases",
)
with app.setup:
- import importlib.util
- import sys
+ import inspect
import tempfile
from pathlib import Path
- from dill import detect
- from dill.source import getsource
import marimo as mo
import numpy as np
import pandas as pd
+ import xgboost as xgb
from scipy.stats import ks_2samp
from sklearn import metrics, model_selection
-
import wandb
-
- pd.set_option("display.max_columns", None)
- data_root = Path(tempfile.gettempdir()) / "wandb-credit-scorecard"
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/credit-scorecards-with-xgboost-and-w-b/credit_scorecards_with_xgboost_and_w_b.py/server)
-
- - -
+
+ # Vehicle Loan Default Prediction with XGBoost and Weights & Biases
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## Data
- """)
- return
-
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/credit-scorecards-with-xgboost-and-w-b/credit_scorecards_with_xgboost_and_w_b.py/server)
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ### AWS S3, Google Cloud Storage and W&B Artifacts
- """)
- return
+ This lesson tracks a classification pipeline from raw data to trained model.
+ We use Weights & Biases Artifacts to record dataset lineage, Tables to inspect sample rows,
+ the XGBoost callback to log training, and Sweeps to compare hyperparameters.
+ The example uses a subset of the
+ [Vehicle Loan Default Prediction dataset](https://www.kaggle.com/sneharshinde/ltfs-av-data)
+ and preprocessing code shared by Scotiabank in the
+ [RAPIDS credit-scorecard example](https://github.com/rapidsai-community/showcase/tree/main/event_notebooks/GTC_2021/credit_scorecard).
+ The teaching objective is to track the data, settings, and results of a
+ classification experiment. The code does not construct or validate a
+ points-based lending scorecard.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- 
+ > On macOS, XGBoost also needs OpenMP. If importing XGBoost reports a missing
+ `libomp.dylib`, install it with `brew install libomp` and restart.
""")
return
@@ -162,889 +83,421 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- Weights and Biases **Artifacts** enable you to log end-to-end training pipelines to ensure your experiments are always reproducible.
+ ## Authenticate with Weights & Biases
- Data privacy is critical to Weights & Biases, so Artifacts can refer to files in private cloud storage such as Amazon S3 or Google Cloud Storage. Self-Managed W&B deployments are also available.
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- By default, W&B stores artifact files in a private Google Cloud Storage bucket located in the United States. All files are encrypted at rest and in transit. For sensitive files, we recommend a private W&B installation or the use of reference artifacts.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- ### Artifacts Reference Example
- **Create an artifact with the S3/GCS metadata**
+ ## Set your team and project
- The artifact only consists of metadata about the S3/GCS object such as its ETag, size, and version ID (if object versioning is enabled on the bucket).
-
- ```
- run = wandb.init()
- artifact = wandb.Artifact('mnist', type='dataset')
- artifact.add_reference('s3://my-bucket/datasets/mnist')
- run.log_artifact(artifact)
- ```
-
- **Download the artifact locally when needed**
-
- W&B will use the metadata recorded when the artifact was logged to retrieve the files from the underlying bucket.
-
- ```
- artifact = run.use_artifact('mnist:latest', type='dataset')
- artifact_dir = artifact.download()
- ```
-
- See [Create an artifact](https://docs.wandb.ai/models/artifacts/construct-an-artifact) for more on using external references and configuring credentials.
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `vehicle_loan_default`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- mo.md(r"""
- ## Authentication
-
- Enter your [W&B API key](https://wandb.ai/authorize), or leave it blank to
- use `WANDB_API_KEY` from the molab Secrets panel or credentials already
- stored in this runtime. A fresh molab session does not inherit credentials
- from your computer.
-
- The entity is the team name in `wandb.ai/
+
+
- ## Why should I use W&B?
+ # Visualize LLM training progress with Weights & Biases Tables
-
- """)
- return
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Why should I use W&B?
+
+
+
+ # Hugging Face + Weights & Biases
- Use [Weights & Biases](https://wandb.com) for machine learning experiment tracking, dataset versioning, and project collaboration.
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/huggingface-wandb/huggingface_wandb.py/server)
-
""")
return
@@ -90,258 +81,206 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Getting started
-
- In this tutorial, we work with the Hugging Face and Weights & Biases libraries, and the GLUE dataset and training script.
- - [Hugging Face Transformers](https://github.com/huggingface/transformers): Natural language models and datasets
- - [Weights & Biases](https://docs.wandb.com/): Experiment tracking and visualization
- - [GLUE dataset](https://gluebenchmark.com/): A language understanding benchmark dataset
- - [GLUE script](https://github.com/huggingface/transformers/blob/master/examples/run_glue.py): Model training script for sequence classification
- """)
- return
+ ## Run with the notebook's dependencies
+ From this notebook's directory, launch it with its declared Python and package versions:
-@app.cell
-def _():
- run_glue_url = (
- "https://raw.githubusercontent.com/huggingface/transformers/"
- "refs/heads/main/examples/pytorch/text-classification/run_glue.py"
- )
- run_glue_path = "run_glue.py"
+ ```bash
+ uvx marimo edit --sandbox huggingface_wandb.py
+ ```
- with fsspec.open(run_glue_url, "rb") as _source:
- with open(run_glue_path, "wb") as _destination:
- _destination.write(_source.read())
- return (run_glue_path,)
+ The training script uses the same Python interpreter as the notebook.
+ Opening this file in an existing marimo environment does not install the
+ dependencies listed at the top of the file. That environment must also
+ contain `evaluate`, `accelerate`, `datasets`, PyTorch, and Transformers.
+ The sandbox command installs the declared dependencies in an isolated environment.
+ """)
+ return
@app.cell(hide_code=True)
def _():
mo.md(r"""
- In the Hugging Face Transformers repo, we've instrumented the Trainer to automatically log training and evaluation metrics to W&B at each logging step.
+ ## Authenticate with Weights & Biases
- Here's an in depth look at how the integration works: [Hugging Face + W&B Report](https://app.wandb.ai/jxmorris12/huggingface-demo/reports/Train-a-model-with-Hugging-Face-and-Weights-%26-Biases--VmlldzoxMDE2MTU).
- """)
- return
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## Authentication
+ ## Set your team and project
- Enter your [W&B API key](https://wandb.ai/authorize) and, if needed, your team or entity. You can leave the key blank when this environment already has W&B credentials.
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `huggingface-demo`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- _api_key_input = mo.ui.text(
- kind="password",
- label="W&B API key (optional)",
- placeholder="Paste a key or use cached credentials",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- label="W&B entity or team (optional)",
- placeholder="Leave blank to use your default entity",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}")
- .batch(api_key=_api_key_input, entity=_entity_input)
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="huggingface-demo", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(wandb_login_form):
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Connect to W&B above before running the training example."),
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _api_key = wandb_login_form.value["api_key"].strip()
- _entity = wandb_login_form.value["entity"].strip()
try:
_login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
- _login_error = None
- except wandb.errors.Error as _error:
+ except (wandb.errors.Error, ValueError):
_login_ok = False
- _login_error = str(_error)
mo.stop(
not _login_ok,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. "
- f"Check the API key and try again. \n\nW&B reported: `{_login_error or 'unknown error'}`"
- ),
- kind="danger",
- ),
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- wandb_settings = {
- "project": "huggingface-demo",
- "entity": _entity or None,
- }
- mo.callout(
- mo.md("Connected to W&B."),
- kind="success",
- )
- return (wandb_settings,)
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
@app.cell(hide_code=True)
def _():
mo.md(r"""
- Optionally, we can set environment variables to customize W&B logging. See [documentation](https://docs.wandb.com/library/integrations/huggingface).
- """)
- return
-
+ ## Download the training script
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## Train the model
- Next, call the downloaded training script [run_glue.py](https://huggingface.co/transformers/examples.html#glue) and see training automatically get tracked to the Weights & Biases dashboard. This script fine-tunes BERT on the Microsoft Research Paraphrase Corpus— pairs of sentences with human annotations indicating whether they are semantically equivalent.
+ Download `run_glue.py` from the same Transformers release installed by this
+ notebook. Matching the versions keeps the script's arguments and APIs
+ compatible with the library. The script is saved in temporary storage so
+ it can be launched in the next steps.
""")
return
-@app.cell(hide_code=True)
-def _():
- gpu_available = torch.cuda.is_available()
-
- (
- mo.callout(
- mo.md(
- f"GPU ready: **{torch.cuda.get_device_name(0)}**. "
- "The training script will use it automatically."
- ),
- kind="success",
- )
- if gpu_available
- else mo.callout(
- mo.md(
- "No GPU is attached to this session. In Molab, click the "
- "notebook specs button in the header, attach a GPU, then save "
- "and restart. Reconnect to W&B after the restart. Training is "
- "paused to prevent an unexpectedly slow CPU run."
- ),
- kind="warn",
- title="GPU not available",
- )
- )
- return (gpu_available,)
-
-
@app.cell
-def _(wandb_settings):
- task_name = "MRPC"
- wandb_run_id = uuid.uuid4().hex
- wandb_run_entity = (
- wandb_settings["entity"] or wandb.Api().default_entity
- )
-
- run_environment = os.environ.copy()
- run_environment.update(
- {
- "WANDB_RUN_ID": wandb_run_id,
- "WANDB_ENTITY": wandb_run_entity,
- "WANDB_PROJECT": wandb_settings["project"],
- }
+def _():
+ transformers_version = version("transformers")
+ run_glue_url = (
+ "https://raw.githubusercontent.com/huggingface/transformers/"
+ f"v{transformers_version}/examples/pytorch/text-classification/run_glue.py"
)
- wandb_run_url = wandb.Settings(
- entity=wandb_run_entity,
- project=wandb_settings["project"],
- run_id=wandb_run_id,
- ).run_url
- return run_environment, task_name, wandb_run_url
-
+ run_glue_path = Path(tempfile.mkdtemp(prefix="huggingface-glue-")) / "run_glue.py"
-@app.cell
-def _(gpu_available, run_environment, run_glue_path, task_name):
- mo.stop(not gpu_available)
-
- subprocess.run(
- [
- sys.executable,
- run_glue_path,
- "--model_name_or_path",
- "bert-base-uncased",
- "--task_name",
- task_name,
- "--do_train",
- "--do_eval",
- "--max_seq_length",
- "256",
- "--per_device_train_batch_size",
- "32",
- "--learning_rate",
- "2e-4",
- "--num_train_epochs",
- "3",
- "--output_dir",
- f"/tmp/{task_name}/",
- "--report_to",
- "wandb",
- "--logging_steps",
- "50",
- ],
- env=run_environment,
- check=True,
- )
- return
+ with urlopen(run_glue_url, timeout=60) as source:
+ run_glue_path.write_bytes(source.read())
+ return (run_glue_path,)
@app.cell(hide_code=True)
-def _(wandb_run_url):
- mo.md(f"""
- ## Visualize results in dashboard
-
- [**Open this training run in W&B**]({wandb_run_url})
-
- Click the link above, or go to [wandb.ai](https://app.wandb.ai) to see your results stream in live. The link to see your run in the browser will appear after all the dependencies are loaded — look for the following output: "**wandb**: View run at [URL to your unique run]"
-
- **Visualize Model Performance**
- It's easy to look across dozens of experiments, zoom in on interesting findings, and visualize highly dimensional data.
+def _():
+ mo.md(r"""
+ ## Configure Weights & Biases for the training process
- 
+ The script runs in a separate Python process. The next cell copies the runtime
+ environment and adds your selected project and team. Existing environment
+ credentials are inherited, and credentials saved by the SDK are available
+ to the child process in the same runtime.
- **Compare Architectures**
- Here's an example comparing [BERT vs DistilBERT](https://app.wandb.ai/jack-morris/david-vs-goliath/reports/Does-model-size-matter%3F-Comparing-BERT-and-DistilBERT-using-Sweeps--VmlldzoxMDUxNzU) — it's easy to see how different architectures effect the evaluation accuracy throughout training with automatic line plot visualizations.
- 
+ `WANDB_WATCH="all"` adds parameter and gradient histograms to the run.
+ Set it to `"false"` if you only need metrics and want to reduce logging
+ overhead. The Trainer creates the Weights & Biases run when training starts.
""")
return
+@app.cell
+def _(entity, project):
+ run_environment = os.environ.copy()
+ run_environment["WANDB_PROJECT"] = project
+ run_environment["WANDB_WATCH"] = "all"
+ if entity is not None:
+ run_environment["WANDB_ENTITY"] = entity
+ return (run_environment,)
+
+
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Track key information effortlessly by default
- Weights & Biases saves a new run for each experiment. Here's the information that gets saved by default:
- - **Hyperparameters**: Settings for your model are saved in Config
- - **Model Metrics**: Time series data of metrics streaming in are saved in Log
- - **Terminal Logs**: Command line outputs are saved and available in a tab
- - **System Metrics**: GPU and CPU utilization, memory, temperature etc.
+ ## Fine-tune BERT on MRPC
+
+ The arguments below select `bert-base-uncased`, tokenize sentence pairs to
+ at most 256 tokens, and train for three epochs with a batch size of 32 and
+ learning rate of `2e-4`. `--logging_steps 50` records training metrics every
+ 50 optimizer steps, while `--eval_strategy epoch` evaluates on the validation
+ set after each epoch. MRPC evaluation reports accuracy and F1.
+
+ `--report_to wandb` enables experiment tracking. The process prints a link
+ to the run as training begins. It uses the same Python environment as the
+ notebook and automatically selects an available accelerator or CPU.
+
+ Each execution saves the final model and metric files in a new temporary
+ output directory. For a quick check, add `--max_train_samples 64` and
+ `--max_eval_samples 32`, reduce `--num_train_epochs` to `1`, and set
+ `--logging_steps` to `1`. Those small samples help check that the workflow runs. Use a larger, fixed
+ validation set when comparing model quality.
""")
return
+@app.cell
+def _(run_environment, run_glue_path):
+ output_dir = tempfile.mkdtemp(prefix="bert-mrpc-")
+ training_command = [
+ sys.executable,
+ str(run_glue_path),
+ "--model_name_or_path", "bert-base-uncased",
+ "--task_name", "mrpc",
+ "--do_train",
+ "--do_eval",
+ "--max_seq_length", "256",
+ "--per_device_train_batch_size", "32",
+ "--learning_rate", "2e-4",
+ "--num_train_epochs", "3",
+ "--logging_steps", "50",
+ "--eval_strategy", "epoch",
+ "--save_strategy", "no",
+ "--output_dir", output_dir,
+ "--run_name", "bert-base-mrpc",
+ "--report_to", "wandb",
+ ]
+
+ subprocess.run(training_command, env=run_environment, check=True)
+ return
+
+
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Learn more!
- - [Documentation](https://docs.wandb.ai/tutorials/huggingface/): docs on the Weights & Biases and Hugging Face integration
- - [Videos](http://wandb.me/youtube): tutorials, interviews with practitioners, and more on our YouTube channel
- - Contact: Message us at contact@wandb.com with questions
+ ## Verify the run and compare experiments
+
+ Open the run link printed by Weights & Biases during training, or find the run named
+ `bert-base-mrpc` in your chosen project. In the run's workspace, inspect
+ `train/loss`, `eval/loss`, `eval/accuracy`, and `eval/f1`. The training step
+ is the horizontal axis, and evaluation points appear after each epoch.
+ Check the run's config for settings such as `learning_rate`,
+ `num_train_epochs`, and `per_device_train_batch_size`. The system metrics
+ show resource usage. Parameter and gradient histograms are logged
+ periodically when `WANDB_WATCH` is `"all"`.
+
+ Try another learning rate or change `--model_name_or_path` to
+ `distilbert-base-uncased`, give the experiment a descriptive `--run_name`,
+ and rerun training. Compare the runs' validation accuracy and F1 alongside
+ training time in the project workspace.
+
+ See the [Hugging Face + Weights & Biases Report](https://app.wandb.ai/jxmorris12/huggingface-demo/reports/Train-a-model-with-Hugging-Face-and-Weights-%26-Biases--VmlldzoxMDE2MTU)
+ for an example dashboard, and the
+ [BERT versus DistilBERT report](https://app.wandb.ai/jack-morris/david-vs-goliath/reports/Does-model-size-matter%3F-Comparing-BERT-and-DistilBERT-using-Sweeps--VmlldzoxMDUxNzU)
+ for a comparison of architectures.
""")
return
diff --git a/marimo/convert/hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py b/marimo/convert/hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py
index 81c40dc8..5af49da0 100644
--- a/marimo/convert/hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py
+++ b/marimo/convert/hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py
@@ -6,14 +6,18 @@
# "tensorflow>=2.16,<3",
# "wandb>=0.18,<1",
# ]
+#
+# [tool.marimo.runtime]
+# on_cell_change = "lazy"
# ///
-"""Optimize a TensorFlow classifier with W&B Sweeps."""
import marimo
-__generated_with = "0.24.2"
-app = marimo.App(width="medium", app_title="TensorFlow hyperparameter sweeps with W&B")
-
+__generated_with = "0.25.1"
+app = marimo.App(
+ width="medium",
+ app_title="TensorFlow hyperparameter sweeps with Weights & Biases",
+)
with app.setup:
import marimo as mo
@@ -28,31 +32,36 @@
@app.cell(hide_code=True)
def _():
mo.md(r"""
- # Weights & Biases Sweep + TensorFlow 2.x
-
- [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py/server)
+
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
-
+
+
# Introduction to Weights & Biases with Keras
[](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/intro-to-weights-biases-keras/intro_to_weights_biases_keras.py/server)
-
+
+
+ # Introduction to Weights & Biases
[](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/intro-to-weights-biases/intro_to_weights_biases.py/server)
- """)
- return
+ A Weights & Biases run records one experiment's configuration, metrics, and outputs.
+ This lesson starts with simulated metrics, then tracks real PyTorch training,
+ prediction tables, and model checkpoints. An optional final example shows
+ how to send a Weights & Biases alert.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
-
+
- - More info and customization details: [Confusion Matrix](https://wandb.ai/wandb/plots/reports/Confusion-Matrix--VmlldzozMDg1NTM)
- - More examples in this W&B project: [Custom Charts](https://app.wandb.ai/demo-team/custom-charts).
+ # Plot a Confusion Matrix with Weights & Biases
- This notebook explores a transfer learning problem: finetuning InceptionV3 with ImageNet weights to identify 10 types of living things (birds, plants, insects, etc) from 10K photos via [iNaturalist 2017](https://github.com/visipedia/inat_comp).
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/log-a-confusion-matrix-with-w-b/log_a_confusion_matrix_with_w_b.py/server)
+
+ Learn to log a confusion matrix after each training epoch and use it to
+ identify which classes a model confuses. The example trains a classification
+ head on a frozen InceptionV3 backbone using a ten-class sample from
+ [iNaturalist 2017](https://github.com/visipedia/inat_comp).
- Note: Hyperparameters like number of epochs and training dataset size are set to minimum values here for demo efficiency. On the full training data, the model should get to the low 80s in validation accuracy within an epoch or so.
+ The small training defaults are intended to demonstrate logging. They do
+ not establish an expected accuracy. Inspect the confusion matrix alongside
+ validation loss and accuracy before deciding what to change.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ Edit these values to compare runs. Training and validation use whole batches,
+ rounded up to cover the requested example counts (up to the dataset size).
+ Increase `num_train` or the epoch count to test whether additional training
+ improves validation performance.
+ The confusion matrix uses the first `num_log_batches` validation batches each
+ epoch. Increase this value to include more examples in the chart.
""")
return
@@ -71,105 +101,100 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Authentication
+ ## Authenticate with Weights & Biases
- Enter a [W&B API key](https://wandb.ai/authorize) and an optional team entity below. Leave the key blank when `WANDB_API_KEY` is configured in the Molab Secrets panel or this runtime already has W&B credentials. If needed, [create a free W&B account](https://wandb.ai/signup).
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- Your entity is the first path segment after `wandb.ai` when you open a W&B workspace. The API key is used only for `wandb.login()` and is never included in the run configuration or notebook output.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
+
+ ## Set your team and project
+
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `confusion_matrix`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- _api_key_input = mo.ui.text(
- kind="password",
- label="W&B API key (optional)",
- placeholder="Paste a key or use configured credentials",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- label="W&B entity or team (optional)",
- placeholder="Leave blank to use your default entity",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}")
- .batch(api_key=_api_key_input, entity=_entity_input)
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="confusion_matrix", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(wandb_login_form):
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Connect to W&B before configuring the experiment."),
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _api_key = wandb_login_form.value["api_key"].strip()
- _entity = wandb_login_form.value["entity"].strip()
try:
_login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
- _login_error = None
- except wandb.errors.Error as _error:
+ except (wandb.errors.Error, ValueError):
_login_ok = False
- _login_error = str(_error)
mo.stop(
not _login_ok,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. Check the API key or "
- "Molab Secrets configuration and submit again.\n\n"
- f"W&B reported: `{_login_error or 'unknown error'}`"
- ),
- kind="danger",
- ),
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- wandb_settings = {
- "project": PROJECT_NAME,
- "entity": _entity or None,
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
+
+
+@app.cell
+def _():
+ training_config = {
+ "num_train": 100,
+ "num_val": 500,
+ "fc_size": 1024,
+ "img_width": 299,
+ "img_height": 299,
+ "batch_size": 32,
+ "pretrain_epochs": 1,
+ "num_log_batches": 15,
+ "random_seed": 23,
}
- _entity_note = (
- f" Runs will be logged to `{wandb_settings['entity']}`."
- if wandb_settings["entity"]
- else " Runs will use your default W&B entity."
- )
- mo.callout(mo.md(f"Connected to W&B.{_entity_note}"), kind="success")
- return (wandb_settings,)
+ return (training_config,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Setup: Access the remote dataset
-
- The iNaturalist sample contains 10,000 training images and 2,000 validation images, evenly distributed across 10 classes of living things like birds, insects, plants, and mammals. **The archive is approximately 3.6 GB, so the first transfer and extraction can take several minutes.**
-
- The `dataset_fs` value below exposes the public archive through an `fsspec` reference filesystem, so you can browse it in marimo's data side panel without transferring the whole file.
-
- Training still requires local image directories. The first time you submit the training form, `ensure_dataset()` streams the archive into a temporary runtime cache and extracts it there. Later submissions reuse that cache; no dataset files are written beside the notebook.
+ ## Prepare the iNaturalist sample
+
+ The [iNaturalist 2017](https://github.com/visipedia/inat_comp) sample contains
+ 10,000 training images and 2,000 validation images across 10 classes.
+ **The archive is about 3.6 GB**, even when training on a small sample.
+ The cells below expose the archive and define a download helper. The call
+ to `ensure_dataset()` downloads and extracts it. Later cells load ImageNet
+ weights and start training. Download and training time depend on your network
+ and hardware.
+
+ `dataset_fs` exposes the public archive in marimo's data side panel without
+ downloading it. `ensure_dataset()` streams it into a temporary runtime cache
+ and reuses the extracted directories on later runs.
""")
return
@app.cell
def _():
- # Expose the remote dataset as a named filesystem so it appears in
- # marimo's data side panel.
dataset_fs = fsspec.filesystem(
"reference",
- fo={
- "version": 1,
- "refs": {"nature_12K.zip": [DATASET_URL]},
- },
+ fo={"version": 1, "refs": {"nature_12K.zip": [DATASET_URL]}},
)
return (dataset_fs,)
@@ -182,7 +207,7 @@ def ensure_dataset(dataset_fs):
validation_path = RUNTIME_CACHE_ROOT / "inaturalist_12K" / "val"
if train_path.is_dir() and validation_path.is_dir():
- return str(train_path), str(validation_path), "Reused the extracted dataset."
+ return str(train_path), str(validation_path)
RUNTIME_CACHE_ROOT.mkdir(parents=True, exist_ok=True)
if not archive_path.exists():
@@ -203,19 +228,64 @@ def ensure_dataset(dataset_fs):
"inaturalist_12K/val."
)
- return str(train_path), str(validation_path), "Streamed and extracted the remote dataset."
+ return str(train_path), str(validation_path)
+
+
+@app.cell
+def _(dataset_fs):
+ train_data, val_data = ensure_dataset(dataset_fs)
+ return train_data, val_data
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Training code
+ ## Build the data generators and model
- Feel free to try different values for "NUM_TRAIN" and "NUM_EPOCHS" below so you can see a variety of confusion matrices (generally better ones with more training examples/longer training time).
+ InceptionV3 expects pixels scaled to `[-1, 1]`, so both generators use its
+ `preprocess_input()` function. Only training images are augmented and shuffled.
+ Validation keeps a fixed order, and class names follow the generator's numeric
+ IDs so each row and column of the confusion matrix has the correct label.
+
+ `build_model()` freezes the ImageNet backbone and trains a new classification
+ head. Only the new head’s weights are updated in this example.
""")
return
+@app.cell
+def _(train_data, training_config, val_data):
+ tf.keras.utils.set_random_seed(training_config["random_seed"])
+
+ train_datagen = ImageDataGenerator(
+ preprocessing_function=preprocess_input,
+ shear_range=0.2,
+ zoom_range=0.2,
+ horizontal_flip=True,
+ )
+ validation_datagen = ImageDataGenerator(preprocessing_function=preprocess_input)
+
+ train_generator = train_datagen.flow_from_directory(
+ train_data,
+ target_size=(training_config["img_height"], training_config["img_width"]),
+ batch_size=training_config["batch_size"],
+ class_mode="categorical",
+ seed=training_config["random_seed"],
+ )
+
+ class_names = list(train_generator.class_indices)
+
+ val_generator = validation_datagen.flow_from_directory(
+ val_data,
+ classes=class_names,
+ target_size=(training_config["img_height"], training_config["img_width"]),
+ batch_size=training_config["batch_size"],
+ class_mode="categorical",
+ shuffle=False,
+ )
+ return class_names, train_generator, val_generator
+
+
@app.function
def build_model(fc_size, num_classes):
"""Load InceptionV3 with ImageNet weights, freeze it,
@@ -244,142 +314,50 @@ def build_model(fc_size, num_classes):
return model
-@app.function
-def pretrain(config_defaults, wandb_settings, train_data, val_data):
- """Main training loop.
-
- This is called "pretrain" because it freezes the InceptionV3 layers of the
- model and only trains the new top layers on the new data. A subsequent
- training phase would unfreeze all the layers and finetune the whole model
- on the new data.
- """
- if wandb.run is not None:
- wandb.finish()
-
- checkpoint_cache = RUNTIME_CACHE_ROOT / "checkpoints"
- checkpoint_cache.mkdir(parents=True, exist_ok=True)
- with tempfile.TemporaryDirectory(
- prefix="confusion-matrix-", dir=checkpoint_cache
- ) as checkpoint_directory:
- checkpoint_path = Path(checkpoint_directory) / "checkpoint.keras"
- # The checkpoint directory remains alive until the run context finishes
- # and W&B has flushed the checkpoint Artifact.
- with wandb.init(
- project=wandb_settings["project"],
- entity=wandb_settings["entity"],
- config=config_defaults,
- ) as run:
- cfg = run.config
-
- # Set the random seed.
- tf.random.set_seed(cfg.random_seed)
-
- # Also set the NumPy seed to control the train/validation data.
- np.random.seed(cfg.random_seed)
-
- # Create train and validation data generators.
- train_datagen = ImageDataGenerator(
- rescale=1.0 / 255,
- shear_range=0.2,
- zoom_range=0.2,
- horizontal_flip=True,
- )
- validation_datagen = ImageDataGenerator(rescale=1.0 / 255)
-
- train_generator = train_datagen.flow_from_directory(
- train_data,
- target_size=(cfg.img_width, cfg.img_height),
- batch_size=cfg.batch_size,
- class_mode="categorical",
- )
- val_generator = validation_datagen.flow_from_directory(
- val_data,
- target_size=(cfg.img_width, cfg.img_height),
- batch_size=cfg.batch_size,
- class_mode="categorical",
- )
-
- # Instantiate the model and callbacks.
- model = build_model(cfg.fc_size, cfg.num_classes)
- callbacks = [
- WandbMetricsLogger(),
- WandbModelCheckpoint(str(checkpoint_path)),
- PRMetrics(
- run,
- val_generator,
- cfg.num_log_batches,
- ),
- ]
-
- # Train!
- history = model.fit(
- train_generator,
- steps_per_epoch=max(1, cfg.num_train // cfg.batch_size),
- epochs=cfg.pretrain_epochs,
- validation_data=val_generator,
- callbacks=callbacks,
- validation_steps=max(1, cfg.num_val // cfg.batch_size),
- )
-
- final_validation_accuracy = history.history.get(
- "val_accuracy", [None]
- )[-1]
- result = {
- "run_url": run.url,
- "final_validation_accuracy": (
- float(final_validation_accuracy)
- if final_validation_accuracy is not None
- else None
- ),
- }
-
- return result
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Log a confusion matrix after each epoch
+
+ The custom Keras callback collects validation images and their labels from
+ the same batches, predicts class probabilities, and takes `argmax` to get
+ class IDs. `wandb.plot.confusion_matrix()` builds a Weights & Biases chart from those IDs
+ and the true labels. `training_run.log()` saves it under the stable key
+ `conf_mat`, making it easy to compare epochs and runs. Diagonal entries are
+ correct predictions. Off-diagonal entries show which classes are confused.
+ The matrix covers the validation batches selected by `num_log_batches`,
+ which may be fewer examples than the validation metrics cover.
+ """)
+ return
@app.class_definition
-class PRMetrics(Callback):
- """Custom callback to compute metrics at the end of each training epoch."""
+class ConfusionMatrixLogger(Callback):
+ """Log predictions for a bounded validation sample after each epoch."""
- def __init__(self, run, generator=None, num_log_batches=1):
+ def __init__(self, training_run, generator, class_names, num_log_batches):
super().__init__()
- self.run = run
+ if num_log_batches < 1:
+ raise ValueError("num_log_batches must be at least 1")
+ self.training_run = training_run
self.generator = generator
- self.num_batches = num_log_batches
-
- # Store full names of classes in class-ID order.
- self.flat_class_names = [
- name
- for name, _class_id in sorted(
- generator.class_indices.items(), key=lambda item: item[1]
- )
- ]
+ self.class_names = class_names
+ self.num_batches = min(num_log_batches, len(generator))
def on_epoch_end(self, epoch, logs=None):
- # Collect validation data and ground truth labels from the generator.
- val_data, val_labels = zip(
+ images, labels = zip(
*(self.generator[index] for index in range(self.num_batches))
)
- val_data, val_labels = np.vstack(val_data), np.vstack(val_labels)
-
- # Use the trained model to generate predictions for the given number
- # of validation data batches (num_batches).
- val_predictions = self.model.predict(val_data, verbose=0)
- ground_truth_class_ids = val_labels.argmax(axis=1)
-
- # Take the argmax for each set of prediction scores to return the class
- # ID of the highest-confidence prediction.
- top_pred_ids = val_predictions.argmax(axis=1)
-
- # Log the confusion matrix. The key "conf_mat" is the ID of the plot;
- # keep it stable so subsequent runs appear on the same plot.
- self.run.log(
+ probabilities = self.model.predict(np.concatenate(images), verbose=0)
+ self.training_run.log(
{
+ "epoch": epoch,
"conf_mat": wandb.plot.confusion_matrix(
probs=None,
- preds=top_pred_ids,
- y_true=ground_truth_class_ids,
- class_names=self.flat_class_names,
- )
+ preds=probabilities.argmax(axis=1),
+ y_true=np.concatenate(labels).argmax(axis=1),
+ class_names=self.class_names,
+ ),
}
)
@@ -387,148 +365,84 @@ def on_epoch_end(self, epoch, logs=None):
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Run the experiment
+ ## Train and save the results to Weights & Biases
- After connecting to W&B above, submit this training form to materialize the remote dataset in the runtime cache if necessary, download the InceptionV3 weights if necessary, train the model, and create one run in the `confusion_matrix` project. A later submission finishes any prior active run before it starts another.
+ `wandb.init()` records the experiment configuration. `WandbMetricsLogger`
+ logs loss and accuracy, `WandbModelCheckpoint` saves a model Artifact each
+ epoch, and `ConfusionMatrixLogger` logs the validation chart.
+
+ The temporary checkpoint directory stays alive until the run finishes and
+ Weights & Biases flushes the Artifact. The run context finishes automatically, including
+ when training raises an error.
""")
return
-@app.cell(hide_code=True)
-def _(wandb_settings):
- _workflow_ready = callable(ensure_dataset) and callable(pretrain)
- _num_train_input = mo.ui.number(
- start=32,
- stop=10_000,
- step=1,
- value=100,
- label="Training examples",
- )
- _num_epochs_input = mo.ui.number(
- start=1,
- stop=20,
- step=1,
- value=1,
- label="Epochs",
- )
- training_form = (
- mo.md("{num_train}\n\n{num_epochs}")
- .batch(
- num_train=_num_train_input,
- num_epochs=_num_epochs_input,
- )
- .form(
- submit_button_label="Prepare data, train, and log to W&B",
- submit_button_disabled=not _workflow_ready,
- bordered=True,
- )
- )
- _destination = (
- f"`{wandb_settings['entity']}/{wandb_settings['project']}`"
- if wandb_settings["entity"]
- else f"`{wandb_settings['project']}` in your default entity"
- )
- mo.vstack(
- [
- mo.md(f"This run will be created in {_destination}."),
- training_form,
- ]
- )
- return (training_form,)
-
-
-@app.cell(hide_code=True)
-def _(training_form, wandb_settings):
- mo.stop(
- training_form.value is None,
- mo.callout(
- mo.md("Review the workflow above, then submit the form to begin."),
- kind="info",
- ),
- )
-
- training_request = {
- "wandb_settings": dict(wandb_settings),
- "num_train": int(training_form.value["num_train"]),
- "num_epochs": int(training_form.value["num_epochs"]),
- }
- return (training_request,)
-
-
@app.cell
-def _(training_request):
- # Experiment configuration saved to W&B
- config_defaults = {
- # Number of images used to train; set low for demo training speed.
- # You can set this up to 10,000 for the full dataset.
- # GOOD CONFIG TO TRY: 100, 500, 1000, 2000
- "num_train": training_request["num_train"],
- # Number of images used to validate; set low for demo training speed.
- # You can set this up to 2,000 for the full dataset.
- "num_val": 500,
- "num_classes": 10,
- "fc_size": 1024,
-
- # InceptionV3 settings
- "img_width": 299,
- "img_height": 299,
- "batch_size": 32,
-
- # Number of epochs; set low for demo training speed.
- # GOOD CONFIG TO TRY: 3, 5, 10
- "pretrain_epochs": training_request["num_epochs"],
- # Number of validation data batches used to compute metrics at the end
- # of each epoch.
- "num_log_batches": 15,
- # Random seed
- "random_seed": 23,
- }
- config_defaults
- return (config_defaults,)
+def _(
+ class_names,
+ entity,
+ project,
+ train_generator,
+ training_config,
+ val_generator,
+):
+ with tempfile.TemporaryDirectory(
+ prefix="confusion-matrix-"
+ ) as checkpoint_directory:
+ with wandb.init(
+ project=project,
+ entity=entity,
+ name="inaturalist-inceptionv3-confusion-matrix",
+ config={**training_config, "num_classes": len(class_names)},
+ ) as training_run:
+ model = build_model(training_config["fc_size"], len(class_names))
+ callbacks = [
+ WandbMetricsLogger(),
+ WandbModelCheckpoint(
+ str(Path(checkpoint_directory) / "checkpoint.keras")
+ ),
+ ConfusionMatrixLogger(
+ training_run,
+ val_generator,
+ class_names,
+ training_config["num_log_batches"],
+ ),
+ ]
+ history = model.fit(
+ train_generator,
+ steps_per_epoch=min(
+ len(train_generator),
+ max(1, int(np.ceil(
+ training_config["num_train"] / training_config["batch_size"]
+ ))),
+ ),
+ epochs=training_config["pretrain_epochs"],
+ validation_data=val_generator,
+ validation_steps=min(
+ len(val_generator),
+ max(1, int(np.ceil(
+ training_config["num_val"] / training_config["batch_size"]
+ ))),
+ ),
+ callbacks=callbacks,
+ )
+ run_url = training_run.url
+ return
@app.cell(hide_code=True)
-def _(config_defaults, dataset_fs):
- train_data, val_data, dataset_status = ensure_dataset(dataset_fs)
- mo.callout(
- mo.md(
- f"{dataset_status} Training will use "
- f"{config_defaults['num_train']:,} images per epoch."
- ),
- kind="success",
- )
- return train_data, val_data
-
-
-@app.cell
-def _(config_defaults, train_data, training_request, val_data):
- # Run this cell to launch your experiment! Charts appear in the run page
- # under "Media" or "Custom Charts", which you may need to expand.
- experiment_result = pretrain(
- config_defaults=config_defaults,
- wandb_settings=training_request["wandb_settings"],
- train_data=train_data,
- val_data=val_data,
- )
- return (experiment_result,)
+def _():
+ mo.md(r"""
+ Open the run in Weights & Biases App UI and find `conf_mat` in the **Custom Charts** panel.
+ Compare the mistakes across classes, then increase the training sample or
+ epoch count and run again.
+ ## Next steps
-@app.cell(hide_code=True)
-def _(experiment_result):
- _accuracy = experiment_result["final_validation_accuracy"]
- _accuracy_note = (
- f" The final validation accuracy was `{_accuracy:.3f}`."
- if _accuracy is not None
- else ""
- )
- mo.callout(
- mo.md(
- f"Training finished.{_accuracy_note} "
- f"[Open the run in W&B]({experiment_result['run_url']}) and find "
- "`conf_mat` under **Custom Charts** to inspect the confusion matrix."
- ),
- kind="success",
- )
+ - More info and customization details: [Confusion Matrix](https://wandb.ai/wandb/plots/reports/Confusion-Matrix--VmlldzozMDg1NTM)
+ - More examples in this Weights & Biases project: [Custom Charts](https://app.wandb.ai/demo-team/custom-charts).
+ """)
return
diff --git a/marimo/convert/log-almost-anything-with-w-b-media/log_almost_anything_with_w_b_media.py b/marimo/convert/log-almost-anything-with-w-b-media/log_almost_anything_with_w_b_media.py
index e87eb763..2db2f74b 100644
--- a/marimo/convert/log-almost-anything-with-w-b-media/log_almost_anything_with_w_b_media.py
+++ b/marimo/convert/log-almost-anything-with-w-b-media/log_almost_anything_with_w_b_media.py
@@ -12,12 +12,11 @@
import marimo
-__generated_with = "0.24.0"
+__generated_with = "0.25.1"
app = marimo.App(auto_download=["html"])
with app.setup:
import io
- import warnings
import fsspec
import marimo as mo
@@ -27,8 +26,6 @@
import soundfile as sf
import wandb
- warnings.filterwarnings("ignore", category=UserWarning)
-
@app.cell(hide_code=True)
def _():
@@ -52,41 +49,21 @@ def _():
Use [Weights & Biases](https://wandb.com) for machine learning experiment tracking, dataset versioning, and project collaboration.
-
+
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/optimize-hugging-face-models-with-weights-biases/optimize_hugging_face_models_with_weights_biases.py/server)
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- [Hugging Face](https://huggingface.co/) provides tools to quickly train neural networks for NLP (Natural Language Processing) on any task (classification, translation, question answering, etc) and any dataset with PyTorch and TensorFlow 2.0.
+ # Track Hugging Face training with Weights & Biases
- Coupled with [Weights & Biases integration](https://docs.wandb.ai/models/integrations/huggingface_transformers), you can quickly train and monitor models for full traceability and reproducibility without any extra line of code! The notebook metadata provides the libraries; use the submitted authentication form below to sign in before starting an experiment.
+ This lesson shows two ways to connect Hugging Face training to Weights & Biases.
+ First, run an existing BERT training script on MRPC. Then build a custom
+ DistilBERT topic-classification workflow with `Trainer`.
- **Note**: To enable logging to W&B, set `report_to` to `wandb` in your `TrainingArguments` or script.
- """)
- return
+ Both examples use `report_to="wandb"` to log training and evaluation metrics.
+ The settings below also enable parameter and gradient histograms and upload
+ the final model as an Artifact. These records help you compare experiments
+ and identify the settings and model files used by each run.
+ ## Environment setup
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- W&B integration with 🤗 Hugging Face can automatically:
- * log your configuration parameters
- * log your losses and metrics
- * log gradients and parameter distributions
- * log your model
- * keep track of your code
- * log your system metrics (GPU, CPU, memory, temperature, etc)
+ The notebook declares its dependencies in the script metadata. Open it with
+ `uvx marimo edit optimize_hugging_face_models_with_weights_biases.py --sandbox`
+ from this notebook's folder to install and use those dependencies.
+ If you use an existing Python environment instead, install the declared
+ dependencies there too. The training subprocess uses the notebook's Python
+ interpreter; it does not install packages from the script metadata.
+
+ The examples use Transformers
+ for models and training, tokenizers to convert text to token IDs, and Datasets
+ to load and prepare examples. The GLUE script also needs Evaluate to compute
+ metrics and Accelerate to run the Trainer.
+
+ Transformers and `run_glue.py` use the same pinned source revision. Keep those
+ versions aligned when updating the notebook so the script’s arguments match
+ the installed library.
""")
return
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## Here's what the W&B interactive dashboard will look like:
+ `WANDB_LOG_MODEL="end"`
+ tells the Transformers callback to log the trained model as a Weights & Biases Artifact.
+ `WANDB_WATCH="all"` enables gradient and parameter histograms. Both examples
+ enable metrics with `report_to="wandb"` and give the run a descriptive name.
- 
+ > Note: Running this notebook downloads BERT and MRPC, then starts a 300-step
+ training run. The custom example downloads Yahoo! Answers and DistilBERT.
""")
return
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## Installation and set-up
-
- We need the following 🤗 Hugging Face libraries:
- * [transformers](https://huggingface.co/transformers/) contains an API for training models and many pre-trained models
- * [tokenizers](https://huggingface.co/docs/tokenizers) is automatically installed by transformers and "tokenize" our data (ie it converts text to sequence of numbers)
- * [datasets](https://huggingface.co/docs/datasets/) contains a rich source of data and common metrics, perfect for prototyping
+ ## Authenticate with Weights & Biases
- The notebook dependency metadata includes these libraries, the current
- Transformers development version required by `run_glue.py`, and `wandb`
- for automatic experiment tracking.
- """)
- return
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Authentication
+ ## Set your team and project
- Enter your [W&B API key](https://wandb.ai/authorize) and, if needed, your
- team or entity. Leave the key blank to use `WANDB_API_KEY` from molab's
- Secrets panel or credentials already configured in this runtime. Merely
- editing the fields does not authenticate, download assets, or create a run.
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `huggingface-demo`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
-def _(mo):
- _api_key_input = mo.ui.text(
- value="",
- kind="password",
- label="W&B API key (blank uses runtime credentials)",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- value="",
- label="W&B entity or team (blank uses your default)",
- full_width=True,
- )
- _project_input = mo.ui.text(
- value="huggingface-demo",
- label="W&B project",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}\n\n{project}")
- .batch(
- api_key=_api_key_input,
- entity=_entity_input,
- project=_project_input,
- )
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+def _():
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="huggingface-demo", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(mo, wandb, wandb_login_form):
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Submit the authentication form to continue."),
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _submitted_login = wandb_login_form.value
- _api_key = _submitted_login["api_key"].strip()
- _requested_entity = _submitted_login["entity"].strip()
- _project = _submitted_login["project"].strip() or "huggingface-demo"
try:
- _login_ok = wandb.login(
- key=_api_key or None,
- relogin=bool(_api_key),
- )
- _resolved_entity = _requested_entity or wandb.Api().default_entity
- _login_error = None
- except wandb.errors.Error as _error:
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
_login_ok = False
- _resolved_entity = None
- _login_error = str(_error)
mo.stop(
- not _login_ok or not _resolved_entity,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. Check the API key and "
- "entity, then submit again.\n\n"
- f"W&B reported: `{_login_error or 'No default entity was found.'}`"
- ),
- kind="danger",
- ),
- )
-
- wandb_settings = {
- "entity": _resolved_entity,
- "project": _project,
- }
- mo.callout(
- mo.md(
- f"Connected to W&B as entity `{_resolved_entity}`. "
- f"Runs will use project `{_project}`."
- ),
- kind="success",
+ not _login_ok,
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- return (wandb_settings,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Configuration tips
-
- W&B integration with Hugging Face can be configured to add extra functionalities:
-
- * auto-logging of final models as artifacts: set the `WANDB_LOG_MODEL` environment variable to `end`
- * log histograms of gradients and parameters: by default gradients are logged, you can also log parameters by setting environment variable `WANDB_WATCH` to `all`
- * set custom run names with `run_name` arg present in scripts or as part of `TrainingArguments`
- * organize runs by project with the `WANDB_PROJECT` environment variable
-
- For more details refer to [W&B + HF integration documentation](https://docs.wandb.ai/models/integrations/huggingface_transformers).
- """)
- return
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Let's log each final trained model.
- """)
- return
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
@app.cell
def _():
- WANDB_LOG_MODEL = "end"
- return (WANDB_LOG_MODEL,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Training a new model the quick way!
-
- When working on a new problem, you should always check [the task guides](https://huggingface.co/docs/transformers/tasks) as there will often be a script that can already solve your task. At a minimum they will be a great source of inspiration for your own custom pipeline.
-
- Let's use the Hugging Face script responsible for training on any GLUE task, such as sequence classification.
- """)
+ os.environ["WANDB_LOG_MODEL"] = "end"
+ os.environ["WANDB_WATCH"] = "all"
return
@app.cell(hide_code=True)
-def _(torch):
- if torch.cuda.is_available():
- torch_device = torch.device("cuda")
- elif torch.backends.mps.is_available():
- torch_device = torch.device("mps")
- else:
- torch_device = torch.device("cpu")
- return (torch_device,)
-
-
-@app.cell(hide_code=True)
-def _(mo, wandb_settings):
- _run_name_input = mo.ui.text(
- value="demo",
- label="W&B run name",
- full_width=True,
- )
- _max_steps_input = mo.ui.number(
- start=30,
- stop=1000,
- value=300,
- step=30,
- label="Maximum training steps",
- )
- quick_training_form = (
- mo.md("{run_name}\n\n{max_steps}")
- .batch(run_name=_run_name_input, max_steps=_max_steps_input)
- .form(
- submit_button_label="Download run_glue.py and train",
- bordered=True,
- )
- )
- _ready_note = (
- f"The submitted run will use W&B project `{wandb_settings['project']}` "
- f"under entity `{wandb_settings['entity']}`."
- )
- mo.vstack([mo.md(_ready_note), quick_training_form])
- return (quick_training_form,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- These scripts are automatically instrumented with logging when `wandb` is installed and logged in.
+ ## Train with an existing Hugging Face script
+
+ The [task guides](https://huggingface.co/docs/transformers/tasks) provide
+ ready-made training workflows. Here we use `run_glue.py` to fine-tune BERT on
+ MRPC, a sentence-pair classification task.
- Just set `report_to` to `wandb` to enable logging through W&B.
+ The next cells download the script from the same source revision as the
+ installed Transformers library, then run it with this Python environment.
+ `--report_to wandb` activates the integration. Weights & Biases prints the run link in
+ the subprocess output. The script finishes its run when the process exits.
- **Note**: This cell can take up to 5 minutes to run.
+ We evaluate and save every 30 steps, retain two checkpoints, and reload the
+ checkpoint with the best accuracy. `WANDB_LOG_MODEL="end"` then logs that
+ model as an Artifact. Files remain in `hf_assets/` and `hf_outputs/` so you
+ can inspect and reuse them. This example can take several minutes.
""")
return
-@app.cell(hide_code=True)
-def _(
- WANDB_LOG_MODEL,
- mo,
- os,
- quick_training_form,
- uuid,
- wandb,
- wandb_settings,
-):
- mo.stop(
- quick_training_form.value is None,
- mo.callout(
- mo.md(
- "Review the settings, then submit the form when you are ready "
- "to download the script and create a W&B run."
- ),
- kind="info",
- ),
- )
- _submitted_quick_run = quick_training_form.value
- _run_name = _submitted_quick_run["run_name"].strip() or "demo"
- _run_id = uuid.uuid4().hex
- quick_run_environment = os.environ.copy()
- quick_run_environment.update(
- {
- "WANDB_ENTITY": wandb_settings["entity"],
- "WANDB_PROJECT": wandb_settings["project"],
- "WANDB_RUN_ID": _run_id,
- "WANDB_NAME": _run_name,
- "WANDB_LOG_MODEL": WANDB_LOG_MODEL,
- }
+@app.cell
+def _():
+ transformers_revision = "02d8fb9784e8f14a1251e4c992cd82a5762417c6"
+ quick_run_name = "bert-mrpc-300-steps"
+ quick_max_steps = 300
+ quick_output_dir = Path("hf_outputs/mrpc")
+ return (
+ quick_max_steps,
+ quick_output_dir,
+ quick_run_name,
+ transformers_revision,
)
- quick_run_url = wandb.Settings(
- entity=wandb_settings["entity"],
- project=wandb_settings["project"],
- run_id=_run_id,
- ).run_url
- quick_training_request = {
- "max_steps": int(_submitted_quick_run["max_steps"]),
- "output_dir": f"/tmp/MRPC-{_run_id}",
- "run_id": _run_id,
- "run_name": _run_name,
- }
- return quick_run_environment, quick_run_url, quick_training_request
@app.cell
-def _(fsspec, quick_training_request):
+def _(transformers_revision):
run_glue_url = (
"https://raw.githubusercontent.com/huggingface/transformers/"
- "refs/heads/main/examples/pytorch/text-classification/run_glue.py"
+ f"{transformers_revision}/examples/pytorch/text-classification/run_glue.py"
)
- run_glue_path = f"/tmp/run_glue-{quick_training_request['run_id']}.py"
- with fsspec.open(run_glue_url, "rb") as _source:
- with open(run_glue_path, "wb") as _destination:
- _destination.write(_source.read())
+ run_glue_path = Path("hf_assets") / f"run_glue-{transformers_revision[:12]}.py"
+ run_glue_path.parent.mkdir(parents=True, exist_ok=True)
+ with fsspec.open(run_glue_url, "rb") as glue_source:
+ run_glue_path.write_bytes(glue_source.read())
return (run_glue_path,)
-@app.cell(hide_code=True)
-def _(mo, quick_run_url):
- mo.callout(
- mo.md(
- f"[Open this training run in W&B]({quick_run_url}) to watch "
- "metrics and system telemetry appear live."
- ),
- kind="info",
- )
- quick_run_link_ready = True
- return (quick_run_link_ready,)
-
-
@app.cell
def _(
- quick_run_environment,
- quick_run_link_ready,
- quick_training_request,
+ entity,
+ project,
+ quick_max_steps,
+ quick_output_dir,
+ quick_run_name,
run_glue_path,
- subprocess,
- sys,
):
- assert quick_run_link_ready
- subprocess.run(
+ # Pass the project and entity to the separate Python process.
+ quick_environment = os.environ.copy()
+ quick_environment["WANDB_PROJECT"] = project
+ if entity:
+ quick_environment["WANDB_ENTITY"] = entity
+ else:
+ quick_environment.pop("WANDB_ENTITY", None)
+
+ quick_training_process = subprocess.run(
[
sys.executable,
- run_glue_path,
- "--report_to",
- "wandb",
- "--model_name_or_path",
- "bert-base-uncased",
- "--task_name",
- "MRPC",
- "--learning_rate",
- "1e-4",
+ str(run_glue_path),
+ "--report_to", "wandb",
+ "--model_name_or_path", "bert-base-uncased",
+ "--task_name", "mrpc",
+ "--learning_rate", "1e-4",
"--do_train",
"--do_eval",
- "--max_steps",
- str(quick_training_request["max_steps"]),
- "--logging_steps",
- "30",
- "--eval_strategy",
- "steps",
- "--output_dir",
- quick_training_request["output_dir"],
- "--run_name",
- quick_training_request["run_name"],
+ "--max_steps", str(quick_max_steps),
+ "--logging_steps", "30",
+ "--eval_strategy", "steps",
+ "--eval_steps", "30",
+ "--save_steps", "30",
+ "--save_total_limit", "2",
+ "--load_best_model_at_end", "true",
+ "--metric_for_best_model", "accuracy",
+ "--output_dir", str(quick_output_dir),
+ "--run_name", quick_run_name,
],
- env=quick_run_environment,
+ env=quick_environment,
check=True,
)
- quick_training_complete = True
- return
-
-
-@app.cell(hide_code=True)
-def _(mo, quick_run_url):
- mo.callout(
- mo.md(
- "Training finished and the script closed its W&B run. "
- f"[Inspect the completed run]({quick_run_url})."
- ),
- kind="success",
- )
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- You just trained a model and can now visualize your metrics in your dashboard!
-
- Use the live run link above to open the
- [run page](https://docs.wandb.ai/models/runs/view-logged-runs) while training is
- in progress or after it finishes.
-
- It should look something like this:
- """)
+ quick_training_complete = quick_training_process.returncode == 0
return
@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- In addition, your model files have been saved and versioned, along with associated metadata (evaluation & training metrics).
-
- Just check the ["Artifacts" tab](https://docs.wandb.ai/models/runs/view-logged-runs#artifacts) on your run page -- it's the one with the ["stacked pucks" icon](https://stackoverflow.com/questions/2822650/why-is-a-database-always-represented-with-a-cylinder).
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## Advanced usage & custom training
-
- Let's create our own logic for a more customized training.
- """)
- return
-
+ Open the Weights & Biases run link printed by the script to inspect training loss, MRPC
+ evaluation metrics, and system telemetry. Open the run's **Artifacts** tab
+ to find the saved model and its evaluation metadata.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### Preparing a dataset
+ ## Advanced usage: customize a Trainer workflow
- The dataset will vary based on the task you work on. Let's work on sequence classification!
+ Now build a topic classifier with DistilBERT and Yahoo! Answers. This path
+ makes dataset preparation, tokenization, model loading, metrics, and the Weights & Biases
+ run lifecycle explicit.
- Our dataset will be composed of sentences and their associated classes. For example if you wanted to identify the subject of a conversation, you could create a dataset such as:
+ A sequence classifier maps text to a class. For example:
- input | class
- --- | ---
- The team scored a goal in the last seconds | sports
- The debate was heated between the 2 parties | politics
- I've never tasted croissants so delicious ! | food
+ | Input | Class |
+ | --- | --- |
+ | The team scored a goal in the last seconds | sports |
+ | The debate was heated between the two parties | politics |
+ | I've never tasted croissants so delicious! | food |
- The objective of our trained model will be to correctly identify the class associated to new sentences.
+ The [Datasets Hub](https://huggingface.co/datasets) provides many starting
+ points. We use [Yahoo! Answers Topic Classification](https://huggingface.co/datasets/community-datasets/yahoo_answers_topics),
+ which has ten topics such as Society & Culture, Science & Mathematics, and
+ Health. We use `question_title` as the input and `topic` as the target.
+ The following cells download and preprocess the full dataset, which can
+ take time and disk space.
""")
return
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### Finding a dataset
-
- If you don't have the right dataset, you can always explore the [Datasets Hub](https://huggingface.co/datasets). The ["topic classification" category](https://huggingface.co/datasets?filter=task_ids:topic-classification) contains many datasets suitable for prototyping this model.
-
- We select ["Yahoo! Answers Topic Classification"](https://huggingface.co/datasets/community-datasets/yahoo_answers_topics) and visualize it with the [Datasets viewer](https://huggingface.co/datasets/community-datasets/yahoo_answers_topics/viewer/default/train).
-
- 
-
- Each topic number represents a unique subject:
-
- * 0:"Society & Culture"
- * 1:"Science & Mathematics"
- * 2:"Health"
- * etc…
-
- They correspond to the output we will try to predict from the model.
- """)
- return
+@app.cell
+def _():
+ model_checkpoint = "distilbert-base-uncased"
+ custom_run_name = "distilbert-yahoo-topics"
+ custom_max_steps = 30_000
+ custom_output_dir = Path("hf_outputs/topic_classification")
+ return (
+ custom_max_steps,
+ custom_output_dir,
+ custom_run_name,
+ model_checkpoint,
+ )
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ### Loading a dataset
+ ### Load and inspect the dataset
- Any dataset from the Datasets Hub can easily be loaded and is automatically downloaded if not present locally.
+ `load_dataset()` downloads data when it is not cached. The resulting
+ `DatasetDict` shows the splits, features, and row counts. Inspect one row,
+ then use the `ClassLabel` feature's `int2str()` and `str2int()` methods to
+ translate between topic IDs and names.
""")
return
-@app.cell(hide_code=True)
-def _(mo, wandb_settings):
- prepare_advanced_button = mo.ui.run_button(
- label="Download the dataset, tokenizer, and model",
- kind="success",
- )
- _ready_note = (
- "This downloads the Yahoo! Answers dataset and DistilBERT assets. "
- "It does not create a W&B run; training has a separate confirmation."
- )
- mo.vstack(
- [
- mo.md(
- f"The custom experiment will use W&B project "
- f"`{wandb_settings['project']}` under entity "
- f"`{wandb_settings['entity']}`.\n\n{_ready_note}"
- ),
- prepare_advanced_button,
- ]
- )
- return (prepare_advanced_button,)
-
-
-@app.cell(hide_code=True)
-def _(mo, prepare_advanced_button):
- mo.stop(
- not prepare_advanced_button.value,
- mo.callout(
- mo.md(
- "Click the button above when you are ready to download the "
- "advanced tutorial assets."
- ),
- kind="info",
- ),
- )
- advanced_setup_request = True
- return (advanced_setup_request,)
-
-
@app.cell
-def _(advanced_setup_request):
- assert advanced_setup_request
- from datasets import load_dataset
-
+def _():
dataset = load_dataset("community-datasets/yahoo_answers_topics")
- return (dataset,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Just printing our dataset object gives us a lot of information.
- """)
- return
-
-
-@app.cell
-def _(advanced_setup_request, dataset):
- assert advanced_setup_request
dataset
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- We can easily access any element.
- """)
- return
-
-
-@app.cell
-def _(advanced_setup_request, dataset):
- assert advanced_setup_request
- dataset['train'][0]
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- `str2int` and `int2str` help us go from class label to their integer mapping.
- """)
- return
-
-
-@app.cell
-def _(advanced_setup_request, dataset):
- assert advanced_setup_request
- dataset['train'].features['topic'].int2str(4)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- For our topic classification task, we use `question_title` as input and try to predict `topic`.
- """)
- return
+ return (dataset,)
@app.cell
-def _(advanced_setup_request, dataset):
- assert advanced_setup_request
- label_list = dataset['train'].unique('topic')
- label_list.sort()
- label_list
- return (label_list,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- This particular dataset is split between 10 different topics, that will be represented by 10 classes from our model output.
- """)
+def _(dataset):
+ dataset["train"][0]
return
@app.cell
-def _(advanced_setup_request, label_list):
- assert advanced_setup_request
- num_labels = len(label_list)
- num_labels
- return (num_labels,)
+def _(dataset):
+ topic_feature = dataset["train"].features["topic"]
+ label_names = topic_feature.names
+ num_labels = topic_feature.num_classes
+ print(f"{num_labels} topics: {label_names}")
+ topic_feature.int2str(4)
+ return label_names, num_labels
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- The "topic" class needs to be renamed to "labels" for the `Trainer` to find it.
+ Rename `topic` to `labels`, the target column expected by `Trainer`. Keep the
+ original dataset available for inspection and give each preparation stage
+ its own name so marimo can track the dependencies.
""")
return
@app.cell
-def _(advanced_setup_request, dataset):
- assert advanced_setup_request
- dataset_1 = dataset.rename_column('topic', 'labels')
- return (dataset_1,)
+def _(dataset):
+ labeled_dataset = dataset.rename_column("topic", "labels")
+ return (labeled_dataset,)
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ### Tokenizing the dataset
+ ### Tokenize the text
- In order to train a neural network, we need to convert our inputs to numbers:
- * the tokenizer divides a sequence of characters into tokens, ie sub-sequences (such as words, characters, sub-words…)
- * each unique token is mapped to a unique integer
+ A tokenizer splits text into tokens and maps them to integer IDs.
+ `AutoTokenizer` loads the tokenizer associated with the model checkpoint.
+ Inspect a single title and its tokenization before applying it in batches
+ with `Dataset.map()`, which caches the processed data locally.
- There are many [types of tokenizers](https://huggingface.co/transformers/tokenizer_summary.html). 🤗 Transformers can auto-select the right `Tokenizer` associated to a specific model.
+ Truncation limits inputs to the model's supported length. During training,
+ passing the tokenizer as `processing_class` lets `Trainer` dynamically pad
+ each batch. The tokenized rows include `input_ids` and `attention_mask`.
""")
return
@app.cell
-def _(advanced_setup_request):
- assert advanced_setup_request
- from transformers import AutoTokenizer
- tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
+def _(model_checkpoint):
+ tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
return (tokenizer,)
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- The tokenizer let us quickly preprocess our data.
- """)
- return
-
-
@app.cell
-def _(advanced_setup_request, dataset_1):
- assert advanced_setup_request
- sample_input = dataset_1['train'][0]['question_title']
- sample_input
- return (sample_input,)
-
-
-@app.cell
-def _(advanced_setup_request, sample_input, tokenizer):
- assert advanced_setup_request
+def _(labeled_dataset, tokenizer):
+ sample_input = labeled_dataset["train"][0]["question_title"]
+ print(sample_input)
tokenizer(sample_input)
return
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- The tokenizer can quickly process an entire dataset and cache the results locally to avoid any future tokenization of the same data.
-
- We leverage `dataset.map(fn)` function which can efficiently apply any function to a dataset. We also take advantage of batch processing which is supported by the tokenizer and makes the operation even faster.
- """)
- return
-
-
-@app.cell
-def _(advanced_setup_request, dataset_1, tokenizer):
- assert advanced_setup_request
- dataset_2 = dataset_1.map(lambda x: tokenizer(x['question_title'], truncation=True), batched=True)
- return (dataset_2,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- We truncate the data to the max length supported by the model. During training, we will pass the tokenizer to pad inputs to the longest sequence of the batch (the model requires same length inputs in a single batch).
-
- Our dataset now contains new keys: `input_ids` (tokens) and `attention_mask` (needed for certain models).
- """)
- return
-
-
@app.cell
-def _(advanced_setup_request, dataset_2):
- assert advanced_setup_request
- dataset_2['train'][0]
- return
+def _(labeled_dataset, tokenizer):
+ tokenized_dataset = labeled_dataset.map(
+ lambda batch: tokenizer(batch["question_title"], truncation=True),
+ batched=True,
+ )
+ tokenized_dataset["train"][0]
+ return (tokenized_dataset,)
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ### Loading a model
-
- Plenty of models are available and can be explored on the [Model Hub](https://huggingface.co/models).
-
- Once a model has been selected, it can be automatically loaded and adapted to one of its supported tasks.
- """)
- return
-
-
-@app.cell
-def _(advanced_setup_request, num_labels):
- assert advanced_setup_request
- from transformers import AutoModelForSequenceClassification
- model = AutoModelForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=num_labels)
- return (model,)
+ ### Load a model and inspect an initial prediction
+ Explore supported architectures on the [Model Hub](https://huggingface.co/models).
+ `AutoModelForSequenceClassification` loads a pretrained encoder and adds a
+ classification head for our ten topics. The new head has not learned the
+ topic labels yet, so its initial predictions are not meaningful topic classifications.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- In this case, we are loading a pre-trained network to which a custom head has been added for sequence classification and presents 10 classes corresponding to the possible topics of this dataset.
-
- Let's make a function to return the topic prediction from a sample question.
+ `get_topic()` tokenizes a sentence, moves its tensors to the model's device,
+ and selects the largest output logit. Evaluation mode disables dropout.
+ `torch.no_grad()` avoids building a gradient graph for inference.
""")
return
@app.cell
-def _(
- advanced_setup_request,
- dataset_2,
- model,
- tokenizer,
- torch,
- torch_device,
-):
- assert advanced_setup_request
- def get_topic(sentence, tokenize=tokenizer, model=model):
- # tokenize the input
- inputs = tokenize(sentence, return_tensors='pt')
- # Ensure the model and inputs are on the selected device.
- inputs = {name: tensor.to(torch_device) for name, tensor in inputs.items()}
- model = model.to(torch_device)
- # get prediction - 10 classes "probabilities" (not really true because they still need to be normalized)
- with torch.no_grad():
- predictions = model(**inputs).logits.cpu().numpy()
- # get the top prediction class and convert it to its associated label
- top_prediction = predictions.argmax().item()
- return dataset_2['train'].features['labels'].int2str(top_prediction)
-
- return (get_topic,)
+def _(label_names, model_checkpoint, num_labels):
+ set_seed(42)
+ untrained_model = AutoModelForSequenceClassification.from_pretrained(
+ model_checkpoint,
+ num_labels=num_labels,
+ id2label=dict(enumerate(label_names)),
+ label2id={label: index for index, label in enumerate(label_names)},
+ )
+ return (untrained_model,)
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Let's test a prediction on a sample sentence.
- """)
- return
+@app.function
+def get_topic(sentence, model, tokenizer):
+ model.eval()
+ model_device = next(model.parameters()).device
+ inputs = tokenizer(sentence, truncation=True, return_tensors="pt")
+ inputs = {name: tensor.to(model_device) for name, tensor in inputs.items()}
+ with torch.no_grad():
+ predicted_id = model(**inputs).logits.argmax(dim=-1).item()
+ return model.config.id2label[predicted_id]
@app.cell
-def _(advanced_setup_request, get_topic):
- assert advanced_setup_request
- get_topic('Why is cheese so much better with wine?')
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Obviously the model has not been trained yet so the results are still random.
- """)
- return
+def _(tokenizer, untrained_model):
+ sample_sentence = "Why is cheese so much better with wine?"
+ initial_topic_prediction = get_topic(sample_sentence, untrained_model, tokenizer)
+ initial_topic_prediction
+ return initial_topic_prediction, sample_sentence
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ### Training the model
-
- We now need to fine-tune the model based on our dataset.
-
- The `Trainer` class let us easily train a model and is very flexible.
+ ### Define an evaluation metric
- **Note:** set `report_to` to `wandb` in `TrainingArguments` to enable logging through W&B.
+ `Trainer` passes model logits and reference labels to `compute_metrics()`.
+ Convert logits to the most likely class with `argmax`, then compute accuracy
+ with NumPy. You can use metrics from the [Evaluate library](https://huggingface.co/docs/evaluate)
+ here for tasks that need other scores.
""")
return
-@app.cell
-def _():
- from transformers import Trainer, TrainingArguments
-
- ADVANCED_MAX_STEPS = 30000
- ADVANCED_RUN_NAME = "custom_training"
- return ADVANCED_MAX_STEPS, ADVANCED_RUN_NAME, Trainer, TrainingArguments
-
-
-@app.cell
-def _(TrainingArguments, advanced_run_request):
- args = TrainingArguments(
- report_to = 'wandb', # enable logging to W&B
- output_dir = advanced_run_request['output_dir'], # temporary output directory
- eval_strategy = 'steps', # check evaluation metrics at intervals
- learning_rate = 5e-5, # we can customize learning rate
- max_steps = advanced_run_request['max_steps'],
- logging_steps = 100, # we will log every 100 steps
- eval_steps = 5000, # we will perform evaluation every 5,000 steps
- save_steps = 10000,
- load_best_model_at_end = True,
- metric_for_best_model = 'accuracy',
- run_name = advanced_run_request['run_name'] # name of the W&B run
- )
- return (args,)
+@app.function
+def compute_metrics(eval_pred):
+ logits, labels = eval_pred
+ predictions = np.argmax(logits, axis=-1)
+ return {"accuracy": float(np.mean(predictions == labels))}
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- For more customization, refer to [`TrainingArguments` documentation](https://huggingface.co/transformers/main_classes/trainer.html#transformers.TrainingArguments).
- """)
- return
-
+ ### Configure training and Weights & Biases logging
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- We can optionally define metrics to calculate in addition to the loss through the `compute_metrics` function.
+ The next cell configures `Trainer` for 30,000 steps. `report_to="wandb"`
+ installs the Transformers Weights & Biases callback, which logs the model configuration,
+ training loss, evaluation accuracy, and other metrics automatically.
- Several [metrics](https://huggingface.co/docs/evaluate) are readily available from the Evaluate library to monitor model performance.
+ Evaluate and save every 5,000 steps so each evaluated model has a checkpoint.
+ At the end, reload the checkpoint with the best accuracy and log it as a
+ model Artifact.
+ To try a shorter run, reduce `custom_max_steps` and the evaluation and save
+ intervals together. Local checkpoints remain in `custom_output_dir`.
""")
return
@app.cell
-def _(advanced_setup_request):
- assert advanced_setup_request
- import evaluate
- import numpy as np
-
- accuracy_metric = evaluate.load("accuracy")
-
- def compute_metrics(eval_pred):
- predictions, labels = eval_pred
- predictions = np.argmax(predictions, axis=1)
- # metrics from the Evaluate library have a `compute` method
- return accuracy_metric.compute(predictions=predictions, references=labels)
-
- return (compute_metrics,)
-
-
-@app.cell(hide_code=True)
-def _(ADVANCED_MAX_STEPS, advanced_setup_request, mo, wandb_settings):
- assert advanced_setup_request
- advanced_training_button = mo.ui.run_button(
- label="Evaluate the baseline and train with W&B",
- kind="success",
- )
- _runtime_note = (
- "The button below starts baseline evaluation and the full "
- f"{ADVANCED_MAX_STEPS:,}-step fine-tuning run."
- )
- mo.vstack(
- [
- mo.callout(
- mo.md(
- f"{_runtime_note}\n\nThis will create a run in "
- f"`{wandb_settings['entity']}/{wandb_settings['project']}`."
- ),
- kind="warn",
- title="Long-running example",
- ),
- advanced_training_button,
- ]
- )
- return (advanced_training_button,)
-
-
-@app.cell(hide_code=True)
-def _(
- ADVANCED_MAX_STEPS,
- ADVANCED_RUN_NAME,
- WANDB_LOG_MODEL,
- advanced_training_button,
- mo,
- tempfile,
- uuid,
- wandb,
- wandb_settings,
-):
- mo.stop(
- not advanced_training_button.value,
- mo.callout(
- mo.md(
- "Click the button above when you are ready to evaluate the "
- "baseline, create a W&B run, and start fine-tuning."
- ),
- kind="info",
- ),
- )
- mo.stop(
- wandb.run is not None,
- mo.callout(
- mo.md(
- "Another W&B run is active in this Python process. Finish it "
- "before starting the custom training example."
- ),
- kind="warn",
- ),
+def _(custom_max_steps, custom_output_dir, custom_run_name):
+ training_args = TrainingArguments(
+ report_to="wandb",
+ output_dir=str(custom_output_dir),
+ eval_strategy="steps",
+ learning_rate=5e-5,
+ max_steps=custom_max_steps,
+ logging_steps=100,
+ eval_steps=5_000,
+ save_steps=5_000,
+ save_total_limit=2,
+ load_best_model_at_end=True,
+ metric_for_best_model="accuracy",
+ run_name=custom_run_name,
)
-
- _advanced_run_id = uuid.uuid4().hex
- advanced_run_request = {
- "max_steps": ADVANCED_MAX_STEPS,
- "output_dir": tempfile.mkdtemp(
- prefix=f"wandb-hf-{_advanced_run_id[:8]}-"
- ),
- "run_id": _advanced_run_id,
- "run_name": ADVANCED_RUN_NAME,
- }
- advanced_run_env = {
- "WANDB_ENTITY": wandb_settings["entity"],
- "WANDB_PROJECT": wandb_settings["project"],
- "WANDB_RUN_ID": _advanced_run_id,
- "WANDB_NAME": ADVANCED_RUN_NAME,
- "WANDB_LOG_MODEL": WANDB_LOG_MODEL,
- }
- advanced_run_url = wandb.Settings(
- entity=wandb_settings["entity"],
- project=wandb_settings["project"],
- run_id=_advanced_run_id,
- ).run_url
- return advanced_run_env, advanced_run_request, advanced_run_url
+ return (training_args,)
@app.cell(hide_code=True)
-def _(advanced_run_url, mo):
- mo.callout(
- mo.md(
- f"[Open this custom training run in W&B]({advanced_run_url}) "
- "to watch evaluation and training metrics appear live."
- ),
- kind="info",
- )
- advanced_run_link_ready = True
- return (advanced_run_link_ready,)
+def _(custom_max_steps):
+ mo.md(f"""
+ ### Evaluate the baseline and train
+ The next action runs baseline evaluation on Yahoo! Answers and starts
+ **{custom_max_steps:,} training steps**. It creates a Weights & Biases run and uploads
+ metrics and the trained model. This can take substantial time and compute.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- The `Trainer` handles all the training & evaluation logic.
+ We load a fresh copy of the pretrained checkpoint for each run. The explicit
+ `with wandb.init(...) as training_run` block connects the callback to that
+ run and finishes it when training exits, including after an error.
""")
return
@app.cell
def _(
- Trainer,
- advanced_run_link_ready,
- args,
- compute_metrics,
- dataset_2,
- model,
+ custom_run_name,
+ entity,
+ initial_topic_prediction,
+ label_names,
+ model_checkpoint,
+ num_labels,
+ project,
+ tokenized_dataset,
tokenizer,
+ training_args,
):
- assert advanced_run_link_ready
- trainer = Trainer(
- model=model, # model to be trained
- args=args, # training args
- train_dataset=dataset_2['train'],
- eval_dataset=dataset_2['test'],
- processing_class=tokenizer, # for padding batched data
- compute_metrics=compute_metrics, # for custom metrics
- )
- return (trainer,)
-
+ set_seed(training_args.seed)
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- We first call `evaluate()` to verify that the untrained model has about 10% accuracy (random predictions over 10 classes). We then start fine-tuning by calling `train()`.
-
- The cell scopes its W&B environment to this run and closes the run in a `finally` block, even if evaluation or training fails.
- """)
- return
-
-
-@app.cell
-def _(advanced_run_env, advanced_run_link_ready, mo, os, trainer, wandb):
- assert advanced_run_link_ready
- if wandb.run is not None:
- raise RuntimeError(
- "Another W&B run is active. Finish it before starting this example."
+ training_model = AutoModelForSequenceClassification.from_pretrained(
+ model_checkpoint,
+ num_labels=num_labels,
+ id2label=dict(enumerate(label_names)),
+ label2id={label: index for index, label in enumerate(label_names)},
+ )
+ with wandb.init(project=project, entity=entity, name=custom_run_name) as training_run:
+ print(f"Initial sample prediction: {initial_topic_prediction}")
+ trainer = Trainer(
+ model=training_model,
+ args=training_args,
+ train_dataset=tokenized_dataset["train"],
+ eval_dataset=tokenized_dataset["test"],
+ processing_class=tokenizer,
+ compute_metrics=compute_metrics,
)
-
- _previous_wandb_env = {
- _key: os.environ.get(_key) for _key in advanced_run_env
- }
- try:
- os.environ.update(advanced_run_env)
initial_metrics = trainer.evaluate()
- mo.output.append(initial_metrics)
+ print(f"Baseline metrics: {initial_metrics}")
training_result = trainer.train()
- finally:
- try:
- if wandb.run is not None:
- wandb.finish()
- finally:
- for _key, _previous_value in _previous_wandb_env.items():
- if _previous_value is None:
- os.environ.pop(_key, None)
- else:
- os.environ[_key] = _previous_value
-
- advanced_training_complete = True
- (initial_metrics, training_result)
- return
+ trained_model = trainer.model
+ training_run_url = training_run.url
+ return (trained_model,)
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- We can monitor losses, metrics, gradients and parameters as the model trains.
+ ### Predict with the trained model
- 
-
- When training is complete, our model is logged and versioned along with its performance as metadata.
-
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- We can now use the trained model for better predictions.
+ The next cell consumes `trained_model`, which is returned only after
+ training and Weights & Biases cleanup finish. It classifies the same sentence as the
+ baseline so you can compare the predictions.
""")
return
@app.cell
-def _(get_topic):
- trained_topic_prediction = get_topic('Why is cheese so much better with wine?')
+def _(sample_sentence, tokenizer, trained_model):
+ trained_topic_prediction = get_topic(sample_sentence, trained_model, tokenizer)
trained_topic_prediction
- return (trained_topic_prediction,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- `wandb.finish()` closes a W&B run explicitly. The guarded training cell above calls it automatically before restoring the previous W&B environment.
- """)
return
@app.cell(hide_code=True)
-def _(advanced_run_url, mo, trained_topic_prediction):
- mo.callout(
- mo.md(
- "Custom training finished and the W&B run is closed. "
- f"The sample's new predicted topic is `{trained_topic_prediction}`. "
- f"[Inspect the completed run]({advanced_run_url})."
- ),
- kind="success",
- )
- return
-
+def _():
+ mo.md(f"""
+ ## Verify and next steps
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Once you're happy with a model, don't forget to [share it with the world](https://huggingface.co/docs/transformers/model_sharing) on the Model Hub!
- """)
- return
+ Inspect `train/loss`, `eval/loss`, and `eval/accuracy`, the recorded configuration,
+ and gradient/parameter histograms. The **Artifacts** tab contains the
+ trained model with its performance metadata.
+ Once you are happy with a model, you can
+ [share it on the Hugging Face Hub](https://huggingface.co/docs/transformers/model_sharing).
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
## Resources
- * [Hugging Face and W&B integration documentation](https://docs.wandb.ai/models/integrations/huggingface) contains tips for making the most of W&B
+ * [Hugging Face and Weights & Biases integration documentation](https://docs.wandb.ai/models/integrations/huggingface) contains tips for making the most of Weights & Biases
* [🤗 Transformers documentation](https://huggingface.co/docs/transformers/) is extremely thorough and full of examples
""")
return
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Questions about W&B
-
- If you have any questions about using W&B to track your model performance and predictions, visit the [W&B Community](https://community.wandb.ai/).
- """)
- return
-
-
if __name__ == "__main__":
app.run()
diff --git a/marimo/convert/optimize-pytorch-lightning-models-with-weights-biases/optimize_pytorch_lightning_models_with_weights_biases.py b/marimo/convert/optimize-pytorch-lightning-models-with-weights-biases/optimize_pytorch_lightning_models_with_weights_biases.py
index dc4d7e28..905e7c02 100644
--- a/marimo/convert/optimize-pytorch-lightning-models-with-weights-biases/optimize_pytorch_lightning_models_with_weights_biases.py
+++ b/marimo/convert/optimize-pytorch-lightning-models-with-weights-biases/optimize_pytorch_lightning_models_with_weights_biases.py
@@ -12,15 +12,19 @@
import marimo
-__generated_with = "0.24.2"
-app = marimo.App(width="medium", app_title="PyTorch Lightning with W&B")
+__generated_with = "0.25.1"
+app = marimo.App(
+ width="medium",
+ app_title="PyTorch Lightning with Weights & Biases",
+)
with app.setup:
+ from pathlib import Path
+
import marimo as mo
import lightning.pytorch as pl
import torch
import wandb
- from pathlib import Path
from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import Callback, ModelCheckpoint
from lightning.pytorch.loggers import WandbLogger
@@ -35,33 +39,38 @@
@app.cell(hide_code=True)
def _():
mo.md(r"""
- # Pytorch Lightning models with Weights & Biases
-
- [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/optimize-pytorch-lightning-models-with-weights-biases/optimize_pytorch_lightning_models_with_weights_biases.py/server)
+
- ```python
- from lightning.pytorch.loggers import WandbLogger
- from lightning.pytorch import Trainer
+
+
- wandb_logger = WandbLogger()
- trainer = Trainer(logger=wandb_logger)
- ```
+ # Track PyTorch Lightning training with Weights & Biases
- W&B integration with Pytorch Lightning can automatically:
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/optimize-pytorch-lightning-models-with-weights-biases/optimize_pytorch_lightning_models_with_weights_biases.py/server)
- * log your configuration parameters
- * log your losses and metrics
- * log your model checkpoints
- * keep track of your code
- * log your system metrics (GPU, CPU, memory, temperature, etc.)
+ Train an MNIST classifier and use `WandbLogger` to record its hyperparameters,
+ metrics, prediction examples, and checkpoints. Lightning organizes the training
+ loop, while the logger sends the values recorded with `self.log()` to Weights & Biases.
- ## Docs
+ This example also adds `watch()` for gradient and parameter histograms, a
+ custom callback for prediction Tables, and `ModelCheckpoint` for saved weights.
+ These require explicit configuration, which the following sections explain.
- You can find the PyTorch Lightning WandbLogger docs [here](https://lightning.ai/docs/pytorch/stable/extensions/generated/lightning.pytorch.loggers.WandbLogger.html) and the Weights & Biases docs [here](https://docs.wandb.ai/models/integrations/lightning)
+ The basic connection is to pass a `WandbLogger` to the Lightning `Trainer`.
+ The training cell below supplies an existing Weights & Biases run so its
+ team, project, name, and lifecycle are controlled by the notebook.
""")
return
@@ -69,73 +78,98 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Authentication
+ ## Authenticate with Weights & Biases
+
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- Connect to W&B so that experiments are associated with your account. Enter your [W&B API key](https://wandb.ai/authorize), or leave it blank to use `WANDB_API_KEY` from the marimo Secrets panel or credentials already available in this runtime. A new molab session does not inherit credentials from your computer.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- Set the team entity to choose where runs are stored. Find it in your W&B team's URL, `wandb.ai/
+
+
+ # Organizing Hyperparameter Sweeps in PyTorch with Weights & Biases
+
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/organizing-hyperparameter-sweeps-in-pytorch-with-w-b/organizing_hyperparameter_sweeps_in_pytorch_with_w_b.py/server)
""")
return
@@ -58,25 +71,18 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- Finding a machine learning model that meets your desired metric (such as model accuracy) is normally a redundant task that can take multiple iterations. To make matters worse, it might be unclear which hyperparameter combinations to use for a given training run.
+ Comparing hyperparameters requires repeated training runs. A sweep defines
+ which settings to vary and which metric to compare, then records the
+ configuration and results of each trial.
- Use W&B Sweeps to create an organized and efficient way to automatically search through combinations of hyperparameter values such as the learning rate, batch size, number of hidden layers, optimizer type and more to find values that optimize your model based on your desired metric.
+ This example samples optimizer, hidden-layer size, dropout, learning rate,
+ and batch size for a small MNIST classifier. It minimizes training loss to
+ demonstrate sweep mechanics. Use a validation metric when selecting settings
+ for a model you intend to deploy.
- In this tutorial you will create a hyperparameter search with W&B PyTorch integration. Follow along with a [video tutorial](https://wandb.me/sweeps-video)!
+ In this tutorial you will create a hyperparameter search with Weights & Biases PyTorch integration. Follow along with a [video tutorial](https://wandb.me/sweeps-video).
-
- ## Sweeps: An Overview
-
- Running a hyperparameter sweep with Weights & Biases is very easy. There are just 3 simple steps:
-
- 1. **Define the sweep:** we do this by creating a dictionary or a [YAML file](https://docs.wandb.ai/models/sweeps/define-sweep-configuration) that specifies the parameters to search through, the search strategy, the optimization metric, and so on.
-
- 2. **Initialize the sweep:** with one line of code we initialize the sweep and pass in the dictionary of sweep configurations:
- `sweep_id = wandb.sweep(sweep_config)`
-
- 3. **Run the sweep agent:** also accomplished with one line of code, we call `wandb.agent()` and pass the `sweep_id` to run, along with a function that defines your model architecture and trains it:
- `wandb.agent(sweep_id, function=train)`
""")
return
@@ -84,276 +90,104 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Authentication
+ ## Authenticate with Weights & Biases
+
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- The setup cell imports W&B, PyTorch, and torchvision from the packages
- declared at the top of this notebook.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- Enter your [W&B API key](https://wandb.ai/authorize), or leave it blank to
- use `WANDB_API_KEY` from the molab Secrets panel or credentials already
- stored in this runtime. A fresh molab session does not inherit credentials
- from your computer.
+ ## Set your team and project
- The entity is the team name in `wandb.ai/
+
- In this notebook, we'll show you how to use W&B Artifacts (🏺)
- to track your ML experiment pipelines (🧪).
- Our sophisticated mathematical models predict the following result:
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/pipeline-versioning-with-w-b-artifacts/pipeline_versioning_with_w_b_artifacts.py/server)
- $$
- {\Huge
- 🧪 + 🏺 = 😃}
- $$
- """)
- return
+ # Pipeline Versioning with Weights & Biases Artifacts
+ Track an MNIST pipeline from raw data to an evaluated model. This Weights & Biases
+ lesson is based on the PyTorch [Basic MNIST Example](https://github.com/pytorch/examples/tree/master/mnist/).
+ Watch the [original video tutorial](https://youtu.be/Hd94gatGMic).
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Follow along with the video tutorial
+ An **Artifact** is a versioned collection of files, such as a dataset or model.
+ Each run in this lesson records one pipeline step. `log_artifact()` records
+ an output, and `use_artifact()` records an input. Together these relationships form a directed
+ acyclic graph showing which data and model versions produced a result.
- Watch the tutorial below, or [open it on YouTube](https://youtu.be/Hd94gatGMic).
- """)
- return
+ We'll create five runs: load data, preprocess data, initialize a model, train,
+ and evaluate. The initialized model is saved separately so later experiments
+ can start from the same weights.
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.Html(r"""
-
-
+ > Running the pipeline downloads MNIST, trains a small CNN for five epochs, and
+ uploads datasets, checkpoints, metrics, and example images.
""")
return
@@ -90,1065 +78,531 @@ def _(mo):
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- ## What are Artifacts and Why Should I Care?
-
- An "artifact", like a Greek [amphora 🏺](https://en.wikipedia.org/wiki/Amphora),
- is a produced object -- the output of a process.
- In ML, the most important artifacts are _datasets_ and _models_.
-
- And, like the [Cross of Coronado](https://indianajones.fandom.com/wiki/Cross_of_Coronado), these important artifacts belong in a museum!
- That is, they should be cataloged and organized
- so that you, your team, and the ML community at large can learn from them.
- After all, those who don't track training are doomed to repeat it.
-
- Using our Artifacts API, you can log `Artifact`s as outputs of W&B `Run`s or use `Artifact`s as input to `Run`s, as in this diagram,
- where a training run takes in a dataset and produces a model.
-
- 
-
- Since one run can use another's output as an input, Artifacts and Runs together form a directed graph -- actually, a bipartite [DAG](https://en.wikipedia.org/wiki/Directed_acyclic_graph)! -- with nodes for `Artifact`s and `Run`s
- and arrows connecting `Run`s to the `Artifact`s they consume or produce.
- """)
- return
+ ## Authenticate with Weights & Biases
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## 0. Setup
- """)
- return
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
+ ## Set your team and project
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- The notebook metadata declares W&B, PyTorch, and torchvision. The visible
- import cell above is the complete shared runtime setup.
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `artifacts-example`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _(mo):
- mo.md(r"""
- ### Authentication
-
- Submit the form before running the pipeline. Leave the API key blank to
- use `WANDB_API_KEY` from molab's Secrets panel or credentials already
- configured in this runtime. Editing an unsubmitted field does not log in,
- download MNIST, or create a W&B run.
- """)
- return
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="artifacts-example", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(mo):
- _wandb_entity = mo.ui.text(
- value="",
- label="W&B entity -- a team you belong to (blank uses your default)",
- full_width=True,
- )
- _wandb_project = mo.ui.text(
- value="artifacts-example",
- label="W&B project",
- full_width=True,
- )
- _wandb_api_key = mo.ui.text(
- value="",
- kind="password",
- label="W&B API key (blank uses runtime credentials)",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}\n\n{project}")
- .batch(
- api_key=_wandb_api_key,
- entity=_wandb_entity,
- project=_wandb_project,
- )
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
-
+def _(mo, wandb, wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
-@app.cell(hide_code=True)
-def _(mo, wandb, wandb_login_form):
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Submit the authentication form to continue."),
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _login_values = wandb_login_form.value
- _api_key = _login_values["api_key"].strip()
+
try:
- _login_ok = wandb.login(
- key=_api_key or None,
- relogin=bool(_api_key),
- )
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
except (wandb.errors.Error, ValueError):
_login_ok = False
+
mo.stop(
not _login_ok,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. Check the API key or "
- "configure `WANDB_API_KEY` in molab's Secrets panel, then submit again."
- ),
- kind="danger",
- ),
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- wandb_session = {
- "entity": _login_values["entity"].strip() or None,
- "project": _login_values["project"].strip() or "artifacts-example",
- }
- mo.callout(
- mo.md(
- f"Connected. Pipeline runs will be written to "
- f"**{wandb_session['project']}** only after you click a step button."
- ),
- kind="success",
- )
- return (wandb_session,)
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## 1. Log a Dataset
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- First, let's define some Artifacts.
-
- This example is based off of this PyTorch
- ["Basic MNIST Example"](https://github.com/pytorch/examples/tree/master/mnist/),
- but could just as easily have been done in TensorFlow, in any other framework,
- or in pure Python.
-
- We start with the `Dataset`s:
- - a `train`ing set, for choosing the parameters,
- - a `validation` set, for choosing the hyperparameters,
- - a `test`ing set, for evaluating the final model
-
- The first cell below defines these three datasets.
- """)
- return
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
@app.cell
-def _(random, torch, torchvision):
- # Ensure deterministic behavior
- torch.backends.cudnn.deterministic = True
- random.seed(0)
+def _(torch):
torch.manual_seed(0)
- torch.cuda.manual_seed_all(0)
-
- # Device configuration
+ torch.backends.cudnn.deterministic = True
if torch.cuda.is_available():
device = torch.device("cuda:0")
elif torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
+ return (device,)
- # Data parameters
- num_classes = 10
- input_shape = (1, 28, 28)
- # drop slow mirror from list of MNIST mirrors
- torchvision.datasets.MNIST.mirrors = [mirror for mirror in torchvision.datasets.MNIST.mirrors
- if not mirror.startswith("http://yann.lecun.com")]
- return device, input_shape, num_classes
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log the raw dataset
+
+ First define the data-loading logic. MNIST provides 60,000 training images and
+ 10,000 test images. We reserve the last 10,000 training images for validation,
+ leaving the official test set for final evaluation. Each split contains image
+ and label tensors.
+ """)
+ return
@app.cell
def _(TensorDataset, torchvision):
- def load(data_root, train_size=50_000):
- """
- # Load the data
- """
-
- # the data, split between train and test sets
- train = torchvision.datasets.MNIST(data_root, train=True, download=True)
- test = torchvision.datasets.MNIST(data_root, train=False, download=True)
- (x_train, y_train), (x_test, y_test) = (
- (train.data, train.targets),
- (test.data, test.targets),
+ def load_mnist(data_root, train_size=50_000):
+ # Skip the slow original mirror, keeping the fallback mirrors.
+ torchvision.datasets.MNIST.mirrors = [
+ mirror for mirror in torchvision.datasets.MNIST.mirrors
+ if not mirror.startswith("http://yann.lecun.com")
+ ]
+ train_data = torchvision.datasets.MNIST(data_root, train=True, download=True)
+ test_data = torchvision.datasets.MNIST(data_root, train=False, download=True)
+ return (
+ TensorDataset(train_data.data[:train_size], train_data.targets[:train_size]),
+ TensorDataset(train_data.data[train_size:], train_data.targets[train_size:]),
+ TensorDataset(test_data.data, test_data.targets),
)
- # split off a validation set for hyperparameter tuning
- x_train, x_val = x_train[:train_size], x_train[train_size:]
- y_train, y_val = y_train[:train_size], y_train[train_size:]
-
- training_set = TensorDataset(x_train, y_train)
- validation_set = TensorDataset(x_val, y_val)
- test_set = TensorDataset(x_test, y_test)
-
- datasets = [training_set, validation_set, test_set]
-
- return datasets
-
- return (load,)
+ return (load_mnist,)
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- This sets up a pattern we'll see repeated in this example:
- the code to log the data as an Artifact is wrapped around the code for
- producing that data.
- In this case, the code for `load`ing the data is
- separated out from the code for `load_and_log`ging the data.
-
- This is good practice!
- """)
- return
+ Now create `loading_run` and an Artifact named `mnist-raw` with type `dataset`.
+ An Artifact name identifies a collection of versions. Its type, such as
+ `dataset` or `model`, describes the kind of files it contains. A run's
+ `job_type` labels its role in the pipeline. Keeping related runs in one project
+ makes their lineage easier to explore.
+ `new_file()` both creates a file inside the Artifact and attaches it. We save
+ one `.pt` file per split and record the source and split sizes as JSON metadata.
+ After `log_artifact()`, `wait()` ensures the upload has completed before the
+ next stage fetches it. We pass its fully qualified, versioned name downstream,
+ rather than the mutable `latest` alias.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- In order to log these datasets as Artifacts,
- we just need to
- 1. create a `Run` with `wandb.init`,
- 2. create an `Artifact` for the dataset, and
- 3. save and log the associated `file`s with `run.log_artifact`.
-
- Check out the example code below, then click **Log raw dataset** to run
- this pipeline step.
+ Each stage uses a temporary directory for downloaded data and Weights & Biases runtime
+ files. Each run finishes before its temporary directory is removed.
+ The separate Weights & Biases Artifact cache can persist across stages.
""")
return
@app.cell
-def _(load, tempfile, torch, wandb):
- def load_and_log(session):
- workspace_context = tempfile.TemporaryDirectory(
- prefix="wandb-artifacts-load-"
- )
- workspace = workspace_context.name
- # 🚀 start a run, with a type to label it and a project it can call home
+def _(entity, load_mnist, project, tempfile, torch, wandb):
+ with tempfile.TemporaryDirectory(prefix="wandb-mnist-load-") as loading_workspace:
with wandb.init(
- project=session["project"],
- entity=session["entity"],
- job_type="load-data",
- reinit="create_new",
- dir=workspace,
- ) as run:
- datasets = load(workspace) # separate code for loading the datasets
- names = ["training", "validation", "test"]
-
- # 🏺 create our Artifact
+ project=project, entity=entity, job_type="load-data",
+ name="mnist-load-data", dir=loading_workspace,
+ ) as loading_run:
+ raw_datasets = load_mnist(loading_workspace)
raw_data = wandb.Artifact(
"mnist-raw", type="dataset",
- description="Raw MNIST dataset, split into train/val/test",
+ description="Raw MNIST dataset, split into train/validation/test",
metadata={
"source": "torchvision.datasets.MNIST",
- "sizes": [len(dataset) for dataset in datasets],
+ "sizes": [len(dataset) for dataset in raw_datasets],
},
)
+ for _split, _dataset in zip(["training", "validation", "test"], raw_datasets):
+ with raw_data.new_file(_split + ".pt", mode="wb") as _file:
+ torch.save(_dataset.tensors, _file)
- for name, data in zip(names, datasets):
- # 🐣 Store a new file in the artifact, and write something into its contents.
- with raw_data.new_file(name + ".pt", mode="wb") as file:
- x, y = data.tensors
- torch.save((x, y), file)
-
- # ✍️ Save the artifact to W&B.
- logged_artifact = run.log_artifact(raw_data).wait()
- result = {
- "artifact_path": logged_artifact.qualified_name,
- "run_url": run.url,
- }
-
- workspace_context.cleanup()
- return result
-
- return (load_and_log,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- log_raw_data = mo.ui.run_button(label="Log raw dataset")
- log_raw_data
- return (log_raw_data,)
-
-
-@app.cell
-def _(load_and_log, log_raw_data, mo, wandb_session):
- mo.stop(
- not log_raw_data.value,
- mo.md("Click **Log raw dataset** to create the first W&B Artifact."),
- )
- raw_data_result = load_and_log(wandb_session)
- mo.callout(
- mo.md(
- "Raw MNIST Artifact logged. "
- f"[Open the W&B run]({raw_data_result['run_url']})."
- ),
- kind="success",
- )
- return (raw_data_result,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### `wandb.init`
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- When we make the `Run` that's going to produce the `Artifact`s,
- we need to state which `project` it belongs to.
-
- Depending on your workflow,
- a project might be as big as `car-that-drives-itself`
- or as small as `iterative-architecture-experiment-117`.
-
- > **Rule of 👍**: if you can, keep all of the `Run`s that share `Artifact`s
- inside a single project. This keeps things simple,
- but don't worry -- `Artifact`s are portable across projects!
-
- To help keep track of all the different kinds of jobs you might run,
- it's useful to provide a `job_type` when making `Runs`.
- This keeps the graph of your Artifacts nice and tidy.
-
- > **Rule of 👍**: the `job_type` should be descriptive and correspond to a single step of your pipeline. Here, we separate out `load`ing data from `preprocess`ing data.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### `wandb.Artifact`
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- To log something as an `Artifact`, we have to first make an `Artifact` object.
-
- Every `Artifact` has a `name` -- that's what the first argument sets.
-
- > **Rule of 👍**: the `name` should be descriptive, but easy to remember and type --
- we like to use names that are hyphen-separated and correspond to variable names in the code.
-
- It also has a `type`. Just like `job_type`s for `Run`s,
- this is used for organizing the graph of `Run`s and `Artifact`s.
-
- > **Rule of 👍**: the `type` should be simple:
- more like `dataset` or `model`
- than `mnist-data-YYYYMMDD`.
-
- You can also attach a `description` and some `metadata`, as a dictionary.
- The `metadata` just needs to be serializable to JSON.
-
- > **Rule of 👍**: the `metadata` should be as descriptive as possible.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### `artifact.new_file` and `run.log_artifact`
- """)
- return
+ uploaded_raw_data = loading_run.log_artifact(raw_data)
+ uploaded_raw_data.wait()
+ raw_artifact_path = uploaded_raw_data.qualified_name
+ print(f"Raw dataset: {raw_artifact_path}")
+ return (raw_artifact_path,)
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- Once we've made an `Artifact` object, we need to add files to it.
-
- You read that right: _files_ with an _s_.
- `Artifact`s are structured like directories,
- with files and sub-directories.
+ ## Consume the raw Artifact and log preprocessed data
- > **Rule of 👍**: whenever it makes sense to do so, split the contents
- of an `Artifact` up into multiple files. This will help if it comes time to scale!
-
- We use the `new_file` method
- to simultaneously write the file and attach it to the `Artifact`.
- Below, we'll use the `add_file` method,
- which separates those two steps.
-
- Once we've added all of our files, we need to `log_artifact` to [wandb.ai](https://wandb.ai).
-
- You'll notice some URLs appeared in the output,
- including one for the Run page.
- That's where you can view the results of the `Run`,
- including any `Artifact`s that got logged.
-
- We'll see some examples that make better use of the other components of the Run page below.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## 2. Use a Logged Dataset Artifact
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- `Artifact`s in W&B, unlike artifacts in museums,
- are designed to be _used_, not just stored.
-
- Let's see what that looks like.
-
- The cell below defines a pipeline step that takes in a raw dataset
- and uses it to produce a `preprocess`ed dataset:
- `normalize`d and shaped correctly.
-
- Notice again that we split out the meat of the code, `preprocess`,
- from the code that interfaces with `wandb`.
+ The next helpers handle tensor files and preprocessing. `read_dataset()` loads
+ only tensors using `weights_only=True`. `preprocess()` converts pixel values to
+ floating point in `[0, 1]` and adds a channel dimension, changing image shapes
+ from `(28, 28)` to `(1, 28, 28)` for the CNN.
""")
return
@app.cell
-def _(TensorDataset, torch):
- def preprocess(dataset, normalize=True, expand_dims=True):
- """
- ## Prepare the data
- """
- x, y = dataset.tensors
+def _(TensorDataset, os, torch):
+ def read_dataset(data_dir, split):
+ images, labels = torch.load(
+ os.path.join(data_dir, split + ".pt"), weights_only=True,
+ )
+ return TensorDataset(images, labels)
+ def preprocess(dataset, normalize=True, expand_dims=True):
+ images, labels = dataset.tensors
if normalize:
- # Scale images to the [0, 1] range
- x = x.type(torch.float32) / 255
-
+ images = images.to(torch.float32) / 255
if expand_dims:
- # Make sure images have shape (1, 28, 28)
- x = torch.unsqueeze(x, 1)
-
- return TensorDataset(x, y)
+ images = images.unsqueeze(1)
+ return TensorDataset(images, labels)
- return (preprocess,)
+ return preprocess, read_dataset
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- Now for the code that instruments this `preprocess` step with `wandb.Artifact` logging.
+ `preprocessing_run.use_artifact()` records the raw dataset as this run's input.
+ `download()` materializes its files under the supplied directory, reusing Weights & Biases's
+ content-addressed cache when possible. Download time depends on the cache and
+ network.
- Note that the example below both `use`s an `Artifact`,
- which is new,
- and `log`s it,
- which is the same as the last step.
- `Artifact`s are both the inputs and the outputs of `Run`s!
+ We save the transformed splits in `mnist-preprocess` and record the preprocessing
+ options in its metadata. The CNN below expects normalized images with an explicit channel dimension,
+ so keep both preprocessing options enabled.
- We use a new `job_type`, `preprocess-data`,
- to make it clear that this is a different kind of job from the previous one.
+ Artifact references look like `entity/project/mnist-raw:v0`. You can also request
+ `mnist-raw:latest` or a custom alias such as `best`, but aliases can move to a new
+ version. This pipeline uses the exact version produced by the preceding cell.
""")
return
@app.cell
-def _(TensorDataset, os, torch):
- def read(data_dir, split):
- filename = split + ".pt"
- x, y = torch.load(
- os.path.join(data_dir, filename),
- weights_only=True,
- )
-
- return TensorDataset(x, y)
-
- return (read,)
-
-
-@app.cell
-def _(os, preprocess, read, tempfile, torch, wandb):
- def preprocess_and_log(raw_artifact_path, session, steps):
- workspace_context = tempfile.TemporaryDirectory(
- prefix="wandb-artifacts-preprocess-"
- )
- workspace = workspace_context.name
+def _(
+ entity,
+ os,
+ preprocess,
+ project,
+ raw_artifact_path,
+ read_dataset,
+ tempfile,
+ torch,
+ wandb,
+):
+ preprocessing_steps = {"normalize": True, "expand_dims": True}
+ with tempfile.TemporaryDirectory(prefix="wandb-mnist-preprocess-") as preprocessing_workspace:
with wandb.init(
- project=session["project"],
- entity=session["entity"],
- job_type="preprocess-data",
- reinit="create_new",
- dir=workspace,
- ) as run:
+ project=project, entity=entity, job_type="preprocess-data",
+ name="mnist-preprocess-data", dir=preprocessing_workspace,
+ ) as preprocessing_run:
+ raw_input = preprocessing_run.use_artifact(raw_artifact_path, type="dataset")
+ raw_directory = raw_input.download(root=os.path.join(preprocessing_workspace, "raw"))
processed_data = wandb.Artifact(
- "mnist-preprocess",
- type="dataset",
- description="Preprocessed MNIST dataset",
- metadata=steps,
- )
-
- # ✔️ declare which artifact we'll be using
- raw_data_artifact = run.use_artifact(raw_artifact_path, type="dataset")
-
- # 📥 if need be, download the artifact
- raw_dataset = raw_data_artifact.download(
- root=os.path.join(workspace, "raw-data")
+ "mnist-preprocess", type="dataset",
+ description="Normalized MNIST images with a channel dimension",
+ metadata=preprocessing_steps,
)
+ for _split in ["training", "validation", "test"]:
+ _dataset = preprocess(read_dataset(raw_directory, _split), **preprocessing_steps)
+ with processed_data.new_file(_split + ".pt", mode="wb") as _file:
+ torch.save(_dataset.tensors, _file)
- for split in ["training", "validation", "test"]:
- raw_split = read(raw_dataset, split)
- processed_dataset = preprocess(raw_split, **steps)
-
- with processed_data.new_file(split + ".pt", mode="wb") as file:
- x, y = processed_dataset.tensors
- torch.save((x, y), file)
-
- logged_artifact = run.log_artifact(processed_data).wait()
- result = {
- "artifact_path": logged_artifact.qualified_name,
- "run_url": run.url,
- }
-
- workspace_context.cleanup()
- return result
-
- return (preprocess_and_log,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- One thing to notice here is that the `steps` of the preprocessing
- are saved with the `preprocessed_data` as `metadata`.
-
- If you're trying to make your experiments reproducible,
- capturing lots of metadata is a good idea!
-
- Also, even though our dataset is a "`large artifact`",
- the `download` step is done in much less than a second.
-
- Expand the markdown cell below for details.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- preprocess_data = mo.ui.run_button(label="Preprocess dataset")
- preprocess_data
- return (preprocess_data,)
-
-
-@app.cell
-def _(mo, preprocess_and_log, preprocess_data, raw_data_result, wandb_session):
- mo.stop(
- not preprocess_data.value,
- mo.md("Click **Preprocess dataset** after logging the raw dataset."),
- )
- steps = {"normalize": True,
- "expand_dims": True}
-
- preprocess_result = preprocess_and_log(
- raw_data_result["artifact_path"],
- wandb_session,
- steps,
- )
- mo.callout(
- mo.md(
- "Preprocessed MNIST Artifact logged. "
- f"[Open the W&B run]({preprocess_result['run_url']})."
- ),
- kind="success",
- )
- return (preprocess_result,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### `run.use_artifact`
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- These steps are simpler. The consumer just needs to know the `name` of the `Artifact`, plus a bit more.
-
- That "bit more" is the `alias` of the particular version of the `Artifact` you want.
-
- By default, the last version to be uploaded is tagged `latest`.
- Otherwise, you can pick older versions with `v0`/`v1`, etc.,
- or you can provide your own aliases, like `best` or `jit-script`.
- Just like [Docker Hub](https://hub.docker.com/) tags,
- aliases are separated from names with `:`,
- so the `Artifact` we want is `mnist-raw:latest`.
-
- > **Rule of 👍**: Keep aliases short and sweet.
- Use custom `alias`es like `latest` or `best` when you want an `Artifact`
- that satisfies some property
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### `artifact.download`
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Now, you may be worrying about the `download` call.
- If we download another copy, won't that double the burden on memory?
-
- Don't worry friend. Before we actually download anything,
- we check to see if the right version is available locally.
- This uses the same technology that underlies [torrenting](https://en.wikipedia.org/wiki/Torrent_file) and [version control with `git`](https://blog.thoughtram.io/git/2014/11/18/the-anatomy-of-a-git-commit.html): hashing.
-
- `artifact.download()` materializes a version under the `root` you provide
- and reuses W&B's content-addressed cache when possible. This notebook passes
- an isolated temporary root to every pipeline stage, then removes it after
- the run and Artifact upload finish. That keeps downloaded tensors and W&B
- runtime files out of the repository.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- In your own workflow, pass a persistent `root=` to `artifact.download()` if
- you want to inspect the materialized files after the stage completes.
- """)
- return
+ uploaded_processed_data = preprocessing_run.log_artifact(processed_data)
+ uploaded_processed_data.wait()
+ processed_artifact_path = uploaded_processed_data.qualified_name
+ print(f"Preprocessed dataset: {processed_artifact_path}")
+ return (processed_artifact_path,)
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- ### The Artifacts page on [wandb.ai](https://wandb.ai)
-
- Now that we've logged and used an `Artifact`,
- let's check out the Artifacts tab on the Run page.
-
- Navigate to the Run page URL from the `wandb` output
- and select the "Artifacts" tab from the left sidebar
- (it's the one with the database icon,
- which looks like three hockey pucks stacked on top of one another).
-
- Click a row in either the "Input Artifacts" table
- or in the "Output Artifacts" table,
- then check out the tabs ("Overview", "Metadata")
- to see everything logged about the `Artifact`.
-
- We particularly like the "Graph View".
- By default, it shows a graph
- with the `type`s of `Artifact`s
- and the `job_type`s of `Run` as the two types of nodes,
- with arrows to represent consumption and production.
- """)
- return
-
+ ## Log an initialized model
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## 3. Log a Model
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- That's enough to see how the API for `Artifact`s works,
- but let's follow this example through to the end of the pipeline
- so we can see how `Artifact`s can improve your ML workflow.
-
- This first cell here builds a DNN `model` in PyTorch -- a really simple ConvNet.
-
- We'll start by just initializing the `model`, not training it.
- That way, we can repeat the training while keeping everything else constant.
+ This small ConvNet has two convolution/pooling blocks, dropout, and a linear
+ classifier. Its constructor calculates the flattened feature size from the
+ input shape and layer settings. Saving the architecture settings alongside the
+ weights lets another run reconstruct the model.
""")
return
@app.cell
-def _(input_shape, num_classes):
+def _():
from math import floor
import torch.nn as nn
class ConvNet(nn.Module):
- def __init__(self, hidden_layer_sizes=[32, 64],
- kernel_sizes=[3],
- activation="ReLU",
- pool_sizes=[2],
- dropout=0.5,
- num_classes=num_classes,
- input_shape=input_shape):
-
- super(ConvNet, self).__init__()
-
+ def __init__(
+ self, hidden_layer_sizes=(32, 64), kernel_sizes=(3,),
+ activation="ReLU", pool_sizes=(2,), dropout=0.5,
+ num_classes=10, input_shape=(1, 28, 28),
+ ):
+ super().__init__()
self.layer1 = nn.Sequential(
- nn.Conv2d(in_channels=input_shape[0], out_channels=hidden_layer_sizes[0], kernel_size=kernel_sizes[0]),
- getattr(nn, activation)(),
- nn.MaxPool2d(kernel_size=pool_sizes[0])
+ nn.Conv2d(input_shape[0], hidden_layer_sizes[0], kernel_sizes[0]),
+ getattr(nn, activation)(),
+ nn.MaxPool2d(pool_sizes[0]),
)
self.layer2 = nn.Sequential(
- nn.Conv2d(in_channels=hidden_layer_sizes[0], out_channels=hidden_layer_sizes[-1], kernel_size=kernel_sizes[-1]),
- getattr(nn, activation)(),
- nn.MaxPool2d(kernel_size=pool_sizes[-1])
- )
- self.layer3 = nn.Sequential(
- nn.Flatten(),
- nn.Dropout(dropout)
+ nn.Conv2d(hidden_layer_sizes[0], hidden_layer_sizes[-1], kernel_sizes[-1]),
+ getattr(nn, activation)(),
+ nn.MaxPool2d(pool_sizes[-1]),
)
-
- fc_input_dims = floor((input_shape[1] - kernel_sizes[0] + 1) / pool_sizes[0]) # layer 1 output size
- fc_input_dims = floor((fc_input_dims - kernel_sizes[-1] + 1) / pool_sizes[-1]) # layer 2 output size
- fc_input_dims = fc_input_dims*fc_input_dims*hidden_layer_sizes[-1] # layer 3 output size
-
- self.fc = nn.Linear(fc_input_dims, num_classes)
+ self.layer3 = nn.Sequential(nn.Flatten(), nn.Dropout(dropout))
+ feature_width = floor((input_shape[1] - kernel_sizes[0] + 1) / pool_sizes[0])
+ feature_width = floor((feature_width - kernel_sizes[-1] + 1) / pool_sizes[-1])
+ self.fc = nn.Linear(feature_width * feature_width * hidden_layer_sizes[-1], num_classes)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
- x = self.fc(x)
- return x
+ return self.fc(x)
return (ConvNet,)
-@app.cell
-def _(ConvNet):
- model_preview = ConvNet()
- model_preview
- return
-
-
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- Here, we're using W&B to track the run,
- and so using the [`wandb.config`](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/configs-in-w-b/configs_in_w_b.py/server)
- object to store all of the hyperparameters.
+ Choose the model settings below. `wandb.init(config=...)` records them with
+ `initialization_run`. We also copy that configuration into the model Artifact’s
+ metadata for later reconstruction.
- The `dict`ionary version of that `config` object is a really useful piece of `metadata`, so make sure to include it!
+ Here we save the checkpoint with `torch.save()` and attach the existing
+ file to `convnet` with `add_file()`. In contrast, `new_file()` above writes directly into
+ the Artifact. The initialized model is independent of the dataset preprocessing
+ stage. Training will consume both the initialized model and preprocessed data.
""")
return
@app.cell
-def _(ConvNet, os, tempfile, torch, wandb):
- def build_model_and_log(config, session):
- workspace_context = tempfile.TemporaryDirectory(
- prefix="wandb-artifacts-initialize-"
- )
- workspace = workspace_context.name
- with wandb.init(
- project=session["project"],
- entity=session["entity"],
- job_type="initialize",
- config=config,
- reinit="create_new",
- dir=workspace,
- ) as run:
- run_config = run.config
-
- model = ConvNet(**run_config)
-
- model_artifact = wandb.Artifact(
- "convnet",
- type="model",
- description="Simple AlexNet style CNN",
- metadata=dict(run_config),
- )
-
- checkpoint_path = os.path.join(workspace, "initialized_model.pth")
- torch.save(model.state_dict(), checkpoint_path)
- # ➕ another way to add a file to an Artifact
- model_artifact.add_file(checkpoint_path, name="initialized_model.pth")
-
- logged_artifact = run.log_artifact(model_artifact).wait()
- result = {
- "artifact_path": logged_artifact.qualified_name,
- "run_url": run.url,
- }
-
- workspace_context.cleanup()
- return result
-
- return (build_model_and_log,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- initialize_model = mo.ui.run_button(label="Initialize model")
- initialize_model
- return (initialize_model,)
+def _():
+ model_config = {
+ "hidden_layer_sizes": [32, 64],
+ "kernel_sizes": [3],
+ "activation": "ReLU",
+ "pool_sizes": [2],
+ "dropout": 0.5,
+ "num_classes": 10,
+ }
+ return (model_config,)
@app.cell
-def _(
- build_model_and_log,
- initialize_model,
- mo,
- preprocess_result,
- wandb_session,
-):
- mo.stop(
- not initialize_model.value,
- mo.md("Click **Initialize model** after preprocessing the dataset."),
- )
- model_config = {"hidden_layer_sizes": [32, 64],
- "kernel_sizes": [3],
- "activation": "ReLU",
- "pool_sizes": [2],
- "dropout": 0.5,
- "num_classes": 10}
-
- _ = preprocess_result
- initialized_model_result = build_model_and_log(
- model_config,
- wandb_session,
- )
- mo.callout(
- mo.md(
- "Initialized model Artifact logged. "
- f"[Open the W&B run]({initialized_model_result['run_url']})."
- ),
- kind="success",
- )
- return (initialized_model_result,)
+def _(ConvNet, entity, model_config, os, project, tempfile, torch, wandb):
+ with tempfile.TemporaryDirectory(prefix="wandb-mnist-initialize-") as initialization_workspace:
+ with wandb.init(
+ project=project, entity=entity, job_type="initialize",
+ name="mnist-initialize-model", config=model_config, dir=initialization_workspace,
+ ) as initialization_run:
+ initial_model = ConvNet(**initialization_run.config)
+ initialized_model_artifact = wandb.Artifact(
+ "convnet", type="model", description="Initialized MNIST ConvNet",
+ metadata=dict(initialization_run.config),
+ )
+ initialized_checkpoint = os.path.join(initialization_workspace, "initialized_model.pth")
+ torch.save(initial_model.state_dict(), initialized_checkpoint)
+ initialized_model_artifact.add_file(initialized_checkpoint, name="initialized_model.pth")
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### `artifact.add_file`
- """)
- return
+ uploaded_initial_model = initialization_run.log_artifact(initialized_model_artifact)
+ uploaded_initial_model.wait()
+ initialized_artifact_path = uploaded_initial_model.qualified_name
+ print(f"Initialized model: {initialized_artifact_path}")
+ return (initialized_artifact_path,)
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- Instead of simultaneously writing a `new_file` and adding it to the `Artifact`,
- as in the dataset logging examples,
- we can also write files in one step
- (here, `torch.save`)
- and then `add` them to the `Artifact` in another.
+ ## Train using dataset and model Artifacts
- > **Rule of 👍**: use `new_file` when you can, to prevent duplication.
+ Before training, define an evaluation helper that computes mean cross-entropy
+ loss and accuracy across a loader without tracking gradients. Training will use
+ it on validation data after each epoch. Final evaluation will use the same
+ logic on the held-out test set.
""")
return
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## 4. Use a Logged Model Artifact
- """)
- return
+@app.cell
+def _(F, device, torch):
+ def evaluate_loader(model, loader):
+ model.eval()
+ total_loss = 0.0
+ correct = 0
+ with torch.no_grad():
+ for images, labels in loader:
+ images, labels = images.to(device), labels.to(device)
+ logits = model(images)
+ total_loss += F.cross_entropy(logits, labels, reduction="sum").item()
+ correct += (logits.argmax(dim=1) == labels).sum().item()
+ return total_loss / len(loader.dataset), 100.0 * correct / len(loader.dataset)
+
+ return (evaluate_loader,)
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- Just like we could call `use_artifact` on a `dataset`,
- we can call it on our `initialized_model`
- to use it in another `Run`.
+ The training cell declares **two Artifact inputs**: the preprocessed dataset and
+ the initialized model. It loads weights and reconstructs the architecture from
+ model metadata, then trains using the editable settings below.
- This time, let's `train` the `model`.
+ `training_run.log()` records training loss and validation loss/accuracy, with the
+ epoch and number of examples seen. Weights & Biases advances its default step on each log
+ call. `model.train()` at the start of each epoch restores dropout after validation
+ switches the model to evaluation mode.
- For more details, see the notebook on
- [instrumenting W&B with PyTorch](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/simple-pytorch-integration/simple_pytorch_integration.py/server).
+ Finally, save a new `trained-model` Artifact. Its metadata records the architecture
+ and training settings, while the run's input edges identify the exact dataset
+ and initialization used. The upload must finish before the evaluation run fetches it.
+ For more training details, see the
+ [PyTorch integration lesson](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/simple-pytorch-integration/simple_pytorch_integration.py/server).
""")
return
@app.cell
def _():
- import torch.nn.functional as F
-
- return (F,)
-
-
-@app.cell
-def _(F, device, evaluate_loader, torch):
- def train(model, train_loader, valid_loader, config, run):
- optimizer = getattr(torch.optim, config.optimizer)(model.parameters())
- example_ct = 0
- for epoch in range(config.epochs):
- # Evaluation switches the model to eval mode, so restore training mode
- # before every epoch (important for the Dropout layer).
- model.train()
- for batch_idx, (data, target) in enumerate(train_loader):
- data, target = data.to(device), target.to(device)
- optimizer.zero_grad()
- output = model(data)
- loss = F.cross_entropy(output, target)
- loss.backward()
- optimizer.step()
-
- example_ct += len(data)
-
- if batch_idx % config.batch_log_interval == 0:
- print(
- "Train Epoch: {} [{}/{} ({:.0%})]\tLoss: {:.6f}".format(
- epoch,
- batch_idx * len(data),
- len(train_loader.dataset),
- batch_idx / len(train_loader),
- loss.item(),
- )
- )
-
- train_log(run, loss, example_ct, epoch)
-
- # evaluate the model on the validation set at each epoch
- loss, accuracy = evaluate_loader(model, valid_loader)
- validation_log(run, loss, accuracy, example_ct, epoch)
-
- return (train,)
+ train_config = {
+ "batch_size": 128,
+ "epochs": 5,
+ "batch_log_interval": 25,
+ "optimizer": "Adam",
+ }
+ return (train_config,)
@app.cell
-def _(F, device, torch):
- def evaluate_loader(model, test_loader):
- model.eval()
- test_loss = 0
- correct = 0
- with torch.no_grad():
- for data, target in test_loader:
- data, target = data.to(device), target.to(device)
- output = model(data)
- test_loss += F.cross_entropy(output, target, reduction="sum")
- pred = output.argmax(dim=1, keepdim=True)
- correct += pred.eq(target.view_as(pred)).sum()
-
- test_loss /= len(test_loader.dataset)
- accuracy = 100.0 * correct / len(test_loader.dataset)
-
- return test_loss, accuracy
-
- return (evaluate_loader,)
-
-
-@app.function
-def train_log(run, loss, example_ct, epoch):
- loss = loss.detach().item()
-
- # where the magic happens
- run.log({"epoch": epoch, "train/loss": loss}, step=example_ct)
- print(f"Loss after " + str(example_ct).zfill(5) + f" examples: {loss:.3f}")
-
+def _(
+ ConvNet,
+ F,
+ device,
+ entity,
+ evaluate_loader,
+ initialized_artifact_path,
+ os,
+ processed_artifact_path,
+ project,
+ read_dataset,
+ tempfile,
+ torch,
+ train_config,
+ wandb,
+):
+ with tempfile.TemporaryDirectory(prefix="wandb-mnist-train-") as training_workspace:
+ with wandb.init(
+ project=project, entity=entity, job_type="train",
+ name="mnist-train-convnet", config=train_config, dir=training_workspace,
+ ) as training_run:
+ training_data_artifact = training_run.use_artifact(processed_artifact_path, type="dataset")
+ training_data_directory = training_data_artifact.download(
+ root=os.path.join(training_workspace, "data")
+ )
+ training_dataset = read_dataset(training_data_directory, "training")
+ validation_dataset = read_dataset(training_data_directory, "validation")
+ train_loader = torch.utils.data.DataLoader(
+ training_dataset, batch_size=training_run.config.batch_size,
+ )
+ validation_loader = torch.utils.data.DataLoader(
+ validation_dataset, batch_size=training_run.config.batch_size,
+ )
-@app.function
-def validation_log(run, loss, accuracy, example_ct, epoch):
- loss = float(loss)
- accuracy = float(accuracy)
+ initialization_input = training_run.use_artifact(initialized_artifact_path, type="model")
+ initialization_directory = initialization_input.download(
+ root=os.path.join(training_workspace, "initialized-model")
+ )
+ training_model_config = dict(initialization_input.metadata)
+ training_run.config.update(training_model_config)
+ training_model = ConvNet(**training_model_config).to(device)
+ training_model.load_state_dict(torch.load(
+ os.path.join(initialization_directory, "initialized_model.pth"),
+ map_location=device, weights_only=True,
+ ))
+
+ optimizer = getattr(torch.optim, training_run.config.optimizer)(training_model.parameters())
+ example_count = 0
+ for epoch in range(training_run.config.epochs):
+ training_model.train()
+ for batch_index, (batch_images, batch_labels) in enumerate(train_loader):
+ batch_images, batch_labels = batch_images.to(device), batch_labels.to(device)
+ optimizer.zero_grad()
+ training_loss = F.cross_entropy(training_model(batch_images), batch_labels)
+ training_loss.backward()
+ optimizer.step()
+ example_count += len(batch_images)
+ if batch_index % training_run.config.batch_log_interval == 0:
+ training_run.log({
+ "epoch": epoch, "examples": example_count,
+ "train/loss": training_loss.item(),
+ })
+
+ validation_loss, validation_accuracy = evaluate_loader(training_model, validation_loader)
+ training_run.log({
+ "epoch": epoch, "examples": example_count,
+ "validation/loss": validation_loss,
+ "validation/accuracy": validation_accuracy,
+ })
+ print(f"Epoch {epoch + 1}: validation accuracy {validation_accuracy:.2f}%")
- # where the magic happens
- run.log(
- {
- "epoch": epoch,
- "validation/loss": loss,
- "validation/accuracy": accuracy,
- },
- step=example_ct,
- )
- print(
- f"Loss/accuracy after "
- + str(example_ct).zfill(5)
- + f" examples: {loss:.3f}/{accuracy:.3f}"
- )
+ trained_model_artifact = wandb.Artifact(
+ "trained-model", type="model", description="Trained MNIST ConvNet",
+ metadata={
+ "model_config": training_model_config,
+ "training_config": train_config,
+ },
+ )
+ trained_checkpoint = os.path.join(training_workspace, "trained_model.pth")
+ torch.save(training_model.state_dict(), trained_checkpoint)
+ trained_model_artifact.add_file(trained_checkpoint, name="trained_model.pth")
+ uploaded_trained_model = training_run.log_artifact(trained_model_artifact)
+ uploaded_trained_model.wait()
+ trained_artifact_path = uploaded_trained_model.qualified_name
+ print(f"Trained model: {trained_artifact_path}")
+ return (trained_artifact_path,)
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- We'll run two separate `Artifact`-producing `Run`s this time.
+ ## Evaluate the trained model and inspect hard examples
- Once the first finishes `train`ing the `model`,
- the `second` will consume the `trained-model` `Artifact`
- by `evaluate`ing its performance on the `test_dataset`.
+ The final run consumes `trained-model` and the same preprocessed dataset version.
+ It evaluates only the test split, keeping test results separate from training and
+ hyperparameter selection.
- Also, we'll pull out the 32 examples on which the network gets the most confused --
- on which the `categorical_crossentropy` is highest.
-
- This is a good way to diagnose issues with your dataset and your model!
+ The next helper ranks individual test images by cross-entropy loss and returns
+ the 32 highest-loss examples (or all examples for a smaller test set). This helps
+ diagnose confident mistakes or ambiguous images.
""")
return
-@app.cell
-def _(evaluate_loader, get_hardest_k_examples):
- def evaluate(model, test_loader):
- """
- ## Evaluate the trained model
- """
-
- loss, accuracy = evaluate_loader(model, test_loader)
- highest_losses, hardest_examples, true_labels, predictions = (
- get_hardest_k_examples(model, test_loader.dataset)
- )
-
- return loss, accuracy, highest_losses, hardest_examples, true_labels, predictions
-
- return (evaluate,)
-
-
@app.cell
def _(F, device, torch):
def get_hardest_k_examples(model, testing_set, k=32):
model.eval()
-
loader = torch.utils.data.DataLoader(testing_set, 1, shuffle=False)
-
- # get the losses and predictions for each item in the dataset
losses = []
predictions = []
with torch.no_grad():
- for data, target in loader:
- data, target = data.to(device), target.to(device)
- output = model(data)
- loss = F.cross_entropy(output, target)
- pred = output.argmax(dim=1, keepdim=True)
-
- losses.append(loss.detach().cpu())
- predictions.append(pred.detach().cpu().squeeze(1))
-
+ for images, labels in loader:
+ images, labels = images.to(device), labels.to(device)
+ logits = model(images)
+ losses.append(F.cross_entropy(logits, labels).detach().cpu())
+ predictions.append(logits.argmax(dim=1).detach().cpu())
losses = torch.stack(losses)
predictions = torch.cat(predictions)
- argsort_loss = torch.argsort(losses)
- hardest_indices = argsort_loss[-k:]
-
- highest_k_losses = losses[hardest_indices]
- hardest_k_examples = testing_set.tensors[0][hardest_indices]
- true_labels = testing_set.tensors[1][hardest_indices]
- predicted_labels = predictions[hardest_indices]
-
- return highest_k_losses, hardest_k_examples, true_labels, predicted_labels
+ hardest_indices = torch.argsort(losses)[-k:]
+ return (
+ losses[hardest_indices],
+ testing_set.tensors[0][hardest_indices],
+ testing_set.tensors[1][hardest_indices],
+ predictions[hardest_indices],
+ )
return (get_hardest_k_examples,)
@@ -1156,263 +610,98 @@ def get_hardest_k_examples(model, testing_set, k=32):
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- These logging functions don't add any new `Artifact` features,
- so we won't comment on them:
- we're just `use`ing, `download`ing,
- and `log`ging `Artifact`s.
+ `evaluation_run.summary.update()` stores the final test loss and accuracy.
+ `evaluation_run.log()` records high-loss images with their predicted label, true
+ label, and loss in the captions. These Media panels complement aggregate metrics
+ with concrete examples of the model's behavior.
+
+ The trained Artifact stores architecture and training settings separately in
+ its metadata. We use `model_config` to reconstruct the CNN. The training
+ settings remain available for inspection.
""")
return
@app.cell
-def _(ConvNet, device, os, read, tempfile, torch, train, wandb):
- def train_and_log(
- config,
- initialized_model_path,
- processed_data_path,
- session,
- ):
- workspace_context = tempfile.TemporaryDirectory(
- prefix="wandb-artifacts-train-"
- )
- workspace = workspace_context.name
- with wandb.init(
- project=session["project"],
- entity=session["entity"],
- job_type="train",
- config=config,
- reinit="create_new",
- dir=workspace,
- ) as run:
- run_config = run.config
-
- data = run.use_artifact(processed_data_path, type="dataset")
- data_dir = data.download(root=os.path.join(workspace, "data"))
-
- training_dataset = read(data_dir, "training")
- validation_dataset = read(data_dir, "validation")
-
- train_loader = torch.utils.data.DataLoader(
- training_dataset,
- batch_size=run_config.batch_size,
- )
- validation_loader = torch.utils.data.DataLoader(
- validation_dataset,
- batch_size=run_config.batch_size,
- )
-
- model_artifact = run.use_artifact(
- initialized_model_path,
- type="model",
- )
- model_dir = model_artifact.download(
- root=os.path.join(workspace, "initialized-model")
- )
- model_path = os.path.join(model_dir, "initialized_model.pth")
- model_config = model_artifact.metadata
- run_config.update(model_config)
-
- model = ConvNet(**model_config)
- model.load_state_dict(
- torch.load(model_path, map_location=device, weights_only=True)
- )
- model = model.to(device)
-
- train(model, train_loader, validation_loader, run_config, run)
-
- trained_model_artifact = wandb.Artifact(
- "trained-model",
- type="model",
- description="Trained NN model",
- metadata=dict(model_config),
- )
-
- trained_model_path = os.path.join(workspace, "trained_model.pth")
- torch.save(model.state_dict(), trained_model_path)
- trained_model_artifact.add_file(
- trained_model_path,
- name="trained_model.pth",
- )
- logged_model = run.log_artifact(trained_model_artifact).wait()
- result = {
- "artifact_path": logged_model.qualified_name,
- "run_url": run.url,
- }
-
- workspace_context.cleanup()
- return result
-
- return (train_and_log,)
-
-
-@app.cell
-def _(ConvNet, device, evaluate, os, read, tempfile, torch, wandb):
- def evaluate_and_log(
- processed_data_path,
- session,
- trained_model_path,
- config=None,
- ):
- workspace_context = tempfile.TemporaryDirectory(
- prefix="wandb-artifacts-evaluate-"
- )
- workspace = workspace_context.name
+def _(
+ ConvNet,
+ device,
+ entity,
+ evaluate_loader,
+ get_hardest_k_examples,
+ os,
+ processed_artifact_path,
+ project,
+ read_dataset,
+ tempfile,
+ torch,
+ trained_artifact_path,
+ wandb,
+):
+ with tempfile.TemporaryDirectory(prefix="wandb-mnist-evaluate-") as evaluation_workspace:
with wandb.init(
- project=session["project"],
- entity=session["entity"],
- job_type="report",
- config=config,
- reinit="create_new",
- dir=workspace,
- ) as run:
- data = run.use_artifact(processed_data_path, type="dataset")
- data_dir = data.download(root=os.path.join(workspace, "data"))
- testing_set = read(data_dir, "test")
-
- test_loader = torch.utils.data.DataLoader(
- testing_set,
- batch_size=128,
- shuffle=False,
+ project=project, entity=entity, job_type="report",
+ name="mnist-evaluate-convnet", dir=evaluation_workspace,
+ ) as evaluation_run:
+ evaluation_data_artifact = evaluation_run.use_artifact(processed_artifact_path, type="dataset")
+ evaluation_directory = evaluation_data_artifact.download(
+ root=os.path.join(evaluation_workspace, "data")
)
+ testing_set = read_dataset(evaluation_directory, "test")
+ test_loader = torch.utils.data.DataLoader(testing_set, batch_size=128, shuffle=False)
- model_artifact = run.use_artifact(trained_model_path, type="model")
- model_dir = model_artifact.download(
- root=os.path.join(workspace, "trained-model")
+ evaluation_model_artifact = evaluation_run.use_artifact(trained_artifact_path, type="model")
+ evaluation_model_directory = evaluation_model_artifact.download(
+ root=os.path.join(evaluation_workspace, "trained-model")
)
- model_path = os.path.join(model_dir, "trained_model.pth")
- model_config = model_artifact.metadata
-
- model = ConvNet(**model_config)
- model.load_state_dict(
- torch.load(model_path, map_location=device, weights_only=True)
- )
- model = model.to(device)
-
- loss, accuracy, highest_losses, hardest_examples, true_labels, preds = (
- evaluate(model, test_loader)
- )
-
- run.summary.update(
- {"loss": float(loss), "accuracy": float(accuracy)}
- )
-
- run.log(
- {
- "high-loss-examples": [
- wandb.Image(
- hard_example.detach()
- .cpu()
- .clamp(0, 1)
- .mul(255)
- .round()
- .to(torch.uint8),
- caption=str(int(pred)) + "," + str(int(label)),
- )
- for hard_example, pred, label in zip(
- hardest_examples,
- preds,
- true_labels,
- )
- ]
- }
+ evaluation_model_config = evaluation_model_artifact.metadata["model_config"]
+ evaluation_model = ConvNet(**evaluation_model_config).to(device)
+ evaluation_model.load_state_dict(torch.load(
+ os.path.join(evaluation_model_directory, "trained_model.pth"),
+ map_location=device, weights_only=True,
+ ))
+ test_loss, test_accuracy = evaluate_loader(evaluation_model, test_loader)
+ highest_losses, hardest_images, true_labels, predicted_labels = get_hardest_k_examples(
+ evaluation_model, testing_set
)
- result = {
- "accuracy": float(accuracy),
- "loss": float(loss),
- "run_url": run.url,
- }
-
- workspace_context.cleanup()
- return result
-
- return (evaluate_and_log,)
-
-
-@app.cell(hide_code=True)
-def _(mo):
- train_and_evaluate = mo.ui.run_button(label="Train and evaluate")
- train_and_evaluate
- return (train_and_evaluate,)
-
-
-@app.cell
-def _(
- evaluate_and_log,
- initialized_model_result,
- mo,
- preprocess_result,
- train_and_evaluate,
- train_and_log,
- wandb_session,
-):
- mo.stop(
- not train_and_evaluate.value,
- mo.md("Click **Train and evaluate** after initializing the model."),
- )
- train_config = {"batch_size": 128,
- "epochs": 5,
- "batch_log_interval": 25,
- "optimizer": "Adam"}
-
- trained_model_result = train_and_log(
- train_config,
- initialized_model_result["artifact_path"],
- preprocess_result["artifact_path"],
- wandb_session,
- )
- evaluation_result = evaluate_and_log(
- preprocess_result["artifact_path"],
- wandb_session,
- trained_model_result["artifact_path"],
- )
- mo.callout(
- mo.md(
- "Training and evaluation complete. "
- f"Validation artifacts: [training run]({trained_model_result['run_url']}); "
- f"[evaluation run]({evaluation_result['run_url']}). "
- f"Test accuracy: **{evaluation_result['accuracy']:.2f}%**."
- ),
- kind="success",
- )
+ evaluation_run.summary.update({"loss": test_loss, "accuracy": test_accuracy})
+ evaluation_run.log({
+ "high-loss-examples": [
+ wandb.Image(
+ image.cpu().clamp(0, 1).mul(255).round().to(torch.uint8),
+ caption=f"predicted={int(prediction)}, true={int(label)}, loss={float(loss):.3f}",
+ )
+ for image, prediction, label, loss in zip(
+ hardest_images, predicted_labels, true_labels, highest_losses,
+ )
+ ]
+ })
+ print(f"Test loss: {test_loss:.4f}; test accuracy: {test_accuracy:.2f}%")
+ print(f"Evaluation run: {evaluation_run.url}")
return
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- You may have noticed a button labeled "Explode". Don't click that, as it will set off a small bomb underneath your humble author's desk in the W&B HQ!
-
- Just kidding. It "explodes" the graph in a much gentler way:
- `Artifact`s and `Run`s become separated at the level of a single instance,
- rather than a `type`:
- the nodes are not `dataset` and `load-data`, but `dataset:mnist-raw:v1` and `load-data:sunny-smoke-1`, and so on.
+ ## Inspect the pipeline in Weights & Biases
- This provides total insight into your pipeline,
- with logged metrics, metadata, and more
- all at your fingertips --
- you're only limited by what you choose to log with us.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Verify and next steps
+ Open the run links printed by Weights & Biases and inspect:
- In the W&B project, open the runs created by this notebook and inspect:
+ - **Artifacts:** `mnist-raw`, `mnist-preprocess`, `convnet`, and `trained-model`.
+ - **Lineage graph:** follow the raw dataset into preprocessing, then follow
+ the processed dataset and initialized model into training. Confirm that
+ evaluation consumes the trained model and the processed dataset.
+ - **Metadata:** dataset sizes, preprocessing settings, and model configuration.
+ - **Training charts:** `train/loss`, `validation/loss`, and `validation/accuracy`.
+ - **Evaluation:** final `loss` and `accuracy`, plus the `high-loss-examples` panel.
- - the **Artifacts** tab for `mnist-raw`, `mnist-preprocess`, `convnet`, and
- `trained-model`;
- - the Artifact graph to confirm each run consumes the previous step's
- output and produces the next Artifact;
- - the training run charts for `train/loss`, `validation/loss`, and
- `validation/accuracy`;
- - the report run summary for `loss` and `accuracy`, plus the
- `high-loss-examples` media panel.
+ The graph can group nodes by Artifact `type` and run `job_type`, or show individual
+ runs and Artifact versions. This lets you trace a result back to its exact inputs.
- Re-run the model initialization or training steps with different
- hyperparameters to create new Artifact versions and compare them in W&B.
+ Edit the model or training settings and rerun the affected stages to create and
+ compare versions. Weights & Biases deduplicates unchanged file contents. For a controlled
+ training comparison, keep the initialized model and dataset versions fixed.
""")
return
diff --git a/marimo/convert/plot-roc-curves-with-w-b/plot_roc_curves_with_w_b.py b/marimo/convert/plot-roc-curves-with-w-b/plot_roc_curves_with_w_b.py
index a7438fdd..4e84edf5 100644
--- a/marimo/convert/plot-roc-curves-with-w-b/plot_roc_curves_with_w_b.py
+++ b/marimo/convert/plot-roc-curves-with-w-b/plot_roc_curves_with_w_b.py
@@ -1,8 +1,7 @@
# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
-# "fsspec[http]",
-# "marimo>=0.24.0",
+# "marimo>=0.24.2",
# "numpy>=1.26",
# "pandas>=2.2",
# "pillow==12.3.0",
@@ -11,26 +10,32 @@
# "tensorflow==2.21.0",
# "wandb==0.30.0",
# ]
+#
+# [tool.marimo.runtime]
+# on_cell_change = "lazy"
# ///
-"""Fine-tune InceptionV3 and log multiclass ROC curves to W&B."""
import marimo
-__generated_with = "0.24.2"
-app = marimo.App(width="medium", app_title="Plot ROC Curves with W&B")
+__generated_with = "0.25.1"
+app = marimo.App(
+ width="medium",
+ app_title="Plot ROC Curves with Weights & Biases",
+)
with app.setup:
import shutil
import tempfile
import zipfile
from pathlib import Path
+ from urllib.request import urlopen
- import fsspec
import marimo as mo
import numpy as np
+ import pandas as pd
import tensorflow as tf
import wandb
- from tensorflow.keras.applications.inception_v3 import InceptionV3
+ from tensorflow.keras.applications.inception_v3 import InceptionV3, preprocess_input
from tensorflow.keras.callbacks import Callback
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
@@ -38,7 +43,6 @@
from wandb.integration.keras import WandbMetricsLogger
DATASET_URL = "https://storage.googleapis.com/wandb_datasets/nature_12K.zip"
- PROJECT_NAME = "custom_roc_curve"
RUNTIME_CACHE_ROOT = (
Path(tempfile.gettempdir()) / "wandb-examples" / "inaturalist-12k"
)
@@ -47,22 +51,40 @@
@app.cell(hide_code=True)
def _():
mo.md(r"""
- # Plot ROC Curves with W&B
+
- Learn how to log [receiver operating characteristic (ROC) curves](https://scikit-learn.org/stable/modules/model_evaluation.html#receiver-operating-characteristic-roc) backed by [Vega](https://vega.github.io/vega/docs/) in [Weights & Biases](https://wandb.ai). This example fine-tunes InceptionV3 with ImageNet weights to identify 10 types of living things from the [iNaturalist 2017 dataset](https://github.com/visipedia/inat_comp).
+
+
- These simple cases explain the basics; W&B's query editor supports more sophisticated custom charts.
+ # Plot ROC Curves with Weights & Biases
- 
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/plot-roc-curves-with-w-b/plot_roc_curves_with_w_b.py/server)
+
+ Learn to log multiclass ROC curves with `wandb.plot.roc_curve()`. This lesson
+ trains a classification head on a frozen, ImageNet-trained InceptionV3 model
+ to identify ten groups of living things from an iNaturalist image sample.
- The defaults intentionally use few training examples and one epoch so you can exercise the workflow. The original tutorial reports validation accuracy in the low 80s after roughly one epoch on the complete training dataset.
+ An ROC curve plots the true-positive rate against the false-positive rate
+ as the decision threshold changes. For multiclass predictions, each curve
+ treats one class as positive and all other classes as negative.
- ## Method: `wandb.plot.roc_curve()`
+ 
- - Read about [ROC curve customization](https://wandb.ai/wandb/plots/reports/Plot-ROC-Curves--VmlldzoyNjk3MDE).
- - Explore more examples in the [Custom Charts project](https://app.wandb.ai/demo-team/custom-charts).
+ The small training defaults let you try the workflow on a CPU. More data
+ and epochs let you test whether additional training improves class separation.
+ A GPU can reduce training time.
""")
return
@@ -70,422 +92,398 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Authentication
+ ## Authenticate with Weights & Biases
+
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- Enter a [W&B API key](https://wandb.ai/authorize) and an optional team entity below. Leave the key blank when `WANDB_API_KEY` is configured in the Molab Secrets panel or this runtime already has W&B credentials. If needed, [create a free W&B account](https://wandb.ai/signup).
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- Your entity is the first path segment after `wandb.ai` when you open a W&B workspace. The API key is used only for `wandb.login()` and is never included in the run configuration or notebook output.
+ ## Set your team and project
+
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `custom_roc_curve`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- _api_key_input = mo.ui.text(
- kind="password",
- label="W&B API key (optional)",
- placeholder="Paste a key or use configured credentials",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- label="W&B entity or team (optional)",
- placeholder="Leave blank to use your default entity",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}")
- .batch(api_key=_api_key_input, entity=_entity_input)
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="custom_roc_curve", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(wandb_login_form):
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Connect to W&B before configuring the experiment."),
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _api_key = wandb_login_form.value["api_key"].strip()
- _entity = wandb_login_form.value["entity"].strip()
try:
_login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
- _login_error = None
- except wandb.errors.Error as _error:
+ except (wandb.errors.Error, ValueError):
_login_ok = False
- _login_error = str(_error)
mo.stop(
not _login_ok,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. Check the API key or "
- "Molab Secrets configuration and submit again.\n\n"
- f"W&B reported: `{_login_error or 'unknown error'}`"
- ),
- kind="danger",
- ),
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- wandb_settings = {
- "project": PROJECT_NAME,
- "entity": _entity or None,
- }
- _entity_note = (
- f" Runs will be logged to `{wandb_settings['entity']}`."
- if wandb_settings["entity"]
- else " Runs will use your default W&B entity."
- )
- mo.callout(mo.md(f"Connected to W&B.{_entity_note}"), kind="success")
- return (wandb_settings,)
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Setup: Access the remote dataset
+ ## Experiment configuration
- The iNaturalist sample contains 10,000 training images and 2,000 validation images, evenly distributed across 10 classes such as birds, insects, plants, and mammals. **The archive is approximately 3.6 GB, so the first transfer and extraction can take several minutes.**
+ Start with 100 training images, 500 validation images, and one epoch. We
+ sample evenly across the ten classes so every class has positive and
+ negative examples for its ROC curve. Keep the seed and validation size
+ fixed when comparing experiments.
- The `dataset_fs` value below exposes the public archive through an `fsspec` reference filesystem, so you can browse it in marimo's data side panel without transferring the whole file.
-
- Training still requires local image directories. The first time you submit the training form, `ensure_dataset()` streams the archive into a temporary runtime cache and extracts it there. Later submissions reuse that cache; no dataset files are written beside the notebook.
+ Passing this dictionary to `wandb.init(config=config)` records the settings
+ with the run. Try increasing `num_train` or `epochs` after your first run.
""")
return
@app.cell
def _():
- # Expose the remote dataset as a named filesystem so it appears in
- # marimo's data side panel.
- dataset_fs = fsspec.filesystem(
- "reference",
- fo={
- "version": 1,
- "refs": {"nature_12K.zip": [DATASET_URL]},
- },
- )
- return (dataset_fs,)
+ config = {
+ "num_train": 100,
+ "num_val": 500,
+ "num_classes": 10,
+ "fc_size": 1024,
+ "img_width": 299,
+ "img_height": 299,
+ "batch_size": 32,
+ "epochs": 1,
+ "random_seed": 23,
+ }
+ return (config,)
-@app.function
-def ensure_dataset(dataset_fs):
- """Materialize and extract the remote dataset once, then return its paths."""
- archive_path = RUNTIME_CACHE_ROOT / "nature_12K.zip"
- train_path = RUNTIME_CACHE_ROOT / "inaturalist_12K" / "train"
- validation_path = RUNTIME_CACHE_ROOT / "inaturalist_12K" / "val"
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Prepare the dataset
- if train_path.is_dir() and validation_path.is_dir():
- return str(train_path), str(validation_path), "Reused the extracted dataset."
+ The [iNaturalist sample](https://github.com/visipedia/inat_comp) contains
+ 10,000 training images and 2,000 validation images across ten classes.
+ **The first run downloads and extracts an approximately 3.6 GB archive,**
+ even with the small training defaults. Later runs reuse the temporary
+ runtime cache.
- RUNTIME_CACHE_ROOT.mkdir(parents=True, exist_ok=True)
- if not archive_path.exists():
- partial_path = archive_path.with_suffix(".zip.part")
- # Keras needs local directories, so stream the remote archive into the
- # runtime cache without loading 3.6 GB into memory.
- with dataset_fs.open("nature_12K.zip", "rb") as source:
- with partial_path.open("wb") as destination:
- shutil.copyfileobj(source, destination)
- partial_path.replace(archive_path)
-
- with zipfile.ZipFile(archive_path) as archive:
- archive.extractall(archive_path.parent)
+ `sample_images` selects the requested number of files from each split,
+ distributing them across classes. The download and sampling helpers are
+ defined at the end of the notebook.
+ """)
+ return
- if not train_path.is_dir() or not validation_path.is_dir():
- raise FileNotFoundError(
- "The archive did not contain inaturalist_12K/train and "
- "inaturalist_12K/val."
- )
- return str(train_path), str(validation_path), "Streamed and extracted the remote dataset."
+@app.cell
+def _(config):
+ train_path, validation_path = ensure_dataset()
+ train_samples = sample_images(train_path, config["num_train"], config["random_seed"])
+ validation_samples = sample_images(
+ validation_path, config["num_val"], config["random_seed"]
+ )
+ class_names = sorted(train_samples["class"].unique())
+ return class_names, train_samples, validation_samples
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Training code
+ ## Prepare image batches
+
+ Augment the training images with shear, zoom, and horizontal flips. Both
+ generators use InceptionV3's [preprocessing function](https://keras.io/api/applications/inceptionv3/inception_v3_preprocessing/),
+ which scales pixel values to the range −1 to 1.
+
+ Validation images are neither augmented nor shuffled. The same selected
+ images, class IDs, and class order are used for validation and ROC logging
+ at every epoch.
+ """)
+ return
+
+
+@app.cell
+def _(class_names, config, train_samples, validation_samples):
+ train_datagen = ImageDataGenerator(
+ preprocessing_function=preprocess_input,
+ shear_range=0.2,
+ zoom_range=0.2,
+ horizontal_flip=True,
+ )
+ validation_datagen = ImageDataGenerator(preprocessing_function=preprocess_input)
+
+ train_generator = train_datagen.flow_from_dataframe(
+ train_samples,
+ x_col="filename",
+ y_col="class",
+ classes=class_names,
+ target_size=(config["img_height"], config["img_width"]),
+ batch_size=config["batch_size"],
+ class_mode="categorical",
+ shuffle=True,
+ seed=config["random_seed"],
+ )
+ validation_generator = validation_datagen.flow_from_dataframe(
+ validation_samples,
+ x_col="filename",
+ y_col="class",
+ classes=class_names,
+ target_size=(config["img_height"], config["img_width"]),
+ batch_size=config["batch_size"],
+ class_mode="categorical",
+ shuffle=False,
+ )
+ return train_generator, validation_generator
+
- The model freezes the ImageNet-trained InceptionV3 base and trains a new classification head. At the end of each epoch, `ROCCurveLogger` predicts several validation batches and sends the ground-truth class IDs and per-class probabilities to `wandb.plot.roc_curve()`.
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Build and compile the model
- Change the number of training examples or epochs in the form to compare ROC curves. The small defaults favor demo speed over accuracy.
+ Freeze the pretrained InceptionV3 base and train a new classification head.
+ The final softmax layer produces one probability per class. Rerun the model
+ creation cell before each new experiment to start with a fresh head.
+ Rerunning only the training cell continues from the current model weights.
""")
return
@app.function
def build_model(fc_size, num_classes):
- """Freeze an ImageNet-trained InceptionV3 base and add a classifier head."""
base = InceptionV3(weights="imagenet", include_top=False)
- for layer in base.layers:
- layer.trainable = False
-
- features = base.get_layer("mixed10").output
- features = GlobalAveragePooling2D()(features)
+ base.trainable = False
+ features = GlobalAveragePooling2D()(base.output)
features = Dense(fc_size, activation="relu")(features)
- predictions = Dense(num_classes, activation="softmax")(features)
+ probabilities = Dense(num_classes, activation="softmax")(features)
+ return Model(inputs=base.input, outputs=probabilities)
+
- model = Model(inputs=base.input, outputs=predictions)
+@app.cell
+def _(config):
+ tf.keras.backend.clear_session()
+ tf.keras.utils.set_random_seed(config["random_seed"])
+ model = build_model(config["fc_size"], config["num_classes"])
model.compile(
optimizer="rmsprop",
loss="categorical_crossentropy",
metrics=["accuracy"],
)
- return model
+ model.summary()
+ return (model,)
-@app.function
-def make_roc_curve_logger(run, generator, num_log_batches):
- """Create a Keras callback that logs multiclass ROC curves each epoch."""
-
- class ROCCurveLogger(Callback):
- def __init__(self):
- super().__init__()
- self.generator = generator
- self.num_batches = num_log_batches
- self.class_names = [
- name
- for name, _class_id in sorted(
- generator.class_indices.items(), key=lambda item: item[1]
- )
- ]
-
- def on_epoch_end(self, epoch, logs=None):
- validation_batches = [
- self.generator[index] for index in range(self.num_batches)
- ]
- validation_data = np.vstack(
- [batch_data for batch_data, _batch_labels in validation_batches]
- )
- validation_labels = np.vstack(
- [batch_labels for _batch_data, batch_labels in validation_batches]
- )
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Log ROC curves at each epoch
- validation_probabilities = self.model.predict(
- validation_data, verbose=0
- )
- ground_truth_class_ids = validation_labels.argmax(axis=1)
-
- # Keep this key stable so curves from multiple runs share a panel.
- run.log(
- {
- "roc_curve": wandb.plot.roc_curve(
- ground_truth_class_ids,
- validation_probabilities,
- labels=self.class_names,
- )
- }
- )
+ The callback passes three values to [`wandb.plot.roc_curve`](https://docs.wandb.ai/models/ref/python/custom-charts/roc_curve):
- return ROCCurveLogger()
+ - Ground-truth class IDs, with shape `(num_images,)`.
+ - Predicted probabilities, with shape `(num_images, num_classes)`.
+ - Class names in the same order as the probability columns.
+ Pass probabilities so the chart can evaluate many decision thresholds.
+ Predicted class IDs alone discard the scores needed to build the curves. The callback processes one batch of images
+ at a time and collects only labels and probabilities for the chart.
+
+ Keep the `roc_curve` logging key consistent so runs can share a chart panel.
+ """)
+ return
-@app.function
-def train_and_log(config, wandb_settings, train_path, validation_path):
- """Train the model and complete one W&B run."""
- if wandb.run is not None:
- wandb.finish()
-
- with wandb.init(
- project=wandb_settings["project"],
- entity=wandb_settings["entity"],
- config=config,
- ) as run:
- run_config = run.config
- tf.random.set_seed(run_config.random_seed)
- np.random.seed(run_config.random_seed)
-
- train_datagen = ImageDataGenerator(
- rescale=1.0 / 255,
- shear_range=0.2,
- zoom_range=0.2,
- horizontal_flip=True,
- )
- validation_datagen = ImageDataGenerator(rescale=1.0 / 255)
- train_generator = train_datagen.flow_from_directory(
- train_path,
- target_size=(run_config.img_width, run_config.img_height),
- batch_size=run_config.batch_size,
- class_mode="categorical",
+@app.class_definition
+class ROCCurveLogger(Callback):
+ def __init__(self, run, generator):
+ super().__init__()
+ self.run = run
+ self.generator = generator
+ self.class_names = sorted(
+ generator.class_indices, key=generator.class_indices.get
)
- validation_generator = validation_datagen.flow_from_directory(
- validation_path,
- target_size=(run_config.img_width, run_config.img_height),
- batch_size=run_config.batch_size,
- class_mode="categorical",
- # Shuffle so the validation subset logged to the ROC chart samples
- # every class instead of the first few directory-ordered classes.
- shuffle=True,
- seed=run_config.random_seed,
+
+ def on_epoch_end(self, epoch, logs=None):
+ class_ids = []
+ probabilities = []
+ for index in range(len(self.generator)):
+ images, labels = self.generator[index]
+ class_ids.append(labels.argmax(axis=1))
+ probabilities.append(self.model.predict_on_batch(images))
+
+ self.run.log(
+ {
+ "roc_curve": wandb.plot.roc_curve(
+ np.concatenate(class_ids),
+ np.concatenate(probabilities),
+ labels=self.class_names,
+ )
+ }
)
- model = build_model(run_config.fc_size, run_config.num_classes)
- callbacks = [
- WandbMetricsLogger(),
- make_roc_curve_logger(
- run,
- validation_generator,
- run_config.num_log_batches,
- ),
- ]
+
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Train
+
+ Create a Weights & Biases run, then pass both callbacks to `model.fit`.
+ `WandbMetricsLogger` records loss and accuracy. `ROCCurveLogger` records
+ per-class ROC curves. Training and validation use all of their selected
+ images, including the final partial batch.
+
+ The context manager finishes the run when training ends.
+ """)
+ return
+
+
+@app.cell
+def _(config, entity, model, project, train_generator, validation_generator):
+ with wandb.init(entity=entity, project=project, config=config) as run:
history = model.fit(
train_generator,
- steps_per_epoch=max(1, run_config.num_train // run_config.batch_size),
- epochs=run_config.pretrain_epochs,
+ epochs=config["epochs"],
validation_data=validation_generator,
- callbacks=callbacks,
- validation_steps=max(1, run_config.num_val // run_config.batch_size),
+ callbacks=[
+ WandbMetricsLogger(),
+ ROCCurveLogger(run, validation_generator),
+ ],
)
-
- final_validation_accuracy = history.history.get("val_accuracy", [None])[-1]
- result = {
- "run_url": run.url,
- "final_validation_accuracy": (
- float(final_validation_accuracy)
- if final_validation_accuracy is not None
- else None
- ),
- }
-
- return result
+ return
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Run the experiment
+ ## Explore your results
- After connecting to W&B above, submit this training form to materialize the remote dataset in the runtime cache if necessary, download the InceptionV3 weights if necessary, train the model, and create one run in the `custom_roc_curve` project. A later submission finishes any prior active run before it starts another.
+ Open the run link printed by Weights & Biases and find the `roc_curve`
+ custom chart. Check that it contains all ten class names. Compare curves
+ across epochs and runs: a curve closer to the upper-left corner indicates
+ better separation of that class from the others.
+
+ Inspect `epoch/loss`, `epoch/val_loss`, and `epoch/val_accuracy` alongside
+ the curves. Increase `num_train` or `epochs`, rebuild the model, and train
+ again while keeping the validation sample fixed.
+
+ See [ROC chart customization](https://wandb.ai/wandb/plots/reports/Plot-ROC-Curves--VmlldzoyNjk3MDE)
+ and the [Custom Charts project](https://app.wandb.ai/demo-team/custom-charts)
+ for more examples.
""")
return
@app.cell(hide_code=True)
-def _(wandb_settings):
- _workflow_ready = callable(ensure_dataset) and callable(train_and_log)
- _num_train_input = mo.ui.number(
- start=32,
- stop=10_000,
- step=1,
- value=100,
- label="Training examples",
- )
- _num_epochs_input = mo.ui.number(
- start=1,
- stop=20,
- step=1,
- value=1,
- label="Epochs",
- )
- training_form = (
- mo.md("{num_train}\n\n{num_epochs}")
- .batch(
- num_train=_num_train_input,
- num_epochs=_num_epochs_input,
- )
- .form(
- submit_button_label="Prepare data, train, and log to W&B",
- submit_button_disabled=not _workflow_ready,
- bordered=True,
- )
- )
- _destination = (
- f"`{wandb_settings['entity']}/{wandb_settings['project']}`"
- if wandb_settings["entity"]
- else f"`{wandb_settings['project']}` in your default entity"
- )
- mo.vstack(
- [
- mo.md(f"This run will be created in {_destination}."),
- training_form,
- ]
- )
- return (training_form,)
+def _():
+ mo.md(r"""
+ ## Helper functions
+ Cache the archive and extracted directories, then select a reproducible
+ sample of images from each class. Use at least ten images per split to
+ include every class, up to 10,000 for training and 2,000 for validation.
+ """)
+ return
-@app.cell(hide_code=True)
-def _(training_form, wandb_settings):
- mo.stop(
- training_form.value is None,
- mo.callout(
- mo.md("Review the workflow above, then submit the form to begin."),
- kind="info",
- ),
- )
- training_request = {
- "wandb_settings": dict(wandb_settings),
- "num_train": int(training_form.value["num_train"]),
- "num_epochs": int(training_form.value["num_epochs"]),
- }
- return (training_request,)
+@app.function
+def dataset_classes_with_images(directory):
+ directory = Path(directory)
+ if not directory.is_dir():
+ return ()
+ class_dirs = sorted(path for path in directory.iterdir() if path.is_dir())
+ if not class_dirs or any(
+ not any(
+ path.is_file() and path.suffix.lower() in {".jpg", ".jpeg", ".png"}
+ for path in class_dir.rglob("*")
+ )
+ for class_dir in class_dirs
+ ):
+ return ()
+ return tuple(path.name for path in class_dirs)
-@app.cell
-def _(training_request):
- # Experiment configuration saved to W&B.
- config_defaults = {
- "num_train": training_request["num_train"],
- "num_val": 500,
- "num_classes": 10,
- "fc_size": 1024,
- "img_width": 299,
- "img_height": 299,
- "batch_size": 32,
- "pretrain_epochs": training_request["num_epochs"],
- "num_log_batches": 15,
- "random_seed": 23,
- }
- config_defaults
- return (config_defaults,)
+@app.function
+def ensure_dataset():
+ """Materialize and extract the remote dataset once, then return its paths."""
+ archive_path = RUNTIME_CACHE_ROOT / "nature_12K.zip"
+ train_path = RUNTIME_CACHE_ROOT / "inaturalist_12K" / "train"
+ validation_path = RUNTIME_CACHE_ROOT / "inaturalist_12K" / "val"
+ train_classes = dataset_classes_with_images(train_path)
+ if len(train_classes) == 10 and train_classes == dataset_classes_with_images(validation_path):
+ return str(train_path), str(validation_path)
-@app.cell(hide_code=True)
-def _(config_defaults, dataset_fs):
- train_path, validation_path, dataset_status = ensure_dataset(dataset_fs)
- mo.callout(
- mo.md(
- f"{dataset_status} Training will use "
- f"{config_defaults['num_train']:,} images per epoch."
- ),
- kind="success",
- )
- return train_path, validation_path
+ RUNTIME_CACHE_ROOT.mkdir(parents=True, exist_ok=True)
+ if not archive_path.exists():
+ partial_path = archive_path.with_suffix(".zip.part")
+ # Keras needs local directories, so stream the remote archive into the
+ # runtime cache without loading 3.6 GB into memory.
+ with urlopen(DATASET_URL, timeout=60) as source:
+ with partial_path.open("wb") as destination:
+ shutil.copyfileobj(source, destination)
+ partial_path.replace(archive_path)
+ with zipfile.ZipFile(archive_path) as archive:
+ archive.extractall(archive_path.parent)
-@app.cell
-def _(config_defaults, train_path, training_request, validation_path):
- experiment_result = train_and_log(
- config=config_defaults,
- wandb_settings=training_request["wandb_settings"],
- train_path=train_path,
- validation_path=validation_path,
- )
- return (experiment_result,)
+ train_classes = dataset_classes_with_images(train_path)
+ if len(train_classes) != 10 or train_classes != dataset_classes_with_images(validation_path):
+ raise ValueError(
+ "The extracted dataset must contain the same ten classes with images "
+ f"in both train and val. Check {archive_path}; replace an incomplete "
+ "archive and rerun dataset preparation."
+ )
+ return str(train_path), str(validation_path)
-@app.cell(hide_code=True)
-def _(experiment_result):
- _accuracy = experiment_result["final_validation_accuracy"]
- _accuracy_note = (
- f" The final validation accuracy was `{_accuracy:.3f}`."
- if _accuracy is not None
- else ""
- )
- mo.callout(
- mo.md(
- f"Training finished.{_accuracy_note} "
- f"[Open the run in W&B]({experiment_result['run_url']}) and find "
- "`roc_curve` under **Custom Charts** to inspect the ROC curves."
- ),
- kind="success",
- )
- return
+
+@app.function
+def sample_images(directory, num_images, seed):
+ class_dirs = sorted(path for path in Path(directory).iterdir() if path.is_dir())
+ if not class_dirs:
+ raise ValueError(f"No class folders found in {directory}. Rerun dataset preparation.")
+ if num_images < len(class_dirs):
+ raise ValueError("Choose at least one image per class.")
+ per_class, remainder = divmod(num_images, len(class_dirs))
+ samples = []
+ for index, class_dir in enumerate(class_dirs):
+ filenames = sorted(
+ str(path)
+ for path in class_dir.rglob("*")
+ if path.is_file() and path.suffix.lower() in {".jpg", ".jpeg", ".png"}
+ )
+ requested = per_class + (index < remainder)
+ if len(filenames) < requested:
+ raise ValueError(
+ f"Class '{class_dir.name}' in {directory} has {len(filenames)} images, "
+ f"but sampling requires {requested}. Rerun dataset preparation if "
+ "files are missing, or reduce the requested sample size."
+ )
+ class_images = pd.DataFrame({"filename": filenames, "class": class_dir.name})
+ samples.append(
+ class_images.sample(n=requested, random_state=seed)
+ )
+ return pd.concat(samples, ignore_index=True)
if __name__ == "__main__":
diff --git a/marimo/convert/report-api-quickstart/report_api_quickstart.py b/marimo/convert/report-api-quickstart/report_api_quickstart.py
index 1bb7ac3f..bf824469 100644
--- a/marimo/convert/report-api-quickstart/report_api_quickstart.py
+++ b/marimo/convert/report-api-quickstart/report_api_quickstart.py
@@ -4,7 +4,6 @@
# "marimo>=0.24",
# "pandas>=2.0",
# "pillow>=10.0",
-# "requests>=2.31",
# "wandb>=0.19.10",
# "wandb-workspaces>=0.1",
# ]
@@ -12,19 +11,19 @@
import marimo
-__generated_with = "0.24.2"
-app = marimo.App(width="medium", app_title="W&B Report API Quickstart")
+__generated_with = "0.25.1"
+app = marimo.App(
+ width="medium",
+ app_title="Weights & Biases Report API Quickstart",
+)
with app.setup:
import math
import random
- import string
- from io import BytesIO
from itertools import product
import marimo as mo
import pandas as pd
- import requests
import wandb
import wandb_workspaces.reports.v2 as wr
from PIL import Image
@@ -33,17 +32,32 @@
@app.cell(hide_code=True)
def _():
mo.md(r"""
- # W&B Report API
+
- Programmatically create, manage, and customize Reports by defining configurations, panel layouts, and runsets with the wandb-workspaces W&B library. Load and modify Reports with URLs, filter and group runs using expressions, and customize run appearances using Report templates.
+
+
+
+ # Weights & Biases Report API Quickstart
+
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/report-api-quickstart/report_api_quickstart.py/server)
- [wandb-workspaces](https://github.com/wandb/wandb-workspaces) is a Python library for programmatically creating and customizing W&B Workspaces and Reports.
+ Use Workspaces and Reports API ([wandb-workspaces](https://github.com/wandb/wandb-workspaces)) to create and
+ customize Weights & Biases Reports. A report combines narrative blocks with panel grids. A grid’s runsets
+ select the runs to include, and its panels visualize data from those runs.
- In this tutorial you will see how to use wandb-workspaces to create and customize W&B Reports.
+ This tutorial covers how to create, save, and load a report. Along the way you will learn how to add common UI elements such as text, media, and panels to your report.
""")
return
@@ -51,962 +65,621 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Authentication
+ ## Authenticate with Weights & Biases
- The notebook dependencies install the W&B SDK and `wandb-workspaces` automatically. Enter your [W&B API key](https://wandb.ai/authorize) and, if needed, a team entity. You can leave the key blank when this environment already has credentials, such as `WANDB_API_KEY` configured in the marimo Secrets panel. If you are new to W&B, [create a free account](https://wandb.ai/signup).
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- Connecting only authenticates and resolves your entity. It does not create runs or Reports. `wandb-workspaces` is currently in Public Preview, and Report objects are uploaded only when the tutorial explicitly calls `report.save()` after you click a creation button.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
+
+ ## Set your team and project
+
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `report-api-quickstart`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- _api_key_input = mo.ui.text(
- kind="password",
- label="W&B API key (optional)",
- placeholder="Paste a key or use configured credentials",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- label="W&B entity or team (optional)",
- placeholder="Leave blank to use your default entity",
- full_width=True,
- )
- _project_input = mo.ui.text(
- value="report-api-quickstart",
- label="Report project",
- full_width=True,
- )
- _lineage_project_input = mo.ui.text(
- value="lineage-example",
- label="Artifact lineage project",
- full_width=True,
- )
- wandb_login_form = (
- mo.md(
- "{api_key}\n\n{entity}\n\n{project}\n\n{lineage_project}"
- )
- .batch(
- api_key=_api_key_input,
- entity=_entity_input,
- project=_project_input,
- lineage_project=_lineage_project_input,
- )
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="report-api-quickstart", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(wandb_login_form):
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Connect to W&B before using the Report API examples."),
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _api_key = wandb_login_form.value["api_key"].strip()
- _submitted_entity = wandb_login_form.value["entity"].strip()
try:
- _login_ok = wandb.login(
- key=_api_key or None,
- relogin=bool(_api_key),
- )
- _login_error = None
- except wandb.errors.Error as _error:
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
_login_ok = False
- _login_error = str(_error)
mo.stop(
not _login_ok,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. Check the API key and "
- f"try again.\n\nW&B reported: `{_login_error or 'unknown error'}`"
- ),
- kind="danger",
- ),
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- ENTITY = _submitted_entity or wandb.Api().default_entity
- mo.stop(
- not ENTITY,
- mo.callout(
- mo.md(
- "W&B did not return a default entity. Enter a team entity in "
- "the form above and reconnect."
- ),
- kind="danger",
- ),
- )
- PROJECT = wandb_login_form.value["project"].strip() or "report-api-quickstart"
- LINEAGE_PROJECT = (
- wandb_login_form.value["lineage_project"].strip() or "lineage-example"
- )
- mo.callout(
- mo.md(
- f"Connected to W&B as `{ENTITY}`. Report examples will use "
- f"`{PROJECT}` and lineage examples will use `{LINEAGE_PROJECT}`."
- ),
- kind="success",
- )
- return ENTITY, LINEAGE_PROJECT, PROJECT
-
-
-@app.function
-def get_image(url):
- response = requests.get(url, timeout=30)
- response.raise_for_status()
- return Image.open(BytesIO(response.content))
-
-
-@app.function
-def log_dummy_data(entity, project):
- run_names = [
- "adventurous-aardvark-1",
- "bountiful-badger-2",
- "clairvoyant-chipmunk-3",
- "dastardly-duck-4",
- "eloquent-elephant-5",
- "flippant-flamingo-6",
- "giddy-giraffe-7",
- "haughty-hippo-8",
- "ignorant-iguana-9",
- "jolly-jackal-10",
- "kind-koala-11",
- "laughing-lemur-12",
- "manic-mandrill-13",
- "neighbourly-narwhal-14",
- "oblivious-octopus-15",
- "philistine-platypus-16",
- "quant-quail-17",
- "rowdy-rhino-18",
- "solid-snake-19",
- "timid-tarantula-20",
- "understanding-unicorn-21",
- "voracious-vulture-22",
- "wu-tang-23",
- "xenic-xerneas-24",
- "yielding-yveltal-25",
- "zooming-zygarde-26",
- ]
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
- opts = ["adam", "sgd"]
- encoders = ["resnet18", "resnet50"]
- learning_rates = [0.01]
- for (i, run_name), (opt, encoder, lr) in zip(
- enumerate(run_names), product(opts, encoders, learning_rates)
- ):
- config = {
- "optimizer": opt,
- "encoder": encoder,
- "learning_rate": lr,
- "momentum": 0.1 * random.random(),
- }
- displacement1 = random.random() * 2
- displacement2 = random.random() * 4
- with wandb.init(
- entity=entity,
- project=project,
- config=config,
- name=run_name,
- reinit="create_new",
- ) as run:
- for step in range(1000):
- run.log(
- {
- "acc": 0.1
- + 0.4
- * (
- math.log(1 + step + random.random())
- + random.random() * run.config.learning_rate
- + random.random()
- + displacement1
- + random.random() * run.config.momentum
- ),
- "val_acc": 0.1
- + 0.4
- * (
- math.log(1 + step + random.random())
- + random.random() * run.config.learning_rate
- - random.random()
- + displacement1
- ),
- "loss": 0.1
- + 0.08
- * (
- 3.5
- - math.log(1 + step + random.random())
- + random.random() * run.config.momentum
- + random.random()
- + displacement2
- ),
- "val_loss": 0.1
- + 0.04
- * (
- 4.5
- - math.log(1 + step + random.random())
- + random.random() * run.config.learning_rate
- - random.random()
- + displacement2
- ),
- }
- )
-
- with wandb.init(
- entity=entity,
- project=project,
- config=config,
- name=run_names[i + 1],
- reinit="create_new",
- ) as run:
- img = get_image(
- "https://www.akc.org/wp-content/uploads/2017/11/Shiba-Inu-standing-in-profile-outdoors.jpg"
- )
- image = wandb.Image(img)
- df = pd.DataFrame(
- {
- "int": [1, 2, 3, 4],
- "float": [1.2, 2.3, 3.4, 4.5],
- "str": ["a", "b", "c", "d"],
- "img": [image] * 4,
- }
- )
- run.log({"img": image, "my-table": df})
-
-
-@app.class_definition
-class Step:
- def __init__(self, j, r, u, o, at=None):
- self.job_type = j
- self.runs = r
- self.uses_per_run = u
- self.outputs_per_run = o
- self.artifact_type = at if at is not None else "model"
- self.artifacts = []
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Create example data for the report
-@app.function
-def create_artifact(name: str, type: str, content: str):
- art = wandb.Artifact(name, type)
- with art.new_file("test-name", mode="w") as f:
- f.write(content)
+ The next cells create synthetic metrics, an image, a Table, and a small
+ Artifact pipeline. These give the report panels data to display without
+ training a model. The run names identify the examples, and the color mapping
+ makes those runs easier to follow across charts.
+ """)
+ return
- img = get_image(
- "https://www.akc.org/wp-content/uploads/2017/11/Shiba-Inu-standing-in-profile-outdoors.jpg"
- )
- image = wandb.Image(img)
- df = pd.DataFrame(
- {
- "int": [1, 2, 3, 4],
- "float": [1.2, 2.3, 3.4, 4.5],
- "str": ["a", "b", "c", "d"],
- "img": [image] * 4,
- }
- )
- art.add(wandb.Table(dataframe=df), "dataframe")
- return art
+@app.cell
+def _():
+ demo_run_colors = {}
-@app.function
-def log_dummy_lineage(entity, lineage_project):
- pipeline = [
- Step("dataset-generator", 1, 0, 3, "dataset"),
- Step("trainer", 4, (1, 2), 3),
- Step("evaluator", 2, 1, 3),
- Step("ensemble", 1, 1, 1),
+ demo_rng = random.Random(0)
+ demo_run_names = [
+ "adventurous-aardvark-1", "bountiful-badger-2",
+ "clairvoyant-chipmunk-3", "dastardly-duck-4", "eloquent-elephant-5",
]
- for i, step in enumerate(pipeline):
- for _ in range(step.runs):
- with wandb.init(
- entity=entity,
- project=lineage_project,
- job_type=step.job_type,
- reinit="create_new",
- ) as run:
- # use
- uses = step.uses_per_run
- if type(uses) == tuple:
- uses = random.choice(list(uses))
-
- if i > 0:
- prev_step = pipeline[i - 1]
- input_artifacts = random.sample(prev_step.artifacts, uses)
- for a in input_artifacts:
- run.use_artifact(a)
- # log output artifacts
- for j in range(step.outputs_per_run):
- # name = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
- name = f"{step.artifact_type}-{j}"
- content = "".join(
- random.choices(string.ascii_lowercase + string.digits, k=12)
- )
- art = create_artifact(name, step.artifact_type, content)
- run.log_artifact(art)
- art.wait()
-
- # save in pipeline
- step.artifacts.append(art)
-
-
-@app.cell(hide_code=True)
-def _(ENTITY, LINEAGE_PROJECT, PROJECT):
- _log_dummy_runs = mo.ui.checkbox(
- value=True,
- label="Log demo runs and Artifact lineage",
- )
- demo_data_form = (
- mo.md("{log_dummy_runs}")
- .batch(log_dummy_runs=_log_dummy_runs)
- .form(
- submit_button_label="Create 13 runs + Artifacts / skip unchecked",
- bordered=True,
- )
- )
- mo.vstack(
- [
- mo.md(
- "The original tutorial can create five metric/media runs in "
- f"`{ENTITY}/{PROJECT}` and eight lineage runs plus Artifacts in "
- f"`{ENTITY}/{LINEAGE_PROJECT}`. Keep the box checked to create "
- "that remote demo data, or clear it to skip the optional setup. "
- "Changing the checkbox does nothing until you submit the form."
- ),
- demo_data_form,
- ]
- )
- return (demo_data_form,)
-
-
-@app.cell(hide_code=True)
-def _(ENTITY, LINEAGE_PROJECT, PROJECT, demo_data_form):
- mo.stop(
- demo_data_form.value is None,
- mo.callout(
- mo.md("Submit the demo-data choice above to continue."),
- kind="info",
- ),
- )
- LOG_DUMMY_RUNS = bool(demo_data_form.value["log_dummy_runs"])
- if LOG_DUMMY_RUNS:
- log_dummy_data(ENTITY, PROJECT)
- log_dummy_lineage(ENTITY, LINEAGE_PROJECT)
- _demo_status = "Created the optional metric/media runs and Artifact lineage."
- else:
- _demo_status = "Skipped the optional demo runs and Artifact lineage."
- mo.callout(mo.md(_demo_status), kind="success")
- return
+ demo_colors = ["#e84118", "#fbc531", "#4cd137", "#00a8ff", "#9c88ff"]
+ return demo_colors, demo_rng, demo_run_colors, demo_run_names
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Quickstart!
+ The next cell creates four runs with synthetic training curves, one media
+ run with an image and Table, and eight runs that log a simulated Artifact
+ pipeline. All of them use your selected project.
+
+ The optimizer and encoder names are example configuration values. No models
+ are trained. The Artifact inputs and outputs create lineage for the report’s
+ lineage panel.
""")
return
-@app.cell(hide_code=True)
-def _(PROJECT):
- create_quickstart_report = mo.ui.run_button(
- label="Create and update the Quickstart Report"
- )
- mo.vstack(
- [
- mo.md(
- f"This action creates one Report in `{PROJECT}`, then saves "
- "each of the changes shown in the next four code cells. A new "
- "button click intentionally performs the sequence again."
- ),
- create_quickstart_report,
- ]
- )
- return (create_quickstart_report,)
+@app.cell
+def _(demo_colors, demo_rng, demo_run_colors, demo_run_names, entity, project):
+ # Log psuedo training metrics such as accuracy and loss
+ for demo_index, (optimizer_name, encoder_name) in enumerate(product(
+ ["adam", "sgd"], ["resnet18", "resnet50"],
+ )):
+ with wandb.init(
+ entity=entity, project=project, name=demo_run_names[demo_index],
+ config={"optimizer": optimizer_name, "encoder": encoder_name,
+ "learning_rate": 0.01, "momentum": demo_rng.random() / 10},
+ ) as metrics_run:
+ for step in range(100):
+ progress = 1 - math.exp(-step / 30)
+ noise = demo_rng.uniform(-0.02, 0.02)
+ metrics_run.log({
+ "acc": 0.2 + 0.75 * progress + noise,
+ "val_acc": 0.2 + 0.7 * progress + noise,
+ "loss": 1.5 * (1 - progress) + 0.1,
+ "val_loss": 1.7 * (1 - progress) + 0.15,
+ })
+ demo_run_colors[metrics_run.id] = demo_colors[demo_index]
+
+
+ # Create a color gradient image
+ demo_image = Image.new("RGB", (64, 64))
+
+ demo_image.putdata([(4 * x, 4 * y, 128) for y in range(64) for x in range(64)])
+ table_image = wandb.Image(demo_image, caption="Synthetic color gradient")
+
+ demo_dataframe = pd.DataFrame({
+ "int": [1, 2, 3, 4], "float": [1.2, 2.3, 3.4, 4.5],
+ "str": ["a", "b", "c", "d"], "img": [table_image] * 4,
+ })
+
+ # Log color gradient image
+ with wandb.init(
+ entity=entity, project=project, name=demo_run_names[-1], job_type="media",
+ ) as media_run:
+ media_run.log({"img": table_image, "my-table": wandb.Table(dataframe=demo_dataframe)})
+ demo_run_colors[media_run.id] = demo_colors[-1]
+ pipeline = [
+ ("dataset-generator", 1, 0, 3, "dataset"),
+ ("trainer", 4, (1, 2), 3, "model"),
+ ("evaluator", 2, 1, 3, "model"),
+ ("ensemble", 1, 1, 1, "model"),
+ ]
+ previous_artifacts = []
-@app.cell(hide_code=True)
-def _(ENTITY, PROJECT, create_quickstart_report):
- mo.stop(
- not create_quickstart_report.value,
- mo.callout(
- mo.md("Click the button above when you are ready to create the Report."),
- kind="info",
- ),
- )
- quickstart_target = {"entity": ENTITY, "project": PROJECT}
- return (quickstart_target,)
+ # Log artifact
+ for job_type, run_count, uses_per_run, output_count, artifact_type in pipeline:
+ stage_artifacts = []
+ for lineage_index in range(run_count):
+ with wandb.init(
+ entity=entity, project=project, job_type=job_type,
+ name=f"report-demo-{job_type}-{lineage_index + 1}",
+ ) as lineage_run:
+ input_count = demo_rng.choice(uses_per_run) if isinstance(uses_per_run, tuple) else uses_per_run
+ for input_artifact in demo_rng.sample(previous_artifacts, input_count):
+ lineage_run.use_artifact(input_artifact)
+ for output_index in range(output_count):
+ lineage_artifact = wandb.Artifact(f"{artifact_type}-{output_index}", type=artifact_type)
+ with lineage_artifact.new_file("stage.txt", mode="w") as stage_file:
+ stage_file.write(f"{job_type}/{lineage_index}/{output_index}/{demo_rng.random()}")
+ lineage_artifact.add(wandb.Table(dataframe=demo_dataframe), "dataframe")
+ uploaded_artifact = lineage_run.log_artifact(lineage_artifact)
+ uploaded_artifact.wait()
+ stage_artifacts.append(uploaded_artifact)
+ previous_artifacts = stage_artifacts
+
+ print("Created demo metrics, media, tables, and Artifact lineage.")
+ return
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Create, save, and load reports
- - NOTE: Reports are not saved automatically to reduce clutter. Explicitly save the report by calling `report.save()`
- """)
- return
+ Use the logged runs as data for the report below.
+ ## Build a Report
-@app.cell
-def _(quickstart_target):
- report = wr.Report(
- entity=quickstart_target["entity"],
- project=quickstart_target["project"],
- title="Quickstart Report",
- description="That was easy!",
- ) # Create
- report.save() # Save
- wr.Report.from_url(report.url) # Load
- return (report,)
-
+ ### Text and image blocks
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ### Add content via blocks
- - Use blocks to add content like text, images, code, and more
- - See `wr.blocks` for all available blocks
+ Define content with `TableOfContents`, headings, paragraphs, and images. These
+ are Report objects, separate from marimo's display elements. The image below is
+ from the original Weights & Biases tutorial. Nothing is uploaded while building these lists.
""")
return
@app.cell
-def _(report):
- report.blocks = [
+def _():
+ text_blocks = [
wr.TableOfContents(),
wr.H1("Text and images example"),
- wr.P("Lorem ipsum dolor sit amet. Aut laborum perspiciatis sit odit omnis aut aliquam voluptatibus ut rerum molestiae sed assumenda nulla ut minus illo sit sunt explicabo? Sed quia architecto est voluptatem magni sit molestiae dolores. Non animi repellendus ea enim internos et iste itaque quo labore mollitia aut omnis totam."),
+ wr.P("Use a Report to explain an experiment, show results, and share conclusions with your team."),
wr.Image("https://api.wandb.ai/files/telidavies/images/projects/831572/8ad61fd1.png", caption="Craiyon generated images"),
- wr.P("Et voluptatem galisum quo facilis sequi quo suscipit sunt sed iste iure! Est voluptas adipisci et doloribus commodi ab tempore numquam qui tempora adipisci. Eum sapiente cupiditate ut natus aliquid sit dolor consequatur?"),
+ wr.P("Panel grids connect this narrative to the underlying experiment data."),
]
- report.save()
- report_with_content = report
- return (report_with_content,)
+ return (text_blocks,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Add charts and more via Panel Grid
- - `PanelGrid` is a special type of block that holds `runsets` and `panels`
- - `runsets` organize data logged to W&B
- - `panels` visualize runset data. For a full set of panels, see `wr.panels`
+ ### Panel grids and run colors
+
+ A `PanelGrid` combines runsets and panels. The first runset selects all runs in
+ the project. The second searches for `elephant`, matching the demo media run.
+ Panels show validation accuracy, accuracy bars, images, and differing run settings.
+
+ `custom_run_colors` maps **run IDs** to colors. The demo-data cell records
+ those IDs as it creates runs. For existing data, edit the mapping with IDs from
+ your own project. The `WeaveBlockArtifact` below embeds Artifact lineage from
+ the selected project. In this API, `WeaveBlockArtifact` names a Report block
+ for Artifact views.
""")
return
@app.cell
-def _(ENTITY, PROJECT, report_with_content):
- pg = wr.PanelGrid(
+def _(demo_run_colors, entity, project):
+ panel_grid = wr.PanelGrid(
runsets=[
- wr.Runset(ENTITY, PROJECT, "First Run Set"),
- wr.Runset(ENTITY, PROJECT, "Elephants Only!", query="elephant"),
+ wr.Runset(entity, project, "All runs"),
+ wr.Runset(entity, project, "Elephants only", query="elephant"),
],
panels=[
wr.LinePlot(x="Step", y=["val_acc"], smoothing_factor=0.8),
wr.BarPlot(metrics=["acc"]),
wr.MediaBrowser(media_keys=["img"], num_columns=1),
- wr.RunComparer(diff_only="split", layout={"w": 24, "h": 9}),
+ wr.RunComparer(diff_only="split", layout=wr.Layout(w=24, h=9)),
],
)
-
- report_with_content.blocks = report_with_content.blocks[:1] + [wr.H1("Panel Grid Example"), pg] + report_with_content.blocks[1:]
- report_with_content.save()
- report_with_panel_grid = report_with_content
- return pg, report_with_panel_grid
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ### Add data lineage with Artifact blocks
- - There are equivalent weave panels as well
- """)
- return
-
-
-@app.cell
-def _(ENTITY, LINEAGE_PROJECT, report_with_panel_grid):
+ panel_grid.custom_run_colors = dict(demo_run_colors)
artifact_lineage = wr.WeaveBlockArtifact(
- entity=ENTITY,
- project=LINEAGE_PROJECT,
- artifact="model-1",
- tab="lineage",
+ entity=entity, project=project, artifact="model-1", tab="lineage",
)
-
- report_with_panel_grid.blocks = report_with_panel_grid.blocks[:1] + [wr.H1("Artifact lineage example"), artifact_lineage] + report_with_panel_grid.blocks[1:]
- report_with_panel_grid.save()
- report_with_lineage = report_with_panel_grid
- return (report_with_lineage,)
+ return artifact_lineage, panel_grid
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Customize run colors
- - Pass in a `dict[run_name, color]`
- """)
- return
+ ### Assemble, save, load, and update
+ `Report(...)` creates a local object. Assign its block list to define the
+ content order, then call `save()` to publish it to Weights & Biases.
-@app.cell
-def _(pg, report_with_lineage):
- pg.custom_run_colors = {
- "adventurous-aardvark-1": "#e84118",
- "bountiful-badger-2": "#fbc531",
- "clairvoyant-chipmunk-3": "#4cd137",
- "dastardly-duck-4": "#00a8ff",
- "eloquent-elephant-5": "#9c88ff",
- }
- report_with_lineage.save()
- quickstart_report = report_with_lineage
- quickstart_report_complete = quickstart_report.url
- return quickstart_report, quickstart_report_complete
-
-
-@app.cell(hide_code=True)
-def _(quickstart_report_complete):
- mo.callout(
- mo.md(f"Quickstart complete: [open the saved Report]({quickstart_report_complete})."),
- kind="success",
- )
- return
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## FAQ
+ The next cell publishes a report, reloads it with `Report.from_url()`, updates
+ its description, and saves the change. Inspect the block definitions before
+ running this cell. Rerunning the whole cell constructs and publishes a new
+ report. To update an existing report, load its URL and save that object.
""")
return
-@app.cell(hide_code=True)
-def _(PROJECT, quickstart_report_complete):
- create_faq_reports = mo.ui.run_button(label="Create the FAQ example Reports")
- mo.vstack(
- [
- mo.md(
- f"This action clones the [Quickstart Report]({quickstart_report_complete}) and creates seven more "
- f"Reports in `{PROJECT}` for the examples below. Each click "
- "intentionally creates a fresh set of eight Reports."
- ),
- create_faq_reports,
- ]
- )
- return (create_faq_reports,)
-
+@app.cell
+def _(artifact_lineage, entity, panel_grid, project, text_blocks):
+ quickstart_report = wr.Report(
+ entity=entity, project=project, title="Quickstart Report",
+ description="A Report built with the Weights & Biases Report API",
+ )
+ quickstart_report.blocks = [
+ text_blocks[0], wr.H1("Panel grid example"), panel_grid,
+ wr.H1("Artifact lineage example"), artifact_lineage, *text_blocks[1:],
+ ]
-@app.cell(hide_code=True)
-def _(create_faq_reports, quickstart_report):
- mo.stop(
- not create_faq_reports.value,
- mo.callout(
- mo.md("Click the button above when you are ready to create the FAQ Reports."),
- kind="info",
- ),
- )
- faq_source_report = quickstart_report
- return (faq_source_report,)
+ quickstart_report.save()
+ loaded_report = wr.Report.from_url(quickstart_report.url)
+ loaded_report.description = "Created, loaded, and updated with the Weights & Biases Report API"
+ loaded_report.save()
+ return (quickstart_report,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### My report is too wide/narrow
- - Change the report's width to the right size for you.
- """)
- return
+ ## FAQ examples: build more local drafts
+ ### Resize panels
-@app.cell
-def _(faq_source_report):
- report2 = faq_source_report.save(clone=True)
- report2.width = "fluid"
- report2.save()
- faq_report_urls = [report2.url]
- return (faq_report_urls,)
+ Report width can be `readable`, `fixed`, or `fluid`. The publishing example below
+ will clone the main Report and set its width to `fluid`.
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ### How do I resize panels?
- - Pass a `dict[dim, int]` to `panel.layout`
- - `dim` is a dimension, which can be `x`, `y` (the coordiantes of the top left corner) `w`, `h` (the size of the panel)
- - You can pass any or all dimensions at once
- - The space between two dots in a panel grid is 2.
+ Panel layout controls position (`x`, `y`) and dimensions (`w`, `h`) in grid units.
+ The following parallel coordinates plot uses the full 24-unit panel-grid width.
+ Configuration fields use a `c::` prefix. For example, `c::encoder` selects
+ the `encoder` value logged in each run’s configuration.
""")
return
@app.cell
-def _(ENTITY, PROJECT, faq_report_urls):
+def _(entity, project):
resize_report = wr.Report(
- entity=ENTITY,
- project=PROJECT,
- title="Resizing panels",
- description="Look at this wide parallel coordinates plot!",
- blocks=[
- wr.PanelGrid(
- panels=[
- wr.ParallelCoordinatesPlot(
- columns=[
- wr.ParallelCoordinatesPlotColumn(metric="Step"),
- wr.ParallelCoordinatesPlotColumn(metric="c::model"),
- wr.ParallelCoordinatesPlotColumn(metric="c::optimizer"),
- wr.ParallelCoordinatesPlotColumn(metric="Step"),
- wr.ParallelCoordinatesPlotColumn(metric="val_acc"),
- wr.ParallelCoordinatesPlotColumn(metric="val_loss"),
- ],
- layout=wr.Layout(w=24, h=9) # Adjusting the layout for the plot size
- ),
- ]
- )
- ],
- )
- resize_report.save()
- faq_report_urls_after_resize = [*faq_report_urls, resize_report.url]
- return (faq_report_urls_after_resize,)
+ entity=entity,
+ project=project,
+ title="Resizing panels",
+ description="Look at this wide parallel coordinates plot!",
+ blocks=[
+ wr.PanelGrid(
+ panels=[
+ wr.ParallelCoordinatesPlot(
+ columns=[
+ wr.ParallelCoordinatesPlotColumn(metric="Step"),
+ wr.ParallelCoordinatesPlotColumn(metric="c::encoder"),
+ wr.ParallelCoordinatesPlotColumn(metric="c::optimizer"),
+ wr.ParallelCoordinatesPlotColumn(metric="Step"),
+ wr.ParallelCoordinatesPlotColumn(metric="val_acc"),
+ wr.ParallelCoordinatesPlotColumn(metric="val_loss"),
+ ],
+ layout=wr.Layout(w=24, h=9) # Adjusting the layout for the plot size
+ ),
+ ]
+ )
+ ],
+ )
+ return (resize_report,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### What blocks are available?
- - See `wr.blocks` for a list of blocks.
- - In an IDE or notebook, you can also do `wr.blocks.
+
- Inside the training loop, the script simulates an accuracy and loss value for each epoch. It then logs those values to W&B with `run.log()`. After the script runs, you can view the logged metrics in the W&B App.
+ # Weights & Biases Run Quickstart
- ### Authentication
+ A Weights & Biases run represents one experiment. It records your configuration, metrics,
+ and results so you can inspect an experiment or compare it with other runs.
- To save your experiment in W&B, you need to authenticate.
+ This lesson simulates training to focus on three Weights & Biases operations:
+ `wandb.init()` starts a run, `training_run.log()` records metrics, and
+ `wandb.Api()` retrieves the results.
- Authenticate with W&B one of two ways: run **`wandb login`** in
- your shell before starting marimo, or paste your key into the
- **W&B API key** field in the form below. Get your key from
- [wandb.ai/authorize](https://wandb.ai/authorize).
+ Running the code cells creates a Weights & Biases run and uploads its simulated metrics.
""")
return
@app.cell(hide_code=True)
def _():
- epochs = mo.ui.slider(
- start=1,
- stop=20,
- step=1,
- value=10,
- label="Epochs",
- show_value=True,
- )
-
- lr = mo.ui.slider(
- start=0.001,
- stop=0.1,
- step=0.001,
- value=0.01,
- label="Learning rate",
- show_value=True,
- )
- api_key = mo.ui.text(
- value="",
- kind="password",
- label="W&B API key (blank uses your shell login)",
- )
- project = mo.ui.text(value="my-awesome-project", label="W&B project")
- entity = mo.ui.text(
- value="",
- label="W&B entity \u2014 a team you belong to (blank uses your default)",
- )
+ mo.md(r"""
+ ## Authenticate with Weights & Biases
- # Batch every control into one form so training only kicks off on submit.
- # `form.value` is None until the user clicks Train model, then becomes a dict
- # keyed by the names below - the training cell gates on that.
- form = (
- mo.md(
- """
- **Training**
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- {epochs}
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- {lr}
+ ## Set your team and project
- **W&B run.**
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `my-awesome-project`. Then select **Connect to Weights & Biases**.
+ """)
+ return
- {api_key}
- {project}
+@app.cell(hide_code=True)
+def _():
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="my-awesome-project", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
- {entity}
- """
- )
- .batch(
- epochs=epochs,
- lr=lr,
- api_key=api_key,
- project=project,
- entity=entity,
- )
- .form(submit_button_label="Start run", bordered=False)
- )
- form
- return (form,)
+@app.cell(hide_code=True)
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
+ mo.stop(
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
+ )
-@app.cell
-def _(form):
- run_path = None
+ try:
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
+ _login_ok = False
mo.stop(
- form.value is None,
- mo.md(
- "Configure the run and click Start run to simulate training and log "
- "accuracy and loss to W&B."
- ),
+ not _login_ok,
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- # Get the data from the form
- cfg = form.value
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
+
+
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Define the Experiment Configuration
+
+ The next cell defines a configuration dictionary. Passing it to
+ `wandb.init(config=config)` records these settings with the run so you can
+ compare experiments by hyperparameter.
+
+ We simulate ten epochs with a fixed random seed. The learning rate is
+ recorded to demonstrate configuration tracking. It does not affect our
+ synthetic metrics. In real training, you would use it to configure an
+ optimizer.
+ """)
+ return
- # Dictionary with hyperparameters
+
+@app.cell
+def _():
config = {
- "epochs": cfg["epochs"],
- "lr": cfg["lr"],
+ "epochs": 10,
+ "lr": 0.01,
+ "seed": 42,
}
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Start a Run and Log Metrics
+
+ The next cell starts `training_run` with `wandb.init()`. Inside the loop,
+ `training_run.log()` records a dictionary containing the epoch, accuracy,
+ and loss. In your own training code, replace the simulated values with
+ measured metrics.
+
+ The context manager finishes the run and completes logging when the block
+ exits. We keep its URL for the dashboard and its
+ `entity/project/run_id` path for the API query that follows.
+ """)
+ return
+
- # Authenticate and start the run. Finish any prior run first (marimo keeps the
- # kernel alive across re-submits). A key pasted into the form wins; otherwise
- # fall back to ambient login (shell `wandb login`, WANDB_API_KEY, or netrc).
- # The key is never written to the run config.
- if wandb.run is not None:
- wandb.finish()
- if cfg["api_key"]:
- wandb.login(key=cfg["api_key"])
-
- with wandb.init(project=cfg["project"], entity=cfg["entity"] or None, config=config) as run:
- offset = random.random() / 5
- print(f"lr: {config['lr']}")
-
- # Simulate a training run
- for epoch in range(1, config['epochs'] + 1):
- acc = 1 - 2**-config['epochs'] - random.random() / config['epochs'] - offset
- loss = 2**-config['epochs'] + random.random() / config['epochs'] + offset
- print(f"epoch={epoch}, accuracy={acc}, loss={loss}")
- run.log(
- {
- "epoch": epoch,
- "accuracy": acc,
- "loss": loss
- }
- )
-
- run_path = f"{run.entity}/{run.project}/{run.id}"
+@app.cell
+def _(config, entity, project):
+ rng = random.Random(config["seed"])
+ offset = rng.random() / 5
+
+ with wandb.init(project=project, entity=entity, config=config) as training_run:
+ for epoch in range(1, config["epochs"] + 1):
+ accuracy = 1 - 2 ** -epoch - rng.random() / config["epochs"] - offset
+ loss = 2 ** -epoch + rng.random() / config["epochs"] + offset
+ training_run.log({
+ "epoch": epoch,
+ "accuracy": accuracy,
+ "loss": loss,
+ })
+ print(f"Epoch {epoch}: accuracy={accuracy:.4f}, loss={loss:.4f}")
+
+ # For demo: Store run path to retrieve run info (see following cell)
+ run_path = f"{training_run.entity}/{training_run.project}/{training_run.id}"
return (run_path,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Get the logged run results from W&B
+ ## Retrieve the Logged Results
+
+ The next cell uses `wandb.Api().run(run_path)` to fetch the completed
+ experiment. `logged_run.config` contains its settings, and
+ `logged_run.summary` includes the latest logged metric values.
+
+ `scan_history()` iterates over every recorded row containing the requested
+ keys. We collect those rows in a list so we can inspect the complete epoch
+ history, rather than a sampled subset.
""")
return
@app.cell
def _(run_path):
- history = []
- remote_run = None
-
- mo.stop(
- run_path is None,
- mo.md("Run an experiment above to load its logged results from W&B."),
- )
+ api = wandb.Api()
+ logged_run = api.run(run_path)
+ history = list(logged_run.scan_history(keys=["epoch", "accuracy", "loss"]))
+ history
+ return
- _api = wandb.Api()
- remote_run = _api.run(run_path)
- history = [
- {
- "epoch": row["epoch"],
- "accuracy": round(row["accuracy"], 4),
- "loss": round(row["loss"], 4),
- }
- for row in remote_run.scan_history(
- keys=["epoch", "accuracy", "loss"]
- )
- ]
- return history, remote_run
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Explore the Run in Weights & Biases
+ Open the Weights & Biases run to view the accuracy and loss charts. Its
+ configuration contains the settings from this notebook, and its summary
+ shows the most recently logged metric values.
-@app.cell
-def _(history, remote_run):
- mo.stop(
- not history,
- mo.md("The run finished, but no metric history is available yet."),
- )
+ Change the seed or epoch count and rerun the experiment. Weights & Biases creates a new
+ run that you can compare with the previous one in your project. In a real
+ training workflow, changing the learning rate would let you compare how that
+ hyperparameter affects measured performance.
- final = history[-1]
-
- mo.vstack(
- [
- mo.callout(
- mo.md(
- f"**Run complete:** "
- f"[{remote_run.name}]({remote_run.url})\n\n"
- f"Final accuracy: **{final['accuracy']:.2%}** \n"
- f"Final loss: **{final['loss']:.4f}**"
- ),
- kind="success",
- ),
- mo.ui.table(history, selection=None),
- ]
- )
+ See the [Weights & Biases quickstart](https://docs.wandb.ai/models/quickstart)
+ for more examples.
+ """)
return
diff --git a/marimo/convert/simple-pytorch-integration/simple_pytorch_integration.py b/marimo/convert/simple-pytorch-integration/simple_pytorch_integration.py
index 516638a1..e473a367 100644
--- a/marimo/convert/simple-pytorch-integration/simple_pytorch_integration.py
+++ b/marimo/convert/simple-pytorch-integration/simple_pytorch_integration.py
@@ -9,11 +9,14 @@
# "tqdm>=4.66",
# "wandb>=0.18",
# ]
+#
+# [tool.marimo.runtime]
+# on_cell_change = "lazy"
# ///
import marimo
-__generated_with = "0.24.2"
+__generated_with = "0.25.1"
app = marimo.App(width="medium", app_title="Simple PyTorch Integration")
with app.setup:
@@ -41,34 +44,36 @@
@app.cell(hide_code=True)
def _():
mo.md(r"""
- # W&B + PyTorch
+
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- Use [Weights & Biases](https://wandb.com) for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## Example Gallery
-
- See examples of projects tracked and visualized with W&B in our [Gallery →](https://app.wandb.ai/gallery)
- """)
- return
+ Change `learning_rate` or `batch_size` in `config`, then run the experiment
+ cell again. This creates a new run. Compare both runs in your selected project to see how the change affected loss and test accuracy.
+ ## Next steps: Test Hyperparameters with Sweeps
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## Advanced Setup
+ We only looked at a single set of hyperparameters in this example.
+ But an important part of most ML workflows is iterating over
+ a number of hyperparameters.
- 1. [Environment variables](https://docs.wandb.ai/models/track/environment-variables): Set API keys in environment variables so you can run training on a managed cluster.
- 2. [Offline mode](https://docs.wandb.ai/models/track/run): Set `WANDB_MODE=offline` to train offline and sync results later.
- 3. [On-prem](https://docs.wandb.ai/platform/hosting): Install W&B in a private cloud or air-gapped servers in your own infrastructure. We have local installations for everyone from academics to enterprise teams.
- 4. [Sweeps](https://docs.wandb.ai/models/sweeps): Set up hyperparameter search quickly with our lightweight tool for tuning.
+ You can use [Weights & Biases Sweeps](https://docs.coreweave.com/products/wandb/sweeps) to automate hyperparameter testing and explore the space of possible models and optimization strategies.
""")
return
@@ -758,33 +479,15 @@ def _():
mo.md(r"""
## Helper functions
- ### Training and reproducibility helpers
-
- The forward pass, backward pass, and optimizer step are standard PyTorch. Resetting numeric seeds for every submitted experiment makes comparisons repeatable on the same device and software stack.
+ Resetting the random seeds before each experiment makes runs easier to
+ compare on the same device and software stack.
""")
return
-@app.function
-def train_batch(images, labels, model, optimizer, criterion):
- images, labels = images.to(device), labels.to(device)
-
- # Forward pass ➡
- outputs = model(images)
- loss = criterion(outputs, labels)
-
- # Backward pass ⬅
- optimizer.zero_grad()
- loss.backward()
-
- # Step with optimizer
- optimizer.step()
- return loss.item()
-
-
@app.function
def seed_everything(seed):
- # Ensure deterministic behavior
+ # Reset random number generators before each experiment.
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
diff --git a/marimo/convert/torchtune-and-wandb/torchtune_and_wandb.py b/marimo/convert/torchtune-and-wandb/torchtune_and_wandb.py
index 90baf76b..b1f64f9e 100644
--- a/marimo/convert/torchtune-and-wandb/torchtune_and_wandb.py
+++ b/marimo/convert/torchtune-and-wandb/torchtune_and_wandb.py
@@ -9,351 +9,230 @@
# "peft==0.21.0",
# "torch==2.6.0",
# "transformers==5.17.0",
-# "trl==1.13.0",
# "wandb==0.30.0",
# ]
# ///
import marimo
-__generated_with = "0.24.2"
-app = marimo.App(width="medium", app_title="torchtune and W&B")
+__generated_with = "0.25.1"
+app = marimo.App(
+ width="medium",
+ app_title="LoRA fine-tuning with Weights & Biases",
+)
-
-@app.cell
-def _():
- import os
- import uuid
+with app.setup:
+ import tempfile
from pathlib import Path
import marimo as mo
import torch
import wandb
from datasets import load_dataset
- from huggingface_hub import HfApi
- from peft import LoraConfig, TaskType
- from transformers import AutoModelForCausalLM, AutoTokenizer
- from trl import SFTConfig, SFTTrainer
-
- return (
- AutoModelForCausalLM,
- AutoTokenizer,
- HfApi,
- LoraConfig,
- Path,
- SFTConfig,
- SFTTrainer,
- TaskType,
- load_dataset,
- mo,
- os,
- torch,
- uuid,
- wandb,
+ from peft import LoraConfig, TaskType, get_peft_model
+ from transformers import (
+ AutoModelForCausalLM, AutoTokenizer, DataCollatorForSeq2Seq,
+ Trainer, TrainingArguments,
)
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/torchtune-and-wandb/torchtune_and_wandb.py/server)
+
- This notebook uses the maintained Hugging Face training stack:
- [Transformers](https://huggingface.co/docs/transformers),
- [TRL](https://huggingface.co/docs/trl), and
- [PEFT](https://huggingface.co/docs/peft). You will stream a conversational
- dataset, train a compact LoRA adapter, and inspect the live metrics and saved
- adapter in W&B.
+
+
- The default model is the Apache-2.0 licensed
- [`HuggingFaceTB/SmolLM2-135M-Instruct`](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct),
- whose model weights are roughly 269 MB.
- """)
- return
+ # Fine-tune SmolLM2 with LoRA and Weights & Biases
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/torchtune-and-wandb/torchtune_and_wandb.py/server)
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Prerequisites
+ This lesson uses [Transformers](https://huggingface.co/docs/transformers) and
+ [PEFT](https://huggingface.co/docs/peft) to fine-tune a small language model
+ with LoRA. You will track the experiment in Weights & Biases, inspect training metrics,
+ and save the trained adapter as a model Artifact.
- Storage use is deliberately bounded: the dataset is streamed instead of fully
- cached, the base model is much smaller than the former 7B checkpoint, and PEFT
- saves only the trained LoRA adapter. Opening the notebook does not download a
- model, start training, or create a W&B run.
+ The public model `HuggingFaceTB/SmolLM2-135M-Instruct` has 135 million
+ parameters. We stream a small sample of conversational data and train for
+ ten steps to demonstrate the integration. This is a short workflow example,
+ not an evaluation of model quality.
+
+ The default model and dataset are public. For a private or gated repository, configure `HF_TOKEN` as well.
""")
return
@app.cell(hide_code=True)
-def _(torch):
- _cuda_available = torch.cuda.is_available()
- _mps_available = (
- torch.backends.mps.is_built() and torch.backends.mps.is_available()
- )
-
- if _cuda_available:
- training_device = "cuda"
- model_dtype = (
- torch.bfloat16
- if torch.cuda.is_bf16_supported()
- else torch.float16
- )
- elif _mps_available:
- training_device = "mps"
- model_dtype = torch.float32
- else:
- training_device = "cpu"
- model_dtype = torch.float32
- return model_dtype, training_device
+def _():
+ mo.md(r"""
+ ## Authenticate with Weights & Biases
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Authentication and remote storage
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- Enter your [W&B API key](https://wandb.ai/authorize), or leave it blank to use
- `WANDB_API_KEY` from molab Secrets or existing credentials. The entity is the
- team name in `wandb.ai/
+
+
[](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/train-and-debug-yolov5-models-with-weights-biases/train_and_debug_yolov5_models_with_weights_biases.py/server)
# Train and Debug Ultralytics YOLO Models with Weights & Biases
- This notebook fine-tunes the current **Ultralytics YOLO26** model on a
- chess-piece dataset and tracks the experiment with Weights & Biases.
- You will:
-
- 1. run pretrained object detection,
- 2. version a custom dataset with a W&B Artifact,
- 3. fine-tune YOLO26 with the built-in W&B callback, and
- 4. inspect metrics, diagnostic plots, and the best-model Artifact.
+ In this lesson, you will use a pretrained object detector, fine-tune it to
+ recognize chess pieces, and use Weights & Biases to inspect its training metrics,
+ predictions, dataset, and model checkpoint.
- The notebook uses the maintained [`ultralytics`](https://pypi.org/project/ultralytics/)
- package and its [W&B integration](https://docs.ultralytics.com/integrations/weights-biases/).
+ The original Colab used YOLOv5. This version uses **YOLO26** through the
+ maintained [`ultralytics`](https://docs.ultralytics.com/) package and its
+ [Weights & Biases integration](https://docs.ultralytics.com/integrations/weights-biases/).
+ The lesson still covers object detection and experiment tracking, using
+ the Python API in place of the original command-line scripts.
- Opening the notebook does not create a W&B run or upload anything. Those
- actions begin only after you submit the training form.
+ Open the notebook in the marimo editor and submit the setup form before
+ running the training example. The prediction and data-preparation cells
+ download their inputs when executed. Review the settings before running
+ training. After changing a setting, rerun the dependent cells.
""")
return
@app.cell(hide_code=True)
-def _(torch):
- if torch.cuda.is_available():
- training_device = 0
- elif torch.backends.mps.is_available():
- training_device = "mps"
- else:
- training_device = "cpu"
+def _():
+ mo.md(r"""
+ ## Authenticate with Weights & Biases
- return (training_device,)
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Prepare the data
+ ## Set your team and project
- The tutorial uses a small public dataset whose images and labels already
- follow Ultralytics' YOLO format. The download is gated so opening the
- notebook does not write files.
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `yolo-wandb-demo`. Then select **Connect to Weights & Biases**.
""")
return
-@app.cell
-def _(Path, io, tempfile, urlopen, yaml, zipfile):
- def download_chess_dataset():
- """Download and validate the public chess-piece dataset."""
- workspace_path = Path(tempfile.mkdtemp(prefix="wandb-yolo26-"))
- dataset_url = "https://public.roboflow.com/ds/1BpjFZe9ST?key=KXD7eDvwTa"
-
- with urlopen(dataset_url, timeout=120) as dataset_response:
- with zipfile.ZipFile(io.BytesIO(dataset_response.read())) as archive:
- for member in archive.infolist():
- member_path = (workspace_path / member.filename).resolve()
- if (
- workspace_path.resolve() != member_path
- and workspace_path.resolve() not in member_path.parents
- ):
- raise ValueError(
- f"Unsafe path in dataset archive: {member.filename}"
- )
- archive.extractall(workspace_path)
-
- data_yaml_path = workspace_path / "data.yaml"
- if not data_yaml_path.exists():
- raise FileNotFoundError("The downloaded dataset did not contain data.yaml")
-
- dataset_config = yaml.safe_load(data_yaml_path.read_text())
- dataset_config.update(
- {
- "path": str(workspace_path),
- "train": "train/images",
- "val": "valid/images",
- }
- )
- data_yaml_path.write_text(yaml.safe_dump(dataset_config, sort_keys=False))
- return workspace_path, data_yaml_path
-
- return (download_chess_dataset,)
+@app.cell(hide_code=True)
+def _():
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="yolo-wandb-demo", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(mo):
- prepare_data_button = mo.ui.run_button(
- label="Download the chess dataset",
- kind="success",
- tooltip="Downloads public files into a temporary workspace",
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
+ mo.stop(
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- prepare_data_button
- return (prepare_data_button,)
+ try:
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
+ _login_ok = False
-@app.cell(hide_code=True)
-def _(download_chess_dataset, mo, prepare_data_button):
mo.stop(
- not prepare_data_button.value,
- mo.callout(
- mo.md("Click the download button when you are ready."),
- kind="info",
- ),
+ not _login_ok,
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- dataset_root, data_yaml_path = download_chess_dataset()
- mo.callout(
- mo.md(f"The chess dataset is ready at `{dataset_root}`."),
- kind="success",
- )
- return data_yaml_path, dataset_root
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
+
+
+@app.cell
+def _():
+ run_name = "chess-yolo26n"
+ config = {
+ "model": "yolo26n.pt",
+ "epochs": 3,
+ "imgsz": 640,
+ "batch": 16,
+ }
+ yolo_settings.update({"wandb": True})
+ return config, run_name
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
## Detect objects with a pretrained model
- YOLO26 is pretrained on the COCO dataset. Before fine-tuning, run it on a
- sample street image to see the base model's predictions. This step downloads
- the small `yolo26n.pt` checkpoint but does not contact W&B.
+ YOLO26 is pretrained on the COCO dataset, which includes objects such as
+ people, buses, and stop signs. If your application needs those classes, a
+ pretrained model may already be useful without additional training.
+
+ Run the detector on a sample street image. This downloads the small
+ `yolo26n.pt` checkpoint. `conf=0.25` filters out predictions whose confidence
+ is below 0.25. The image shows the remaining boxes, labels, and scores.
""")
return
-@app.cell(hide_code=True)
-def _(mo):
- run_detection_button = mo.ui.run_button(
- label="Run pretrained YOLO26 detection",
- kind="success",
- )
- run_detection_button
- return (run_detection_button,)
-
-
@app.cell
-def _(YOLO, mo, run_detection_button, training_device, yolo_assets):
- mo.stop(
- not run_detection_button.value,
- mo.callout(
- mo.md("Click the detection button above to run inference."),
- kind="info",
- ),
- )
-
- _prediction_model = YOLO("yolo26n.pt")
- _prediction = _prediction_model.predict(
+def _(config):
+ pretrained_model = YOLO(config["model"])
+ prediction = pretrained_model.predict(
source=str(yolo_assets / "bus.jpg"),
- imgsz=640,
+ imgsz=config["imgsz"],
conf=0.25,
- device=training_device,
+ device=device,
save=False,
- verbose=False,
)[0]
- _detection_rows = _prediction.summary()
- _annotated_image = _prediction.plot()[:, :, ::-1]
-
- mo.vstack(
- [
- mo.image(
- _annotated_image,
- width=600,
- ),
- mo.ui.table(_detection_rows, label="Detections"),
- ]
- )
+ # Ultralytics plots use BGR; PIL expects RGB.
+ Image.fromarray(prediction.plot()[:, :, ::-1])
return
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## Connect to Weights & Biases
-
- Enter your [W&B API key](https://wandb.ai/authorize) and, if needed, your
- team or entity. Leave the key blank to use `WANDB_API_KEY` from molab's
- Secrets panel or credentials already configured in the runtime.
-
- Submitting this form authenticates the SDK. It does not create a run.
+ ## Prepare the chess-piece dataset
+
+ Chess pieces are not among the COCO classes. To recognize them, we will
+ fine-tune the pretrained detector on
+ [chessboard images with labeled bounding boxes](https://public.roboflow.com/object-detection/chess-full).
+ Starting from pretrained weights lets us reuse features the model has
+ already learned.
+
+ The download contains images, YOLO-format text labels, and `data.yaml`.
+ The YAML file tells the trainer where the training and validation images
+ are and how class IDs map to names. The helper at the end downloads the
+ dataset into `data/chess` and fixes its paths for this runtime.
""")
return
-@app.cell(hide_code=True)
-def _(mo):
- _api_key_input = mo.ui.text(
- value="",
- kind="password",
- label="W&B API key (blank uses runtime credentials)",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- value="",
- label="W&B entity or team (blank uses your default)",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}")
- .batch(api_key=_api_key_input, entity=_entity_input)
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+@app.cell
+def _():
+ dataset_root, data_yaml_path = download_chess_dataset()
+ print(data_yaml_path.read_text())
+ return data_yaml_path, dataset_root
@app.cell(hide_code=True)
-def _(mo, wandb, wandb_login_form):
- mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Submit the authentication form to continue."),
- kind="info",
- ),
- )
-
- _submitted_login = wandb_login_form.value
- _api_key = _submitted_login["api_key"].strip()
- _requested_entity = _submitted_login["entity"].strip()
- try:
- _login_ok = wandb.login(
- key=_api_key or None,
- relogin=bool(_api_key),
- )
- _resolved_entity = _requested_entity or wandb.Api().default_entity
- _login_error = None
- except wandb.errors.Error as _error:
- _login_ok = False
- _resolved_entity = None
- _login_error = str(_error)
-
- mo.stop(
- not _login_ok or not _resolved_entity,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. Check the API key and "
- "entity, then submit again.\n\n"
- f"W&B reported: `{_login_error or 'No default entity was found.'}`"
- ),
- kind="danger",
- ),
- )
-
- wandb_settings = {"entity": _resolved_entity}
- mo.callout(
- mo.md(f"Connected to W&B as entity `{_resolved_entity}`."),
- kind="success",
- )
- return (wandb_settings,)
+def _():
+ mo.md(r"""
+ ## Fine-tune the detector and track the experiment
+ Weights & Biases helps connect a model's behavior to the data and settings that produced
+ it. In the next cell:
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## Fine-tune YOLO26 and track the run
+ 1. `wandb.init` starts a run and records our configuration.
+ 2. `wandb.Artifact` and `run.log_artifact` version the dataset used by this run.
+ 3. `YOLO.train` fine-tunes the model. The built-in Weights & Biases integration logs
+ losses, validation metrics, prediction images, and diagnostic plots.
+ With `save=True`, it also logs the best checkpoint as a model Artifact.
- The current Ultralytics integration automatically logs training and
- validation metrics, diagnostic images, performance curves, system metrics,
- and the best checkpoint as a model Artifact.
+ The run is closed when training completes or raises an error. Its link
+ appears in the cell output while training is underway.
- Dataset versioning is explicit below: the notebook creates a W&B dataset
- Artifact with `wandb.Artifact` and `run.log_artifact` before training.
- This keeps the W&B workflow visible and reusable with other training
- libraries.
+ The default three epochs demonstrate the workflow. A GPU can reduce training
+ time for longer runs. Training uses CUDA when available and otherwise CPU,
+ including on Macs, to avoid an observed MPS error during loss computation.
+ CPU training can take longer. Local training
+ outputs are stored under `runs/chess`, outside the versioned dataset.
""")
return
-@app.cell(hide_code=True)
-def _(dataset_root, mo, wandb_settings):
- _project_input = mo.ui.text(
- value="yolo-wandb-demo",
- label="W&B project",
- full_width=True,
- )
- _run_name_input = mo.ui.text(
- value="chess-yolo26n",
- label="W&B run name",
- full_width=True,
- )
- _epochs_input = mo.ui.number(
- start=1,
- stop=20,
- value=3,
- step=1,
- label="Training epochs",
- )
- yolo_training_form = (
- mo.md("{project}\n\n{run_name}\n\n{epochs}")
- .batch(
- project=_project_input,
- run_name=_run_name_input,
- epochs=_epochs_input,
- )
- .form(
- submit_button_label="Train YOLO26 and log to W&B",
- bordered=True,
- )
- )
- mo.vstack(
- [
- mo.md(
- f"The dataset is ready at `{dataset_root}`. The run will be "
- f"logged under entity `{wandb_settings['entity']}`."
- ),
- yolo_training_form,
- ]
- )
- return (yolo_training_form,)
-
-
@app.cell
-def _(
- data_yaml_path,
- dataset_root,
- mo,
- training_device,
- wandb,
- wandb_settings,
- yolo_training_form,
-):
- mo.stop(
- yolo_training_form.value is None,
- mo.callout(
- mo.md(
- "Review the settings, then submit the form when you are ready "
- "to create a W&B run and upload the dataset Artifact."
- ),
- kind="info",
- ),
- )
-
- _submitted_training = yolo_training_form.value
- _project = _submitted_training["project"].strip() or "yolo-wandb-demo"
- _run_name = _submitted_training["run_name"].strip() or "chess-yolo26n"
- _epochs = int(_submitted_training["epochs"])
-
- if wandb.run is not None:
- wandb.finish()
-
- _run = wandb.init(
- entity=wandb_settings["entity"],
- project=_project,
- name=_run_name,
+def _(config, data_yaml_path, dataset_root, entity, project, run_name):
+ with wandb.init(
+ entity=entity,
+ project=project,
+ name=run_name,
job_type="training",
- config={
- "model": "yolo26n.pt",
- "dataset": "chess-pieces",
- "epochs": _epochs,
- "device": str(training_device),
- },
- )
- training_run_url = _run.url
-
- _artifact_name = f"{_project}-chess-dataset".replace("/", "-").replace(
- " ", "-"
- )
- _dataset_artifact = wandb.Artifact(
- name=_artifact_name,
- type="dataset",
- description="Chess-piece object-detection dataset in YOLO format",
- metadata={"data_yaml": data_yaml_path.name},
- )
- _dataset_artifact.add_dir(str(dataset_root))
- _run.log_artifact(_dataset_artifact, aliases=["latest"])
-
- yolo_training_request = {
- "epochs": _epochs,
- "run_name": _run_name,
- "local_project_dir": str(dataset_root / "runs"),
- }
- return training_run_url, yolo_training_request
-
-
-@app.cell(hide_code=True)
-def _(mo, training_run_url):
- mo.callout(
- mo.md(
- f"[Open the live W&B run]({training_run_url}) while YOLO26 trains."
- ),
- kind="info",
- )
- training_run_link_ready = True
- return (training_run_link_ready,)
-
+ config=config,
+ ) as run:
+ dataset_artifact = wandb.Artifact(
+ name="chess-pieces",
+ type="dataset",
+ description="Chess-piece images and YOLO bounding-box labels",
+ )
+ dataset_artifact.add_dir(str(dataset_root))
+ run.log_artifact(dataset_artifact, aliases=["latest"])
-@app.cell
-def _(
- Path,
- YOLO,
- data_yaml_path,
- training_device,
- training_run_link_ready,
- wandb,
- yolo_settings,
- yolo_training_request,
-):
- assert training_run_link_ready
- _original_wandb_setting = yolo_settings["wandb"]
- yolo_settings.update({"wandb": True})
- try:
- _training_model = YOLO("yolo26n.pt")
- _training_results = _training_model.train(
+ training_model = YOLO(run.config.model)
+ training_results = training_model.train(
data=str(data_yaml_path),
- epochs=yolo_training_request["epochs"],
- imgsz=640,
+ epochs=run.config.epochs,
+ imgsz=run.config.imgsz,
+ batch=run.config.batch,
device=training_device,
- project=yolo_training_request["local_project_dir"],
- name=yolo_training_request["run_name"],
- exist_ok=True,
+ workers=0,
+ project=str(Path("runs/chess").resolve()),
+ name=run.id,
plots=True,
save=True,
)
- finally:
- if wandb.run is not None:
- wandb.finish()
- yolo_settings.update({"wandb": _original_wandb_setting})
-
- training_output_dir = Path(_training_results.save_dir)
- training_metrics = {
- _name: round(float(_value), 4)
- for _name, _value in _training_results.results_dict.items()
- }
- return training_metrics, training_output_dir
+
+ training_results.results_dict
+ return
@app.cell(hide_code=True)
-def _(mo, training_metrics, training_output_dir, training_run_url):
- mo.vstack(
- [
- mo.callout(
- mo.md(
- "Training finished and the W&B run is closed. "
- f"[Inspect the completed run]({training_run_url})."
- ),
- kind="success",
- ),
- mo.ui.table(
- [
- {"metric": _name, "value": _value}
- for _name, _value in training_metrics.items()
- ],
- label="Final validation metrics",
- ),
- mo.md(f"Local outputs: `{training_output_dir}`"),
- ]
- )
+def _():
+ mo.md(r"""
+ ## Verify and compare your experiments
+
+ Open the Weights & Biases run link and inspect:
+
+ - **Config:** the model, epoch count, image size, and batch size you chose.
+ - **Charts:** training losses, precision, recall, mAP50, and mAP50-95.
+ Losses describe training progress. Mean average precision (mAP) measures detection quality on
+ held-out images, including whether the predicted boxes match the labels.
+ - **Images and plots:** validation predictions, confusion matrix, and
+ performance curves. Look for missed pieces or confused classes.
+ - **Artifacts:** the `chess-pieces` dataset and the best model checkpoint.
+ These record the input data and the model produced by this experiment.
+
+ Increase `epochs` in the configuration, then rerun setup and training.
+ Compare the runs in the same Weights & Biases project. Does validation mAP improve, and
+ do the prediction images show fewer mistakes?
+ """)
return
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## Verify and next steps
+ ## Resume interrupted training
+
+ Ultralytics saves `last.pt` in the run's local `weights` directory. To resume
+ an **unfinished** training job, use its actual checkpoint path:
+
+ ```python
+ model_to_resume = YOLO("runs/chess/
+
+
# Using Keras Evaluation Callbacks with Weights & Biases
[](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/use-wandbevalcallback-in-your-keras-workflow/use_wandbevalcallback_in_your_keras_workflow.py/server)
- This notebook introduces `WandbEvalCallback`, an abstract callback you can subclass to visualize datasets and model predictions using W&B Tables. See the [current Keras integration](https://docs.wandb.ai/models/integrations/keras).
+ This notebook introduces `WandbEvalCallback`, an abstract callback you can subclass to visualize datasets and model predictions using Weights & Biases Tables. See the [Keras integration](https://docs.wandb.ai/models/integrations/keras).
- Use Weights & Biases for machine learning experiment tracking, dataset
- versioning, and project collaboration.
- A CPU is sufficient; a GPU speeds up training. The first
- training submission downloads Fashion-MNIST and ImageNet weights.
+
+ This example can run on CPU. A GPU can reduce training time. Running the data and model
+ cells downloads Fashion-MNIST and ImageNet weights the first time.
+
+ Run the cells from top to bottom in the marimo editor. Editing a cell
+ marks its dependents as stale until you run them again.
""")
return
@@ -46,109 +72,92 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Authentication
-
- Enter your [W&B API key](https://wandb.ai/authorize), or leave it blank to use
- `WANDB_API_KEY` from the marimo Secrets panel or credentials already stored
- in this runtime. A fresh molab session does not inherit your local login.
- If needed, [create a free W&B account](https://wandb.ai/signup).
-
- The entity is the team name in your W&B project URL:
- `wandb.ai/
+
-
-
+ # Using WandbMetricsLogger in your Keras workflow
- Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/use-wandbmetriclogger-in-your-keras-workflow/use_wandbmetriclogger_in_your_keras_workflow.py/server)
-
+
-
-
+ # Using WandbModelCheckpoint in your Keras workflow
+
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/use-wandbmodelcheckpoint-in-your-keras-workflow/use_wandbmodelcheckpoint_in_your_keras_workflow.py/server)
- Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+ This notebook introduces the `WandbModelCheckpoint` callback. It saves Keras
+ model checkpoints and logs them as versioned Weights & Biases [Artifacts](https://docs.wandb.ai/models/artifacts).
+ Use `WandbMetricsLogger` alongside it to track training and validation metrics.
-
+
- Squeezing the best performance out of tree-based models requires
- [selecting the right hyperparameters](https://blog.cambridgespark.com/hyperparameter-tuning-in-xgboost-4ff9100a3b2f).
- How many `early_stopping_rounds`? What should the `max_depth` of a tree be?
+ # Using Weights & Biases Sweeps with XGBoost
- Searching through high dimensional hyperparameter spaces to find the most performant model can get unwieldy very fast.
- Hyperparameter sweeps provide an organized and efficient way to conduct a battle royale of models and crown a winner.
- They enable this by automatically searching through combinations of hyperparameter values to find the most optimal values.
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/using-w-b-sweeps-with-xgboost/using_w_b_sweeps_with_xgboost.py/server)
- In this tutorial we'll see how you can run sophisticated hyperparameter sweeps on XGBoost models in 3 easy steps using Weights and Biases.
+ Hyperparameter sweeps search combinations of model settings and record
+ their performance. In this lesson, compare XGBoost's tree and linear
+ boosters while exploring learning rate, tree depth, and row sampling.
- For a teaser, check out the plots below:
- """)
- return
+ There are three steps:
+ 1. **Define the sweep:** choose the search strategy, parameters, and metric.
+ 2. **Initialize the sweep:** register the configuration with `wandb.sweep`.
+ 3. **Run an agent:** call `wandb.agent` with a training function and run count.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- 
+ This notebook runs on a CPU. On macOS, install the OpenMP runtime with
+ `brew install libomp` before starting the notebook if it is not installed.
""")
return
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## Sweeps: An Overview
+ ## Authenticate with Weights & Biases
- Running a hyperparameter sweep with Weights & Biases is very easy. There are just 3 simple steps:
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- 1. **Define the sweep:** we do this by creating a dictionary-like object that specifies the sweep: which parameters to search through, which search strategy to use, which metric to optimize.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- 2. **Initialize the sweep:** with one line of code we initialize the sweep and pass in the dictionary of sweep configurations:
- `sweep_id = wandb.sweep(sweep_config)`
+ ## Set your team and project
- 3. **Run the sweep agent:** also accomplished with one line of code, we call `wandb.agent()` and pass the `sweep_id` along with a function that defines your model architecture and trains it:
- `wandb.agent(sweep_id, function=train)`
-
- And voila! That's all there is to running a hyperparameter sweep!
-
- In the notebook below, we'll walk through these 3 steps in more detail.
-
- We highly encourage you to fork this notebook, tweak the parameters, or try the model with your own dataset!
-
- ### Resources
- - [Sweeps docs →](https://docs.wandb.com/library/sweeps)
- - [Launching from the command line →](https://www.wandb.com/articles/hyperparameter-tuning-as-easy-as-1-2-3)
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `XGBoost-sweeps`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- mo.md(r"""
- ## Authentication
-
- To run this tutorial, authenticate with W&B and choose where to create the
- sweep. Get an API key from [wandb.ai/authorize](https://wandb.ai/authorize),
- or leave the key blank to use `WANDB_API_KEY` from marimo's Secrets panel,
- an existing W&B settings file, or a prior `wandb login` in this runtime.
-
- If you work in a team, enter its entity slug. You can find it in the URL of
- your team's W&B workspace; leave it blank to use your default entity.
-
- When running locally on macOS, XGBoost also needs the OpenMP runtime. Install
- it once with `brew install libomp`. Hosted Linux environments generally
- provide the required runtime with the XGBoost wheel.
-
- Submitting the form creates a new Sweep Controller and then starts the
- requested number of W&B runs. Editing fields without submitting does not
- create anything in W&B. Editing the `sweep_config` or `train` code below
- resets the form so you can review and submit those changes explicitly.
- Submitting again creates and runs a new sweep.
- """)
- return
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="XGBoost-sweeps", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(sweep_config, train):
- # Check the native XGBoost runtime before creating a Sweep Controller.
- # On macOS, the XGBoost wheel requires libomp from Homebrew.
- try:
- from xgboost import XGBClassifier as _XGBClassifier # noqa: F401
- except (ImportError, OSError, ValueError) as _error:
- _xgboost_error = str(_error)
- else:
- _xgboost_error = None
-
- _macos_help = (
- " On macOS, run `brew install libomp`, then restart this notebook."
- if "libomp.dylib" in (_xgboost_error or "")
- else ""
- )
- mo.stop(
- _xgboost_error is not None,
- mo.callout(
- mo.md(
- "**XGBoost could not load its native runtime.**"
- f"{_macos_help}\n\n"
- "The sweep has not been created, so fixing the runtime and "
- "resubmitting will not leave failed runs behind."
- ),
- kind="danger",
- ),
- )
-
- _api_key_input = mo.ui.text(
- value="",
- kind="password",
- label="W&B API key (blank uses runtime credentials)",
- )
- _entity_input = mo.ui.text(
- value="",
- label="W&B entity (blank uses your default)",
- )
- _project_input = mo.ui.text(value="XGBoost-sweeps", label="W&B project")
- _run_count_input = mo.ui.number(
- start=1,
- stop=50,
- value=25,
- step=1,
- label="Number of sweep runs",
- )
-
- sweep_form = (
- mo.md(
- """
- {api_key}
-
- {entity}
-
- {project}
-
- {run_count}
- """
- )
- .batch(
- api_key=_api_key_input,
- entity=_entity_input,
- project=_project_input,
- run_count=_run_count_input,
- )
- .form(
- submit_button_label="Create sweep and run agent",
- bordered=False,
- )
- )
-
- mo.vstack(
- [
- mo.md(
- f"Ready to submit a `{sweep_config['method']}` sweep using "
- f"the `{train.__name__}` training function."
- ),
- sweep_form,
- ]
- )
- return (sweep_form,)
-
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
-@app.cell(hide_code=True)
-def _(sweep_form):
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- sweep_form.value is None,
- mo.md(
- "Review the sweep configuration below, then submit the form above "
- "when you are ready to create the Sweep Controller and launch its runs."
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _submitted = sweep_form.value
- _api_key = _submitted["api_key"].strip()
try:
- _login_ok = wandb.login(
- key=_api_key or None,
- relogin=bool(_api_key),
- )
- _login_error = None
- except wandb.errors.Error as _error:
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
_login_ok = False
- _login_error = str(_error)
mo.stop(
not _login_ok,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. "
- f"Check the API key and try again.\n\nW&B reported: `{_login_error or 'unknown error'}`"
- ),
- kind="danger",
- ),
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- sweep_request = {
- "entity": _submitted["entity"].strip() or None,
- "project": _submitted["project"].strip() or "XGBoost-sweeps",
- "run_count": int(_submitted["run_count"]),
- }
- mo.callout(
- mo.md("Connected to W&B. The submitted sweep will start now."),
- kind="success",
- )
- return (sweep_request,)
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## 1. Define the Sweep
+ ## Define the sweep
- Weights & Biases sweeps give you powerful levers to configure your sweeps exactly how you want them, with just a few lines of code. The sweeps config can be defined as
- [a dictionary or a YAML file](https://docs.wandb.ai/guides/sweeps/configuration).
+ A sweep configuration can be a [Python dictionary or YAML file](https://docs.wandb.ai/models/sweeps/define-sweep-configuration).
+ Here we use a dictionary with three main fields:
- Let's walk through some of them together:
- * **Metric** – This is the metric the sweeps are attempting to optimize. Metrics can take a `name` (this metric should be logged by your training script) and a `goal` (`maximize` or `minimize`).
- * **Search Strategy** – Specified using the `"method"` key. We support several different search strategies with sweeps.
- * **Grid Search** – Iterates over every combination of hyperparameter values.
- * **Random Search** – Iterates over randomly chosen combinations of hyperparameter values.
- * **Bayesian Search** – Creates a probabilistic model that maps hyperparameters to probability of a metric score, and chooses parameters with high probability of improving the metric. The objective of Bayesian optimization is to spend more time in picking the hyperparameter values, but in doing so trying out fewer hyperparameter values.
- * **Parameters** – A dictionary containing the hyperparameter names, and discrete values, a range, or distributions from which to pull their values on each iteration.
+ - **Metric:** the logged metric to optimize and whether to maximize or
+ minimize it. Our training function logs validation accuracy as `accuracy`.
+ - **Method:** `random` samples combinations, `grid` tries every combination,
+ and `bayes` uses previous results to choose promising settings.
+ - **Parameters:** the values, ranges, or distributions to search.
- You can find a list of all configuration options [here](https://docs.wandb.com/library/sweeps/configuration).
+ Both boosters use `learning_rate`. The [tree parameters](https://xgboost.readthedocs.io/en/stable/parameter.html#parameters-for-tree-booster)
+ `max_depth` and `subsample` affect only `gbtree`, so the training function
+ applies them only to tree models. Their sampled values are still recorded
+ in every run’s config. Filter to tree runs when comparing those parameters.
""")
return
@@ -300,256 +173,164 @@ def _():
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ## 2. Initialize the Sweep
- """)
- return
+ ## Initialize the sweep
+ `wandb.sweep` registers this configuration with Weights & Biases and returns a sweep ID.
+ The sweep service supplies parameter choices to agents and tracks the
+ metrics their runs report.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Calling `wandb.sweep` starts a Sweep Controller --
- a centralized process that provides settings of the `parameters` to any who query it
- and expects them to return performance on `metrics` via `wandb` logging.
+ Running this cell creates a new sweep. To add trials to the same sweep,
+ rerun only the agent cell below with the existing `sweep_id`.
""")
return
@app.cell
-def _(sweep_config, sweep_request):
- sweep_id = wandb.sweep(
- sweep=sweep_config,
- entity=sweep_request["entity"],
- project=sweep_request["project"],
- )
+def _(entity, project, sweep_config):
+ sweep_id = wandb.sweep(sweep=sweep_config, entity=entity, project=project)
return (sweep_id,)
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- ### Define your training process
- Before we can run the sweep,
- we need to define a function that creates and trains the model --
- the function that takes in hyperparameter values and spits out metrics.
-
- We'll also need `wandb` to be integrated into our script.
- There's three main components:
- * `wandb.init()` – Initialize a new W&B run. Each run is single execution of the training script.
- * `wandb.config` – Save all your hyperparameters in a config object. This lets you use [our app](https://wandb.ai) to sort and compare your runs by hyperparameter values.
- * `wandb.log()` – Logs metrics and custom objects – these can be images, videos, audio files, HTML, plots, point clouds etc.
- """)
- return
+ ### Prepare the dataset
+ Download the Pima Indians Diabetes CSV once and reuse it across trials.
+ The first eight columns are features and the last column is the binary
+ target. Each trial uses the same split seed and held-out fraction so its
+ accuracy can be compared with other trials.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- We also need the Pima Indians Diabetes dataset. Each sweep run reads this
- small public CSV directly from its source, so the notebook does not rely on
- a repository checkout or leave a downloaded data file behind.
+ Because the sweep uses this held-out split to select hyperparameters, it
+ serves as validation data. Use a separate test set for a final evaluation.
""")
return
-@app.function
-def train():
- # XGBoost model for Pima Indians dataset
- from xgboost import XGBClassifier
+@app.cell
+def _():
+ with urlopen(DATASET_URL, timeout=30) as dataset_response:
+ dataset = loadtxt(dataset_response, delimiter=",")
+ X, y = dataset[:, :8], dataset[:, 8]
+ return X, y
- if wandb.run is not None:
- wandb.finish()
- config_defaults = {
- "booster": "gbtree",
- "max_depth": 3,
- "learning_rate": 0.1,
- "subsample": 1,
- "seed": 117,
- "test_size": 0.33,
- }
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ### Define the training function
- # The sweep agent overrides these defaults with each sampled configuration.
- with wandb.init(config=config_defaults) as run:
- config = run.config
-
- # Load data and split it into predictors and targets.
- with urlopen(DATASET_URL, timeout=30) as dataset_response:
- dataset = loadtxt(dataset_response, delimiter=",")
- X, y = dataset[:, :8], dataset[:, 8]
-
- # Split data into train and test sets.
- X_train, X_test, y_train, y_test = train_test_split(
- X,
- y,
- test_size=config.test_size,
- random_state=config.seed,
- )
-
- # Fit the model on the training data.
- model = XGBClassifier(
- booster=config.booster,
- max_depth=config.max_depth,
- learning_rate=config.learning_rate,
- subsample=config.subsample,
- )
- model.fit(X_train, y_train)
-
- # Make predictions on the test set.
- y_pred = model.predict(X_test)
- predictions = [round(value) for value in y_pred]
-
- # Evaluate and log the result.
- accuracy = accuracy_score(y_test, predictions)
- print(f"Accuracy: {accuracy:.0%}")
- run.log({"accuracy": accuracy})
+ The agent calls `train()` once per trial. Inside the function:
+ - `wandb.init(config=config_defaults)` starts a run. Values supplied by the
+ sweep override the defaults.
+ - `run.config` provides that trial's hyperparameters.
+ - `run.log({"accuracy": accuracy})` reports the metric named in the sweep
+ configuration.
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ## 3. Run the Sweep with an agent
+ The context manager finishes each run before the next trial begins.
""")
return
+@app.cell
+def _(X, y):
+ def train():
+ config_defaults = {
+ "booster": "gbtree",
+ "max_depth": 3,
+ "learning_rate": 0.1,
+ "subsample": 1,
+ "seed": 117,
+ "test_size": 0.33,
+ }
+
+ with wandb.init(config=config_defaults) as run:
+ config = run.config
+ X_train, X_valid, y_train, y_valid = train_test_split(
+ X,
+ y,
+ test_size=config.test_size,
+ random_state=config.seed,
+ )
+
+ model_params = {
+ "booster": config.booster,
+ "learning_rate": config.learning_rate,
+ "random_state": config.seed,
+ "n_jobs": 2,
+ }
+ if config.booster == "gbtree":
+ model_params.update(
+ max_depth=config.max_depth,
+ subsample=config.subsample,
+ )
+ model = XGBClassifier(**model_params)
+ model.fit(X_train, y_train)
+
+ predictions = model.predict(X_valid)
+ accuracy = accuracy_score(y_valid, predictions)
+ print(f"Validation accuracy: {accuracy:.1%}")
+ run.log({"accuracy": accuracy})
+
+ return (train,)
+
+
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
- Now, we call `wandb.agent` to start up our sweep.
+ ## Run the sweep agent
- You can call `wandb.agent` on any machine where you're logged into W&B that has
- - the `sweep_id`,
- - the dataset and `train` function
+ The agent requests hyperparameters from Weights & Biases and calls `train` for each trial.
+ `count` limits how many runs this agent launches. Without a limit or a
+ manual stop, a random-search agent can keep requesting trials. Start with a smaller `run_count` for a quick check.
- and that machine will join the sweep!
-
- > _Note_: a `random` sweep will by default run forever,
- trying new parameter combinations until the cows come home --
- or until you [turn the sweep off from the app UI](https://docs.wandb.ai/ref/app/features/sweeps).
- You can prevent this by providing the total `count` of runs you'd like the `agent` to complete.
+ Agents on other machines can join the same sweep using its ID, entity,
+ project, and the same training function and dataset. The count applies to
+ each agent, rather than the whole sweep.
""")
return
@app.cell
-def _(sweep_id, sweep_request):
+def _(entity, project, sweep_id, train):
+ run_count = 25
wandb.agent(
sweep_id,
function=train,
- count=sweep_request["run_count"],
- )
- sweep_complete = {
- "entity": sweep_request["entity"],
- "project": sweep_request["project"],
- "run_count": sweep_request["run_count"],
- "sweep_id": sweep_id,
- }
- return (sweep_complete,)
-
-
-@app.cell(hide_code=True)
-def _(sweep_complete):
- mo.callout(
- mo.md(
- f"**Sweep complete.** Agent `{sweep_complete['sweep_id']}` finished "
- f"{sweep_complete['run_count']} runs in project "
- f"`{sweep_complete['project']}`."
- ),
- kind="success",
+ count=run_count,
+ entity=entity,
+ project=project,
)
return
@app.cell(hide_code=True)
-def _(mo):
+def _():
mo.md(r"""
## Visualize your results
- """)
- return
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- Now that your sweep is finished, it's time to look at the results.
+ Click on the link returned by Weights & Biases in the previous cell. The link will look similar to the following:
- Weights & Biases will generate a number of useful plots for you automatically.
- """)
- return
+ ```bash
+ 🧹 View sweep at https://wandb.ai/wandb/XGBoost-sweeps/sweeps/wna9genj
+ ```
+ Navigate to the **Sweep** panel. By default, Weights & Biases plots [parallel coordinates](https://docs.coreweave.com/products/wandb/app/features/panels/parallel-coordinates#parallel-coordinates), [scatter plot, and a parameter importance plot](https://docs.coreweave.com/products/wandb/sweeps/visualize-sweep-results#scatter-plot-and-parameter-importance).
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### Parallel coordinates plot
+ ## Verify and next steps
- This plot maps hyperparameter values to model metrics. It’s useful for honing in on combinations of hyperparameters that led to the best model performance.
+ Check that the agent's runs finished and that each run has an `accuracy`
+ value and the sampled parameters in **Config**. Inspect the parallel
+ coordinates and parameter importance panels, filtering to tree runs when
+ evaluating `max_depth` and `subsample`.
- This plot seems to indicate that using a tree as our learner slightly,
- but not mind-blowingly,
- outperforms using a simple linear model as our learner.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- ### Hyperparameter importance plot
-
- The hyperparameter importance plot shows which hyperparameter values had the biggest impact
- on your metrics.
-
- We report both the correlation (treating it as a linear predictor)
- and the feature importance (after training a random forest on your results)
- so you can see which parameters had the biggest effect
- and whether that effect was positive or negative.
-
- Reading this chart, we see quantitative confirmation
- of the trend we noticed in the parallel coordinates chart above:
- the largest impact on validation accuracy came from the choice of
- learner, and the `gblinear` learners were generally worse than `gbtree` learners.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(mo):
- mo.md(r"""
- 
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- These visualizations can help you save both time and resources running expensive hyperparameter optimizations by honing in on the parameters (and value ranges) that are the most important, and thereby worthy of further exploration.
- """)
- return
-
-
-@app.cell(hide_code=True)
-def _(sweep_complete):
- mo.md(f"""
- ## Verify your sweep
-
- In W&B, open project `{sweep_complete["project"]}` and select the sweep with
- ID `{sweep_complete["sweep_id"]}`. Confirm that it contains
- {sweep_complete["run_count"]} completed runs, then inspect the parallel
- coordinates and hyperparameter importance panels to see which settings had
- the strongest relationship with `accuracy`.
+ Try narrowing the learning-rate values or changing `method` to `bayes`.
+ Run the sweep-creation cell again, then the agent cell, to compare the new
+ sweep with the first one. See the [Sweeps documentation](https://docs.wandb.ai/models/sweeps)
+ for more configuration options.
""")
return
diff --git a/marimo/convert/w-b-tables-quickstart/w_b_tables_quickstart.py b/marimo/convert/w-b-tables-quickstart/w_b_tables_quickstart.py
index 70d6e3ca..b8b1cc5d 100644
--- a/marimo/convert/w-b-tables-quickstart/w_b_tables_quickstart.py
+++ b/marimo/convert/w-b-tables-quickstart/w_b_tables_quickstart.py
@@ -1,5 +1,7 @@
# /// script
+# requires-python = ">=3.10"
# dependencies = [
+# "marimo>=0.24",
# "torch==2.14.0",
# "torchvision==0.29.0",
# "wandb==0.30.0",
@@ -8,8 +10,11 @@
import marimo
-__generated_with = "0.24.0"
-app = marimo.App(auto_download=["html"])
+__generated_with = "0.25.1"
+app = marimo.App(
+ width="medium",
+ app_title="Weights & Biases Tables Quickstart",
+)
with app.setup:
import marimo as mo
@@ -21,18 +26,6 @@
import torchvision.transforms as T
import wandb
- WANDB_PROJECT = "table-quickstart"
-
-
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- # View & analyze model predictions during training
-
- [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/w-b-tables-quickstart/w_b_tables_quickstart.py/server)
- """)
- return
-
@app.cell(hide_code=True)
def _():
@@ -54,32 +47,16 @@ def _():
- This quickstart guide covers how to track, visualize, and compare model predictions over the course of training, using PyTorch on MNIST data.
-
- With [W&B Tables](https://docs.wandb.com/datasets-and-predictions):
- 1. Log metrics, images, text, etc. to a `wandb.Table()` during model training or evaluation
- 2. View, sort, filter, group, join, interactively query, and explore these tables
- 3. Compare model predictions or results: dynamically across specific images, hyperparameters/model versions, or time steps.
-
- ## Examples
- ### Compare predicted scores for specific images
-
- [Live example: compare predictions after 1 vs 5 epochs of training →](https://wandb.ai/stacey/table-quickstart/reports/CNN-2-Progress-over-Training-Time--Vmlldzo3NDY5ODU#compare-predictions-after-1-vs-5-epochs)
- 
- The histograms compare per-class scores between the two models. The top green bar in each histogram represents model "CNN-2, 1 epoch" (id 0), which only trained for 1 epoch. The bottom purple bar represents model "CNN-2, 5 epochs" (id 1), which trained for 5 epochs. The images are filtered to cases where the models disagree. For example, in the first row, the "4" gets high scores across all the possible digits after 1 epoch, but after 5 epochs it scores highest on the correct label and very low on the rest.
-
- ### Focus on top errors over time
- [Live example →](https://wandb.ai/stacey/table-quickstart/reports/CNN-2-Progress-over-Training-Time--Vmlldzo3NDY5ODU#top-errors-over-time)
-
- See incorrect predictions (filter to rows where "guess" != "truth") on the full test data. Note that there are 229 wrong guesses after 1 training epoch, but only 98 after 5 epochs.
- 
+ # View & analyze model predictions during training
- ### Compare model performance and find patterns
+ [](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/w-b-tables-quickstart/w_b_tables_quickstart.py/server)
- [See full detail in a live example →](https://wandb.ai/stacey/table-quickstart/reports/CNN-2-Progress-over-Training-Time--Vmlldzo3NDY5ODU#false-positives-grouped-by-guess)
+ Train an MNIST classifier and log a Table of labeled prediction examples
+ after each epoch. Use the images, predictions, and class probabilities to
+ inspect mistakes that an aggregate accuracy score cannot explain.
- Filter out correct answers, then group by the guess to see examples of misclassified images and the underlying distribution of true labels—for two models side-by-side. A model variant with 2X the layer sizes and learning rate is on the left, and the baseline is on the right. Note that the baseline makes slightly more mistakes for each guessed class.
- 
+ Running the following code cells downloads MNIST, trains a model, and logs
+ metrics and prediction tables to Weights & Biases.
""")
return
@@ -87,82 +64,71 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Authentication
+ ## Authenticate with Weights & Biases
- Enter your [W&B API key](https://wandb.ai/authorize) and, if needed, your team or entity. You can leave the key blank when this environment already has W&B credentials, including a `WANDB_API_KEY` configured in the Molab Secrets panel.
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
- If you do not yet have an account, [sign up for W&B](https://wandb.ai/signup). Your entity is the first path segment after `wandb.ai` when you open a W&B workspace.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
- This example runs in Molab as a convenient hosted environment, but you can run the same training code anywhere and visualize its metrics and predictions in W&B.
+ ## Set your team and project
+
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `table-quickstart`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- _api_key_input = mo.ui.text(
- kind="password",
- label="W&B API key (optional)",
- placeholder="Paste a key or use cached credentials",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- label="W&B entity or team (optional)",
- placeholder="Leave blank to use your default entity",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}")
- .batch(api_key=_api_key_input, entity=_entity_input)
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="table-quickstart", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(wandb_login_form):
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- "Submit the form above to authenticate before creating a run.",
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _api_key = wandb_login_form.value["api_key"].strip()
- _entity = wandb_login_form.value["entity"].strip()
-
try:
- wandb.login(key=_api_key or None, relogin=bool(_api_key))
- except wandb.errors.Error:
- mo.stop(
- True,
- mo.callout(
- "W&B authentication failed. Check the API key or your Molab Secrets configuration, then submit the form again.",
- kind="danger",
- ),
- )
-
- wandb_entity = _entity or None
- _entity_note = (
- f" Runs will be logged to `{wandb_entity}`."
- if wandb_entity
- else " Runs will use your default W&B entity."
- )
- mo.callout(
- mo.md(f"W&B is connected.{_entity_note}"),
- kind="success",
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
+ _login_ok = False
+
+ mo.stop(
+ not _login_ok,
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- return (wandb_entity,)
+
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## 0. Setup
+ ## Set the Training and Table Sizes
- The dependencies are declared in the notebook metadata. The next cells define the PyTorch data loader, model, and training schedule. MNIST is downloaded only after you explicitly start the training workflow.
+ The next cell selects a GPU when available and defines the model's
+ hyperparameters. One epoch keeps the demo short. Increase `EPOCHS` to compare
+ predictions at different points in training.
+
+ We evaluate all 10,000 test images, but log only 32 images from each of the
+ first ten test batches. That creates a table with 320 rows per epoch and
+ keeps image uploads small.
""")
return
@@ -170,72 +136,57 @@ def _():
@app.cell
def _():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
- # Create train and test data loaders.
- def get_dataloader(is_train, batch_size, slice=5):
- "Get a training data loader."
- ds = torchvision.datasets.MNIST(
- root=".",
- train=is_train,
- transform=T.ToTensor(),
- download=True,
- )
- loader = torch.utils.data.DataLoader(
- dataset=ds,
- batch_size=batch_size,
- shuffle=True if is_train else False,
- pin_memory=torch.cuda.is_available(),
- num_workers=2,
- )
- return loader
-
- return device, get_dataloader
+ EPOCHS = 1
+ BATCH_SIZE = 32
+ LEARNING_RATE = 0.001
+ NUM_CLASSES = 10
+ L1_SIZE = 32
+ L2_SIZE = 64
+ CONV_KERNEL_SIZE = 5
+ NUM_BATCHES_TO_LOG = 10
+ NUM_IMAGES_PER_BATCH = 32
+ return (
+ BATCH_SIZE,
+ CONV_KERNEL_SIZE,
+ EPOCHS,
+ L1_SIZE,
+ L2_SIZE,
+ LEARNING_RATE,
+ NUM_BATCHES_TO_LOG,
+ NUM_CLASSES,
+ NUM_IMAGES_PER_BATCH,
+ device,
+ )
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## 1. Define the model and training schedule
+ ## Define the Classifier
+
+ The next cell defines a small convolutional network for MNIST. It returns
+ ten logits per image, one for each digit. During evaluation, we apply
+ softmax to convert those logits into per-class probabilities for the table.
- * Set the number of epochs to run, where each epoch consists of a training step and a validation (test) step. Optionally configure the amount of data to log per test step. Here the number of batches and number of images per batch to visualize are set low to simplify the demo.
- * Define a simple convolutional neural net (following [pytorch-tutorial](https://github.com/yunjey/pytorch-tutorial) code).
- * Load in train and test sets using PyTorch
+ This is ordinary PyTorch model code. The Weights & Biases integration happens later,
+ when we start a run and log its metrics and predictions.
""")
return
@app.cell
-def _(device):
- # Number of epochs to run. Each epoch includes a training step and a
- # test step, so this sets the number of prediction tables to log.
- EPOCHS = 1
-
- # Keep the amount of prediction data low to simplify the demo.
- NUM_BATCHES_TO_LOG = 10
- NUM_IMAGES_PER_BATCH = 32
-
- # Training configuration and hyperparameters.
- NUM_CLASSES = 10
- BATCH_SIZE = 32
- LEARNING_RATE = 0.001
- L1_SIZE = 32
- L2_SIZE = 64
- # Changing this may require changing the shape of adjacent layers.
- CONV_KERNEL_SIZE = 5
-
-
- # Define a two-layer convolutional neural network.
+def _(CONV_KERNEL_SIZE, L1_SIZE, L2_SIZE):
class ConvNet(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.layer1 = nn.Sequential(
- nn.Conv2d(1, L1_SIZE, CONV_KERNEL_SIZE, stride=1, padding=2),
+ nn.Conv2d(1, L1_SIZE, CONV_KERNEL_SIZE, stride=1, padding=CONV_KERNEL_SIZE // 2),
nn.BatchNorm2d(L1_SIZE),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.layer2 = nn.Sequential(
- nn.Conv2d(L1_SIZE, L2_SIZE, CONV_KERNEL_SIZE, stride=1, padding=2),
+ nn.Conv2d(L1_SIZE, L2_SIZE, CONV_KERNEL_SIZE, stride=1, padding=CONV_KERNEL_SIZE // 2),
nn.BatchNorm2d(L2_SIZE),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
@@ -248,64 +199,80 @@ def forward(self, x):
out = out.reshape(out.size(0), -1)
return self.fc(out)
- model = ConvNet(NUM_CLASSES).to(device)
- criterion = nn.CrossEntropyLoss()
- optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
-
- model
- return (
- BATCH_SIZE,
- CONV_KERNEL_SIZE,
- EPOCHS,
- L1_SIZE,
- L2_SIZE,
- LEARNING_RATE,
- NUM_BATCHES_TO_LOG,
- NUM_CLASSES,
- NUM_IMAGES_PER_BATCH,
- criterion,
- model,
- optimizer,
- )
+ return (ConvNet,)
@app.cell(hide_code=True)
def _():
- mo.vstack(
- [
- mo.md(r"""
- ## 2. Run training and log test predictions
-
- For every epoch, run a training step and a test step. For each test step, create a `wandb.Table()` in which to store test predictions. These can be visualized, dynamically queried, and compared side by side in your browser.
-
- The button below downloads MNIST, trains the model for the configured number of epochs, and creates a W&B run containing metrics and prediction tables.
- """),
- mo.callout(
- mo.md(
- "For faster training in molab, open **Configure compute** and select a GPU before starting the run."
- ),
- kind="info",
- ),
- ]
- )
+ mo.md(r"""
+ ## Load MNIST
+
+ The next cell downloads MNIST and creates training and test data loaders.
+ Training batches are shuffled. Test batches stay in a fixed order so each
+ logged image has the same identifier across epochs and experiments.
+
+ The test batch size is twice the training batch size. We use
+ `num_workers=0` so loading also works in notebook runtimes without separate
+ worker processes.
+ """)
return
@app.cell
-def _():
- run_training = mo.ui.run_button(
- label="Train model and log predictions to W&B",
- kind="success",
- tooltip="Downloads MNIST and creates a W&B run",
+def _(BATCH_SIZE):
+ train_dataset = torchvision.datasets.MNIST(
+ root="./data", train=True, transform=T.ToTensor(), download=True,
+ )
+ test_dataset = torchvision.datasets.MNIST(
+ root="./data", train=False, transform=T.ToTensor(), download=True,
)
- run_training
- return (run_training,)
+ train_loader = torch.utils.data.DataLoader(
+ train_dataset,
+ batch_size=BATCH_SIZE,
+ shuffle=True,
+ num_workers=0,
+ pin_memory=torch.cuda.is_available(),
+ )
+ test_loader = torch.utils.data.DataLoader(
+ test_dataset,
+ batch_size=2 * BATCH_SIZE,
+ shuffle=False,
+ num_workers=0,
+ pin_memory=torch.cuda.is_available(),
+ )
+ return test_loader, train_loader
+
+
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Train and Log Prediction Tables
+
+ The next cell starts a Weights & Biases run with `wandb.init()` and records the training
+ configuration. Each execution builds a fresh model and optimizer.
+
+ During training, `training_run.log()` sends the loss to Weights & Biases. At the end of
+ each epoch, we evaluate the model and build a new `wandb.Table`:
+
+ - `columns` defines the schema: image ID, image, predicted digit, true digit,
+ and ten class probabilities.
+ - `add_data()` adds one row per image. `wandb.Image()` makes the image
+ viewable in Weights & Biases. The remaining values let you filter errors and compare
+ class scores.
+ - `training_run.log()` uploads accuracy and the table under
+ `test_predictions`. Logging a fresh table under the same key each epoch
+ lets you inspect predictions over time.
+
+ The context manager finishes the run after training and evaluation complete.
+ """)
+ return
@app.cell
def _(
BATCH_SIZE,
CONV_KERNEL_SIZE,
+ ConvNet,
EPOCHS,
L1_SIZE,
L2_SIZE,
@@ -313,27 +280,12 @@ def _(
NUM_BATCHES_TO_LOG,
NUM_CLASSES,
NUM_IMAGES_PER_BATCH,
- criterion,
device,
- get_dataloader,
- model,
- optimizer,
- run_training,
- wandb_entity,
+ entity,
+ project,
+ test_loader,
+ train_loader,
):
- mo.stop(
- not run_training.value,
- mo.callout(
- mo.md(
- "Click **Train model and log predictions to W&B** when you are ready to download MNIST and start the run."
- ),
- kind="info",
- ),
- )
-
- train_loader = get_dataloader(is_train=True, batch_size=BATCH_SIZE)
- test_loader = get_dataloader(is_train=False, batch_size=2 * BATCH_SIZE)
-
training_config = {
"epochs": EPOCHS,
"batch_size": BATCH_SIZE,
@@ -341,47 +293,15 @@ def _(
"l1_size": L1_SIZE,
"l2_size": L2_SIZE,
"conv_kernel": CONV_KERNEL_SIZE,
- "img_count": min(10000, NUM_IMAGES_PER_BATCH * NUM_BATCHES_TO_LOG),
+ "num_batches_to_log": NUM_BATCHES_TO_LOG,
+ "num_images_per_batch": NUM_IMAGES_PER_BATCH,
}
- # ✨ W&B: Initialize a run and log the training configuration.
- with wandb.init(
- project=WANDB_PROJECT,
- entity=wandb_entity,
- config=training_config,
- ) as run:
- run_url = run.url
-
- def log_test_predictions(
- images,
- labels,
- outputs,
- predicted,
- test_table,
- log_counter,
- ):
- "Log predictions and per-class scores for one test batch."
- scores = F.softmax(outputs.data, dim=1)
- log_scores = scores.cpu().numpy()
- log_images = images.cpu().numpy()
- log_labels = labels.cpu().numpy()
- log_preds = predicted.cpu().numpy()
-
- for image_id, (image, label, prediction, score) in enumerate(
- zip(log_images, log_labels, log_preds, log_scores)
- ):
- row_id = f"{image_id}_{log_counter}"
- test_table.add_data(
- row_id,
- wandb.Image(image),
- prediction,
- label,
- *score,
- )
- if image_id + 1 == NUM_IMAGES_PER_BATCH:
- break
-
- total_step = len(train_loader)
+ model = ConvNet(NUM_CLASSES).to(device)
+ criterion = nn.CrossEntropyLoss()
+ optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+ with wandb.init(project=project, entity=entity, config=training_config) as training_run:
for epoch in range(EPOCHS):
model.train()
for step, (images, labels) in enumerate(train_loader):
@@ -394,17 +314,11 @@ def log_test_predictions(
loss.backward()
optimizer.step()
- # ✨ W&B: Log loss over training steps for live visualization.
- run.log({"loss": loss.item()})
+ training_run.log({"loss": loss.item()})
if (step + 1) % 100 == 0:
print(
- "Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}".format(
- epoch + 1,
- EPOCHS,
- step + 1,
- total_step,
- loss.item(),
- )
+ f"Epoch {epoch + 1}/{EPOCHS}, "
+ f"batch {step + 1}/{len(train_loader)}, loss: {loss.item():.4f}"
)
columns = ["id", "image", "guess", "truth"]
@@ -412,68 +326,74 @@ def log_test_predictions(
test_table = wandb.Table(columns=columns)
model.eval()
- log_counter = 0
+ correct = 0
+ total = 0
with torch.no_grad():
- correct = 0
- total = 0
- for images, labels in test_loader:
+ for batch_index, (images, labels) in enumerate(test_loader):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
- _, predicted = torch.max(outputs.data, 1)
-
- if log_counter < NUM_BATCHES_TO_LOG:
- log_test_predictions(
- images,
- labels,
- outputs,
- predicted,
- test_table,
- log_counter,
- )
- log_counter += 1
+ predicted = outputs.argmax(dim=1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
+ if batch_index < NUM_BATCHES_TO_LOG:
+ scores = F.softmax(outputs, dim=1).cpu()
+ for image_index in range(min(NUM_IMAGES_PER_BATCH, len(images))):
+ image_id = batch_index * test_loader.batch_size + image_index
+ test_table.add_data(
+ image_id,
+ wandb.Image(images[image_index].cpu()),
+ predicted[image_index].item(),
+ labels[image_index].item(),
+ *scores[image_index].tolist(),
+ )
+
accuracy = 100 * correct / total
- # ✨ W&B: Log accuracy and the prediction table for this epoch.
- run.log(
- {
- "epoch": epoch,
- "acc": accuracy,
- "test_predictions": test_table,
- }
- )
+ training_run.log({
+ "epoch": epoch + 1,
+ "acc": accuracy,
+ "test_predictions": test_table,
+ })
print(
- "Test accuracy on the 10,000 MNIST test images: "
- f"{accuracy:.2f}%"
+ f"Epoch {epoch + 1}: test accuracy {accuracy:.2f}% "
+ f"on {total:,} images; {len(test_table.data)} table rows"
)
- training_result = {
- "run_url": run_url,
- "accuracy": accuracy,
- "epochs": EPOCHS,
- }
- return (training_result,)
+ run_url = training_run.url
+ print(f"View metrics and prediction tables: {run_url}")
+ return
-@app.cell
-def _(training_result):
- _run_url = training_result["run_url"]
- _accuracy = training_result["accuracy"]
- mo.callout(
- mo.md(
- f"""
- Training is complete. [Open the W&B run]({_run_url}) to explore the logged loss, accuracy, and prediction tables.
-
- Final test accuracy: **{_accuracy:.2f}%**
- """
- ),
- kind="success",
- title="W&B Tables run ready",
- )
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Explore Your Predictions in Weights & Biases
+
+ Open the run URL printed by the training cell and find the `test_predictions` Table.
+ Each row connects an image to its true label, predicted label, and class
+ probabilities.
+
+ Filter to rows where `guess != truth` to inspect mistakes. Group by
+ `guess` to find patterns in misclassified digits, or compare the
+ `score_0` through `score_9` columns to inspect the model's confidence.
+
+ With multiple epochs, use the table's step selector to compare snapshots.
+ Across experiments, stable image IDs let you match predictions for the same
+ test image. Try changing the learning rate or layer sizes, rerun training,
+ and compare the runs in your project. Repeatedly using the test results to
+ choose settings turns this split into validation data. Keep an independent
+ test set for a final evaluation in your own project.
+
+ ## Next steps
+
+ - See the [Weights & Biases Tables guide](https://docs.wandb.ai/models/tables)
+ for more ways to query and compare your results.
+ - Explore the [CNN-2: Progress over Training Time report](https://wandb.ai/stacey/table-quickstart/reports/CNN-2-Progress-over-Training-Time--Vmlldzo3NDY5ODU#false-positives-grouped-by-guess) for examples on how to compare, filter, and group predictions.
+ """)
return
+
if __name__ == "__main__":
app.run()
diff --git a/marimo/convert/wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_time_to_live_ttl_walkthrough.py b/marimo/convert/wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_time_to_live_ttl_walkthrough.py
index 9b859d51..683c0d3f 100644
--- a/marimo/convert/wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_time_to_live_ttl_walkthrough.py
+++ b/marimo/convert/wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_time_to_live_ttl_walkthrough.py
@@ -10,12 +10,15 @@
import marimo
-__generated_with = "0.24.0"
-app = marimo.App(width="medium", app_title="W&B Artifact TTL Walkthrough")
-
+__generated_with = "0.25.1"
+app = marimo.App(
+ width="medium",
+ app_title="Weights & Biases Artifact TTL Walkthrough",
+)
with app.setup:
import random
+ import shutil
import tarfile
import tempfile
from datetime import timedelta
@@ -31,17 +34,19 @@
@app.cell(hide_code=True)
def _():
mo.md(r"""
- # Weights & Biases Artifacts Time-to-live (TTL) Walkthrough
+ # Weights & Biases Artifacts: Time-to-Live (TTL)
[](https://molab.marimo.io/github/wandb/examples/blob/main/marimo/convert/wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_time_to_live_ttl_walkthrough.py/server)
- W&B Artifacts supports setting time-to-live policies on each version of an Artifact. The following examples show the use of TTL policy in a common Artifact logging workflow. We'll cover:
+ Use artifact TTL policies to schedule data retention. This lesson covers
+ setting a policy before upload, changing a policy on an existing version,
+ updating a collection of versions, and following artifact lineage.
- - Setting a TTL policy when creating an Artifact
- - Retroactively setting TTL for a specific Artifact alias
- - Using the W&B API to set a TTL for all versions of an Artifact
+ TTL is measured from an artifact version's creation time. Expiration
+ marks the version for deletion. Storage cleanup happens separately. Linking a
+ version to Registry deactivates its TTL.
- TTL applies to user-generated Artifacts. Linking an Artifact version to W&B Registry deactivates its TTL policy, so use the examples below on versions whose lifecycle you intend to manage directly.
+ Running the code cells downloads Imagenette and uploads sample artifacts.
""")
return
@@ -49,741 +54,331 @@ def _():
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Setup
- Let's do a few things before we get started. Below we will:
-
- - Connect to W&B
- - Download and sample a dataset
+ ## Authenticate with Weights & Biases
- The notebook dependencies install the required libraries automatically.
- """)
- return
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ### Authentication
+ ## Set your team and project
- Enter your [W&B API key](https://wandb.ai/authorize) and the team entity where you have permission to create runs and Artifacts. You can leave the key blank when this environment already has credentials, such as `WANDB_API_KEY` configured in the marimo Secrets panel. If the entity is blank, W&B uses your default entity.
-
- Changing these fields does nothing until you click **Connect to W&B**. Connecting authenticates and resolves the target path, but it does not create or modify any W&B object.
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `artifacts-ttl-demo`. Then select **Connect to Weights & Biases**.
""")
return
@app.cell(hide_code=True)
def _():
- _api_key_input = mo.ui.text(
- kind="password",
- label="W&B API key (optional)",
- placeholder="Paste a key or use configured credentials",
- full_width=True,
- )
- _entity_input = mo.ui.text(
- label="W&B entity or team (optional)",
- placeholder="Leave blank to use your default entity",
- full_width=True,
- )
- _project_input = mo.ui.text(
- value="artifacts-ttl-demo",
- label="W&B project",
- full_width=True,
- )
- wandb_login_form = (
- mo.md("{api_key}\n\n{entity}\n\n{project}")
- .batch(
- api_key=_api_key_input,
- entity=_entity_input,
- project=_project_input,
- )
- .form(submit_button_label="Connect to W&B", bordered=True)
- )
- wandb_login_form
- return (wandb_login_form,)
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="artifacts-ttl-demo", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
@app.cell(hide_code=True)
-def _(wandb_login_form):
+def _(wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
mo.stop(
- wandb_login_form.value is None,
- mo.callout(
- mo.md("Connect to W&B before running a remote Artifact step."),
- kind="info",
- ),
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
)
- _api_key = wandb_login_form.value["api_key"].strip()
- _submitted_entity = wandb_login_form.value["entity"].strip()
try:
- _login_ok = wandb.login(
- key=_api_key or None,
- relogin=bool(_api_key),
- )
- _login_error = None
- except wandb.errors.Error as _error:
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
_login_ok = False
- _login_error = str(_error)
mo.stop(
not _login_ok,
- mo.callout(
- mo.md(
- "W&B authentication did not complete. Check the API key and "
- f"try again.\n\nW&B reported: `{_login_error or 'unknown error'}`"
- ),
- kind="danger",
- ),
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
)
- entity = _submitted_entity or wandb.Api().default_entity
- mo.stop(
- not entity,
- mo.callout(
- mo.md(
- "W&B did not return a default entity. Enter a team entity in "
- "the form above and reconnect."
- ),
- kind="danger",
- ),
- )
- project = wandb_login_form.value["project"].strip() or "artifacts-ttl-demo"
- wandb_settings = {"entity": entity, "project": project}
- mo.callout(
- mo.md(f"Connected. Remote steps will target `{entity}/{project}`."),
- kind="success",
- )
- return (wandb_settings,)
+ entity = _entity
+ project = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return entity, project
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Image Sampling
- For the purposes of the walkthrough, we will sample from the Imagenette dataset and organize them into training and validation directories in our notebook session. The block below:
+ ## Prepare Sample Images
- - Creates folders for our sampled images if they don't already exist
- - Selects a random sample of images from the Imagenette dataset
- - Organizes the samples into training and validation directories
+ [Imagenette](https://github.com/fastai/imagenette) contains images from ten
+ ImageNet classes. The next cell downloads and extracts it into a temporary
+ workspace, then copies 5% of the training and validation images into sample
+ folders. Sampling leaves the extracted dataset intact, so rerunning the cell
+ does not repeatedly shrink it.
- *Note: we overwrite the files every time we execute this so we get new Artifact versions.*
+ We call the training sample "real" and the validation sample "fake" to
+ illustrate different retention policies. Both contain actual photographs.
+ The names describe roles in this example, not whether the images are generated.
""")
return
-@app.cell(hide_code=True)
+@app.cell
def _():
- prepare_imagenette = mo.ui.run_button(label="Download and sample Imagenette")
- mo.vstack(
- [
- mo.md(
- "This downloads and extracts Imagenette locally, then keeps a "
- "random 5% sample. It does not contact W&B."
- ),
- prepare_imagenette,
- ]
- )
- return (prepare_imagenette,)
-
+ workspace = Path(tempfile.mkdtemp(prefix="wandb-ttl-"))
+ archive = workspace / "imagenette.tgz"
-@app.cell(hide_code=True)
-def _(prepare_imagenette):
- mo.stop(
- not prepare_imagenette.value,
- mo.callout(
- mo.md("Click the button above when you are ready to prepare the dataset."),
- kind="info",
- ),
- )
- imagenette_url = "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz"
- imagenette_workspace = Path(
- tempfile.mkdtemp(prefix="wandb-ttl-imagenette-")
+ urlretrieve(
+ "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz",
+ archive,
)
- return imagenette_url, imagenette_workspace
-
-
-@app.cell
-def _(imagenette_url, imagenette_workspace):
- imagenette_archive = imagenette_workspace / "imagenette.tgz"
- urlretrieve(imagenette_url, imagenette_archive)
- return (imagenette_archive,)
-
-@app.function
-def untar_file(file_path, dest_path):
- destination = Path(dest_path).resolve()
- destination.mkdir(parents=True, exist_ok=True)
- with tarfile.open(file_path, "r:gz") as tar:
+ # Validate archive entries before extracting downloaded files.
+ with tarfile.open(archive, "r:gz") as tar:
members = tar.getmembers()
for member in members:
- member_path = (destination / member.name).resolve()
- if not member_path.is_relative_to(destination):
- raise ValueError(f"Unsafe path in archive: {member.name}")
+ destination = (workspace / member.name).resolve()
+ if not destination.is_relative_to(workspace.resolve()):
+ raise ValueError(f"Unsafe archive path: {member.name}")
if not (member.isfile() or member.isdir()):
raise ValueError(f"Unsupported archive entry: {member.name}")
- tar.extractall(destination, members=members)
+ tar.extractall(workspace, members=members)
+ train_source_dir = workspace / "sample" / "train"
+ val_source_dir = workspace / "sample" / "val"
+ rng = random.Random(1)
-@app.cell
-def _(imagenette_archive):
- untar_file(imagenette_archive, imagenette_archive.parent)
- dataset_dir = imagenette_archive.parent / "imagenette2-160"
- return (dataset_dir,)
+ for split, sample_dir in [("train", train_source_dir), ("val", val_source_dir)]:
+ _images = sorted((workspace / "imagenette2-160" / split).rglob("*.JPEG"))
+ selected = rng.sample(_images, max(1, len(_images) // 20))
+ for _image in selected:
+ target = sample_dir / _image.parent.name / _image.name
+ target.parent.mkdir(parents=True, exist_ok=True)
+ shutil.copy2(_image, target)
+ return train_source_dir, val_source_dir, workspace
@app.cell(hide_code=True)
def _():
mo.md(r"""
- We are going to use Imagenette dataset for this example. [Imagenette](https://github.com/fastai/imagenette) is a subset of 10 easily classified classes from Imagenet (tench, English springer, cassette player, chain saw, church, French horn, garbage truck, gas pump, golf ball, parachute). It was created by Jeremy Howard and is a great dataset to experiment with.
+ Preview five images from the training sample before uploading them as an
+ artifact. This cell only displays local files.
""")
return
@app.cell
-def _(dataset_dir):
- # let's keep 5% of the images
- for image in dataset_dir.rglob("*.JPEG"):
- if random.random() > 0.05:
- image.unlink()
-
- # we get two image folders: train and validation
- train_source_dir = dataset_dir / "train"
- val_source_dir = dataset_dir / "val"
- return train_source_dir, val_source_dir
+def _(train_source_dir):
+ _images = sorted(train_source_dir.rglob("*.JPEG"))[:5]
+ figure, axes = plt.subplots(1, len(_images), figsize=(15, 4), squeeze=False)
+ for axis, _image_path in zip(axes[0], _images):
+ with Image.open(_image_path) as _image:
+ axis.imshow(_image)
+ axis.axis("off")
+ figure.tight_layout()
+ plt.close(figure)
+ figure
+ return
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Image Preview
- Quick block to view some of the images in the sampled dataset.
- """)
- return
+ ## Set TTL When Creating an Artifact
+ The next cell starts a run, creates dataset artifacts with
+ `wandb.Artifact()`, adds files with `add_dir()`, and sets a ten-day TTL
+ before calling `raw_data_run.log_artifact()`.
-@app.function
-def show_sample_images(img_dir, num_images=5):
- images = list(img_dir.rglob("*.JPEG"))[:num_images]
- fig, axes = plt.subplots(1, len(images), figsize=(15, 5), squeeze=False)
+ We create two versions of `real-raw`, with aliases `extended` and
+ `compliant`, so the next example can give them different policies. Different
+ directory names inside their manifests ensure these are distinct versions,
+ even though they contain the same images. A TTL applies to a version, so
+ all aliases pointing to that version share its retention policy.
- # Iterate over the images and display them
- for i, img_path in enumerate(images):
- img = Image.open(img_path)
- axes[0, i].imshow(img)
- axes[0, i].axis("off") # Turn off axis labels
-
- plt.tight_layout()
- return fig
+ We also log `fake-raw` with a ten-day TTL for the bulk-update example.
+ `wait()` ensures each upload finishes before later cells fetch it.
+ """)
+ return
@app.cell
-def _(train_source_dir):
- sample_figure = show_sample_images(train_source_dir)
- sample_figure
- return
+def _(entity, project, train_source_dir, val_source_dir):
+ with wandb.init(entity=entity, project=project, job_type="raw-data") as raw_data_run:
+ extended = wandb.Artifact("real-raw", type="dataset")
+ extended.add_dir(str(train_source_dir), name="extended-snapshot")
+ extended.ttl = timedelta(days=10)
+ logged_extended = raw_data_run.log_artifact(extended, aliases=["extended"])
+ logged_extended.wait()
+
+ compliant = wandb.Artifact("real-raw", type="dataset")
+ compliant.add_dir(str(train_source_dir), name="compliant-snapshot")
+ compliant.ttl = timedelta(days=10)
+ logged_compliant = raw_data_run.log_artifact(compliant, aliases=["compliant", "latest"])
+ logged_compliant.wait()
+
+ fake = wandb.Artifact("fake-raw", type="dataset")
+ fake.add_dir(str(val_source_dir))
+ fake.ttl = timedelta(days=10)
+ logged_fake = raw_data_run.log_artifact(fake)
+ logged_fake.wait()
+
+ project_path = f"{raw_data_run.entity}/{raw_data_run.project}"
+ raw_alias_paths = {
+ "extended": f"{project_path}/real-raw:extended",
+ "compliant": f"{project_path}/real-raw:compliant",
+ }
+ return project_path, raw_alias_paths
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Setting TTL on New Artifacts
- Below we create two new Artifacts for our real and fake data. Because we have internal retention policies in hypothetical organization we'd like to remove any Artifact that has real data (potentially containing personal data). Below we:
+ ## Change TTL on Existing Versions
- - Create a W&B Run to track the logging of these raw data Artifacts
- - Set the ttl attribute on the real raw data
- - Log our two Artifacts
+ Suppose we can now keep one version for a year and another indefinitely.
+ The next cell fetches each version by alias with `wandb.Api().artifact()`,
+ changes its `ttl`, and calls `save()` to persist the change.
- > We will use the train dataset as our real data and the validation dataset as our fake data.
+ A `timedelta` sets the retention duration. Setting `ttl = None` deactivates TTL. Changing
+ the duration does not reset the version's creation date. These metadata
+ updates do not require a new run or another file upload.
""")
return
-@app.cell(hide_code=True)
-def _(train_source_dir, val_source_dir, wandb_settings):
- log_raw_artifacts_button = mo.ui.run_button(
- label="Log TTL demo Artifact versions"
- )
- mo.vstack(
- [
- mo.md(
- f"This creates a `raw-data` run in "
- f"`{wandb_settings['entity']}/{wandb_settings['project']}`, "
- "logs two distinct `real-raw` versions with a 10-day TTL, "
- "assigns them the `extended` and `compliant` aliases, and "
- "logs `fake-raw` without a custom TTL. The files come from "
- f"`{train_source_dir}` and `{val_source_dir}`."
- ),
- log_raw_artifacts_button,
- ]
- )
- return (log_raw_artifacts_button,)
-
-
-@app.function
-def log_raw_artifacts(entity, project, train_source_dir, val_source_dir):
- with wandb.init(
- entity=entity,
- project=project,
- job_type="raw-data",
- reinit="create_new",
- ) as run:
- raw_extended_art = wandb.Artifact(
- "real-raw",
- type="dataset",
- description="Raw Imagenette sample approved for extended retention",
- )
- raw_extended_art.add_dir(train_source_dir, name="extended-snapshot")
- raw_extended_art.ttl = timedelta(days=10)
- logged_extended_art = run.log_artifact(
- raw_extended_art,
- aliases=["extended"],
- )
- logged_extended_art.wait()
-
- raw_compliant_art = wandb.Artifact(
- "real-raw",
- type="dataset",
- description="Raw Imagenette sample approved for indefinite retention",
- )
- raw_compliant_art.add_dir(train_source_dir, name="compliant-snapshot")
- raw_compliant_art.ttl = timedelta(days=10)
- logged_compliant_art = run.log_artifact(
- raw_compliant_art,
- aliases=["compliant", "latest"],
- )
- logged_compliant_art.wait()
-
- raw_fake_art = wandb.Artifact(
- "fake-raw",
- type="dataset",
- description="Raw sample from val Imagenette",
- )
-
- raw_fake_art.add_dir(val_source_dir)
- logged_fake_art = run.log_artifact(raw_fake_art)
- logged_fake_art.wait()
- run_url = run.url
- alias_paths = {
- "extended": f"{entity}/{project}/real-raw:extended",
- "compliant": f"{entity}/{project}/real-raw:compliant",
- }
- return run_url, alias_paths
+@app.cell
+def _(raw_alias_paths):
+ _api = wandb.Api()
+ extended_artifact = _api.artifact(raw_alias_paths["extended"])
+ extended_artifact.ttl = timedelta(days=365)
+ extended_artifact.save()
+ compliant_artifact = _api.artifact(raw_alias_paths["compliant"])
+ compliant_artifact.ttl = None
+ compliant_artifact.save()
-@app.cell(hide_code=True)
-def _(
- log_raw_artifacts_button,
- train_source_dir,
- val_source_dir,
- wandb_settings,
-):
- mo.stop(
- not log_raw_artifacts_button.value,
- mo.callout(
- mo.md("Click the button above when you are ready to create the Artifact versions."),
- kind="info",
- ),
- )
- raw_artifacts_run_url, raw_alias_paths = log_raw_artifacts(
- wandb_settings["entity"],
- wandb_settings["project"],
- train_source_dir,
- val_source_dir,
- )
- mo.callout(
- mo.md(f"Artifacts logged: [open the W&B run]({raw_artifacts_run_url})."),
- kind="success",
- )
- return raw_alias_paths, raw_artifacts_run_url
+ print(f"extended: {_api.artifact(raw_alias_paths['extended']).ttl}")
+ print(f"compliant: {_api.artifact(raw_alias_paths['compliant']).ttl}")
+ return (compliant_artifact,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Updating/Retroactively Setting TTL on Artifacts
- In our hypothetical organization we've been given approval to retain a specific version of our data indefinitely. We've also been given approval to extend the retention date of an additional dataset. Below we'll:
+ ## Update All Versions in a Collection
+
+ `wandb.Api().artifacts()` lists the versions in an artifact collection.
+ The next cell queries `fake-raw` directly, then deactivates each version's
+ TTL with the same `ttl = None` and `save()` pattern.
- - Extend the TTL of an Artifact tagged with the `extended` alias
- - Remove the TTL of an Artifact tagged with the `compliant` alias
- - Programmatically check the status of these two Artifacts
+ This applies to every version of `fake-raw` in the selected project,
+ including versions from previous executions of the lesson.
- The preceding logging step creates two distinct `real-raw` versions and assigns these aliases. Retroactive TTL changes take effect when the code calls `artifact.save()`.
+ Team admins can also configure [default TTL policies](https://docs.wandb.ai/models/artifacts/ttl#set-default-ttl-policies-for-a-team)
+ for artifacts without a custom policy.
""")
return
-@app.cell(hide_code=True)
-def _(raw_alias_paths, raw_artifacts_run_url, wandb_settings):
- update_alias_ttls_button = mo.ui.run_button(
- label="Update TTLs for real-raw:extended and :compliant"
- )
- mo.vstack(
- [
- mo.md(
- f"This reads those two aliases from "
- f"`{wandb_settings['entity']}/{wandb_settings['project']}`. "
- "It saves only TTL values that need to change. The source "
- f"Artifact upload completed in [this run]({raw_artifacts_run_url}). "
- f"The exact inputs are `{raw_alias_paths['extended']}` and "
- f"`{raw_alias_paths['compliant']}`."
- ),
- update_alias_ttls_button,
- ]
- )
- return (update_alias_ttls_button,)
-
-
-@app.function
-def update_alias_ttls(entity, project, alias_paths):
- extended_path = alias_paths["extended"]
- compliant_path = alias_paths["compliant"]
- api = wandb.Api()
- extended_current = api.artifact(extended_path)
- compliant_current = api.artifact(compliant_path)
- extended_ttl = timedelta(days=365)
- needs_extended_update = extended_current.ttl != extended_ttl
- needs_compliant_update = compliant_current.ttl is not None
-
- changed = []
- run_url = None
- if needs_extended_update or needs_compliant_update:
- with wandb.init(
- entity=entity,
- project=project,
- job_type="modify-ttl",
- reinit="create_new",
- ) as run:
- extended_art = run.use_artifact(extended_path)
- if needs_extended_update:
- extended_art.ttl = extended_ttl # Delete in a year
- extended_art.save()
- changed.append("real-raw:extended")
-
- compliant_art = run.use_artifact(compliant_path)
- if needs_compliant_update:
- compliant_art.ttl = None
- compliant_art.save()
- changed.append("real-raw:compliant")
-
- print(extended_art.ttl)
- print(compliant_art.ttl)
- run_url = run.url
- return changed, run_url
-
-
-@app.cell(hide_code=True)
-def _(raw_alias_paths, update_alias_ttls_button, wandb_settings):
- mo.stop(
- not update_alias_ttls_button.value,
- mo.callout(
- mo.md("Click the button above when you are ready to update the aliases."),
- kind="info",
- ),
- )
- updated_aliases, alias_update_run_url = update_alias_ttls(
- wandb_settings["entity"],
- wandb_settings["project"],
- raw_alias_paths,
- )
- if updated_aliases:
- _alias_message = (
- f"Updated `{', '.join(updated_aliases)}`. "
- f"[Open the W&B run]({alias_update_run_url})."
- )
- else:
- _alias_message = "Both aliases already had the requested TTL values; no run or write was created."
- mo.callout(mo.md(_alias_message), kind="success")
+@app.cell
+def _(project_path):
+ _api = wandb.Api()
+ cleared_versions = []
+ for _artifact in _api.artifacts(type_name="dataset", name=f"{project_path}/fake-raw"):
+ _artifact.ttl = None
+ _artifact.save()
+ cleared_versions.append(_artifact.qualified_name)
+ print(f"{_artifact.name}: {_artifact.ttl}")
return
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ## Use W&B Import/Export API to Iterate Artifact Versions and Set TTL
- Let's say we've received approval to retain all of the data within a given Artifact and we'd like to remove all TTL policies for every version of an Artifact. Below we:
+ ## Record Lineage During Preprocessing
+
+ The next cell consumes the `compliant` version of `real-raw`, resizes its
+ images to 64 by 64 pixels, and logs a new `real-prepro` artifact.
- - Use the W&B API to get a list of all Runs in a project
- - Get a list of all versions of a specific Artifact (e.g. `fake-raw`)
- - Iterate over each version and remove any existing TTL policy associated with the version
+ `preprocessing_run.use_artifact()` records the input relationship, and
+ `preprocessing_run.log_artifact()` records the output. Together they give Weights & Biases a lineage
+ graph connecting the source data, preprocessing run, and processed data.
""")
return
-@app.cell(hide_code=True)
-def _(raw_artifacts_run_url, wandb_settings):
- clear_fake_raw_ttls_button = mo.ui.run_button(
- label="Clear TTL from every fake-raw version"
- )
- mo.vstack(
- [
- mo.md(
- f"This searches `{wandb_settings['entity']}/{wandb_settings['project']}` "
- "and saves only versions whose TTL is currently set. The "
- f"source upload completed in [this run]({raw_artifacts_run_url})."
- ),
- clear_fake_raw_ttls_button,
- ]
- )
- return (clear_fake_raw_ttls_button,)
-
-
-@app.function
-def clear_fake_raw_ttls(entity, project):
- # Artifact metadata extraction
- api = wandb.Api()
-
- # Define entity and project
- runs = api.runs(entity + "/" + project)
-
- version_names = []
- for run in runs:
- for artifact in iter(run.logged_artifacts()):
- if "fake-raw" in artifact.name:
- # Can be edited to just display individual elements
- version_names.append(artifact.qualified_name)
-
- unique_versions = sorted(set(version_names))
- versions_to_clear = [
- version
- for version in unique_versions
- if api.artifact(version).ttl is not None
- ]
- run_url = None
- if versions_to_clear:
- with wandb.init(
- entity=entity,
- project=project,
- job_type="modify-ttl",
- reinit="create_new",
- ) as run:
- for version in versions_to_clear:
- version_art = run.use_artifact(version)
- version_art.ttl = None
- version_art.save()
- print(version_art.ttl)
- run_url = run.url
- return unique_versions, versions_to_clear, run_url
-
-
-@app.cell(hide_code=True)
-def _(clear_fake_raw_ttls_button, wandb_settings):
- mo.stop(
- not clear_fake_raw_ttls_button.value,
- mo.callout(
- mo.md("Click the button above when you are ready to clear the TTL policies."),
- kind="info",
- ),
- )
- fake_raw_versions, cleared_versions, clear_ttls_run_url = clear_fake_raw_ttls(
- wandb_settings["entity"], wandb_settings["project"]
- )
- if cleared_versions:
- _clear_message = (
- f"Cleared TTL from {len(cleared_versions)} of "
- f"{len(fake_raw_versions)} discovered versions. "
- f"[Open the W&B run]({clear_ttls_run_url})."
- )
- else:
- _clear_message = (
- f"Found {len(fake_raw_versions)} `fake-raw` versions, all already "
- "without a custom TTL. No run or write was created."
- )
- mo.callout(mo.md(_clear_message), kind="success")
- return
+@app.cell
+def _(compliant_artifact, entity, project, workspace):
+ with tempfile.TemporaryDirectory(dir=workspace) as _data_dir:
+ output_dir = Path(_data_dir) / "processed"
+ output_dir.mkdir()
+ with wandb.init(entity=entity, project=project, job_type="preprocessing") as preprocessing_run:
+ source = preprocessing_run.use_artifact(compliant_artifact.qualified_name)
+ source_dir = Path(source.download(root=str(Path(_data_dir) / "raw")))
+ for _image_path in source_dir.rglob("*.JPEG"):
+ _target = output_dir / _image_path.relative_to(source_dir)
+ _target.parent.mkdir(parents=True, exist_ok=True)
+ with Image.open(_image_path) as _image:
+ _image.resize((64, 64)).save(_target)
+
+ processed = wandb.Artifact("real-prepro", type="dataset")
+ processed.add_dir(str(output_dir))
+ processed_artifact = preprocessing_run.log_artifact(processed)
+ processed_artifact.wait()
+ print(f"View preprocessing run: {preprocessing_run.url}")
+ return (processed_artifact,)
@app.cell(hide_code=True)
def _():
mo.md(r"""
- > To apply a TTL policy to all artifacts within a team's projects, team admins can set default TTL policies for their team. The default will be applied to both existing and future artifacts logged to projects as long as no custom policies have been set. To learn more about configuring a team default TTL, visit [this](https://docs.wandb.ai/models/artifacts/ttl#set-default-ttl-policies-for-a-team) section of the W&B documentation.
- """)
- return
+ ## Follow Lineage to Set Downstream TTL
+ TTL policies do not automatically propagate to derived artifacts. The next
+ cell fetches our source version and uses `used_by()` to find its consuming
+ runs. For each run, `logged_artifacts()` lists its outputs. We select dataset
+ outputs, then give them a 30-day TTL and save the changes.
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ## Traverse an Artifact Graph to Set Downstream TTL
- In this last section, we'll do some preprocessing on our images and log those as downstream Artifacts. Once again we'll use the W&B Import/Export API to set a TTL policy on our downstream images for images that originated from our "real" dataset.
+ This follows one dependency step from the source Artifact, through each
+ consuming run, to that run’s output Artifacts. A longer pipeline would repeat this traversal for each output.
""")
return
-@app.cell(hide_code=True)
-def _():
- mo.md(r"""
- ### Preprocess and log a new Artifact
- """)
+@app.cell
+def _(compliant_artifact, processed_artifact):
+ # Referencing the uploaded output orders this cell after preprocessing.
+ print(f"Preprocessing complete: {processed_artifact.qualified_name}")
+ _api = wandb.Api()
+ _source = _api.artifact(compliant_artifact.qualified_name)
+ downstream_ttls = {}
+ for consumer_run in _source.used_by():
+ for _artifact in consumer_run.logged_artifacts():
+ if _artifact.type == "dataset" and _artifact.qualified_name not in downstream_ttls:
+ _artifact.ttl = timedelta(days=30)
+ _artifact.save()
+ downstream_ttls[_artifact.qualified_name] = _artifact.ttl
+ print(f"{_artifact.qualified_name}: {_artifact.ttl}")
return
-@app.cell(hide_code=True)
-def _(raw_artifacts_run_url, wandb_settings):
- preprocess_artifact_button = mo.ui.run_button(
- label="Preprocess real-raw:latest and log real-prepro"
- )
- mo.vstack(
- [
- mo.md(
- f"This downloads `real-raw:latest` from "
- f"`{wandb_settings['entity']}/{wandb_settings['project']}`, "
- "resizes its images, and creates a new downstream Artifact "
- "version. Every click intentionally creates a new version. "
- f"The source upload completed in [this run]({raw_artifacts_run_url})."
- ),
- preprocess_artifact_button,
- ]
- )
- return (preprocess_artifact_button,)
-
-
-@app.function
-def preprocess_image(image_path):
- "Resize the image to 64x64"
- return Image.open(image_path).resize((64, 64))
-
-
-@app.function
-def preprocess_and_log_artifact(entity, project):
- with tempfile.TemporaryDirectory() as data_dir:
- real_prepro_dir = Path(data_dir) / "prepro" / "real"
- real_prepro_dir.mkdir(parents=True, exist_ok=True)
-
- with wandb.init(
- entity=entity,
- project=project,
- job_type="preprocessing",
- reinit="create_new",
- ) as run:
- real_art = run.use_artifact(f"{entity}/{project}/real-raw:latest")
- real_images = Path(
- real_art.download(root=Path(data_dir) / "real-raw")
- )
-
- for image_path in real_images.rglob("*.JPEG"):
- print(f"Preprocessing {image_path.name}")
- preprocessed_image = preprocess_image(image_path)
- preprocessed_image.save(real_prepro_dir / image_path.name)
-
- prepro_real_art = wandb.Artifact(
- "real-prepro",
- type="dataset",
- description="Preprocessed images from Imagenette",
- )
-
- prepro_real_art.add_dir(real_prepro_dir)
- logged_prepro_art = run.log_artifact(prepro_real_art)
- logged_prepro_art.wait()
- run_url = run.url
- return run_url
-
-
-@app.cell(hide_code=True)
-def _(preprocess_artifact_button, wandb_settings):
- mo.stop(
- not preprocess_artifact_button.value,
- mo.callout(
- mo.md("Click the button above when you are ready to create the downstream Artifact."),
- kind="info",
- ),
- )
- preprocessing_run_url = preprocess_and_log_artifact(
- wandb_settings["entity"], wandb_settings["project"]
- )
- mo.callout(
- mo.md(f"Downstream Artifact logged: [open the W&B run]({preprocessing_run_url})."),
- kind="success",
- )
- return (preprocessing_run_url,)
-
-
@app.cell(hide_code=True)
def _():
mo.md(r"""
- ### Traverse the Artifact Graph and Set TTL
- Let's take a look at the original real dataset and traverse downstream runs and Artifacts to set a TTL policy on anything that originated from the real dataset.
- """)
- return
-
+ ## Inspect TTL Policies
-@app.cell(hide_code=True)
-def _(preprocessing_run_url, wandb_settings):
- set_downstream_ttls_button = mo.ui.run_button(
- label="Assign random TTLs to downstream dataset versions"
- )
- mo.vstack(
- [
- mo.md(
- f"This traverses consumers of `real-raw:latest` in "
- f"`{wandb_settings['entity']}/{wandb_settings['project']}` and "
- "saves a new random 1–100 day TTL on each downstream dataset "
- "version. Every click intentionally changes those policies. "
- f"The preprocessing upload completed in [this run]({preprocessing_run_url})."
- ),
- set_downstream_ttls_button,
- ]
- )
- return (set_downstream_ttls_button,)
-
-
-@app.function
-def set_downstream_ttls(entity, project):
- api = wandb.Api()
-
- # For demo purposes we'll just do this on the latest version of the real dataset
- artifact = api.artifact(f"{entity}/{project}/real-raw:latest")
- consumer_runs = artifact.used_by()
-
- # Same pattern from above to get all downstream versions
- version_names = []
- for run in consumer_runs:
- for artifact in iter(run.logged_artifacts()):
- # filter for datasets only
- if artifact.type == "dataset":
- # Can be edited to just display individual elements
- version_names.append(artifact.qualified_name)
-
- unique_versions = sorted(set(version_names))
- changed_ttls = {}
- run_url = None
- if unique_versions:
- with wandb.init(
- entity=entity,
- project=project,
- job_type="modify-ttl",
- reinit="create_new",
- ) as run:
- for version in unique_versions:
- version_art = run.use_artifact(version)
- # set ttl to a random integer so we can see changes in the UI after we run this
- ttl_days = random.randint(1, 100)
- version_art.ttl = timedelta(days=ttl_days)
- version_art.save()
- changed_ttls[version] = ttl_days
- run_url = run.url
- return changed_ttls, run_url
+ Open the project link printed earlier and select **Artifacts**. Check the
+ versions and their TTL policies: `real-raw:extended` has 365 days,
+ `real-raw:compliant` and `fake-raw` have TTL deactivated, and downstream
+ datasets have 30 days. The lineage view shows the preprocessing relationship.
-
-@app.cell(hide_code=True)
-def _(set_downstream_ttls_button, wandb_settings):
- mo.stop(
- not set_downstream_ttls_button.value,
- mo.callout(
- mo.md("Click the button above when you are ready to change downstream TTL policies."),
- kind="info",
- ),
- )
- downstream_ttls, downstream_ttl_run_url = set_downstream_ttls(
- wandb_settings["entity"], wandb_settings["project"]
- )
- if downstream_ttls:
- _downstream_message = (
- f"Updated {len(downstream_ttls)} downstream dataset versions. "
- f"[Open the W&B run]({downstream_ttl_run_url})."
- )
- else:
- _downstream_message = "No downstream dataset versions were found, so no run or write was created."
- mo.callout(mo.md(_downstream_message), kind="success")
+ The reusable pattern is to fetch a version, assign `artifact.ttl`, and call
+ `artifact.save()`. See the [TTL documentation](https://docs.wandb.ai/models/artifacts/ttl)
+ for permissions, team defaults, and expiration behavior.
+ """)
return
diff --git a/marimo/convert/zoo-wandb/zoo_wandb.py b/marimo/convert/zoo-wandb/zoo_wandb.py
index 4331db04..8451598a 100644
--- a/marimo/convert/zoo-wandb/zoo_wandb.py
+++ b/marimo/convert/zoo-wandb/zoo_wandb.py
@@ -36,13 +36,18 @@ def _(mo):
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- The goal of this notebook is to demonstrate how you and other members of your organization can use W&B Registry to track, share, and use dataset and model artifacts in your machine learning workflows. By the end of this notebook, you will know how to use W&B to:
+ Use Weights & Biases Registry to share versioned datasets and models across
+ teams. This lesson converts the UCI Zoo dataset to tensors, logs and links
+ those files to a Registry collection, trains a classifier, and retrieves the
+ saved model and test data for inference.
- 1. Create [collections](https://docs.wandb.ai/guides/registry/create_collection) within [W&B Registry](https://docs.wandb.ai/guides/registry)
- 2. Make dataset and model artifacts available to other members of your organization, and
- 3. Download your trained model and dataset artifacts from the registry for inference
+ The classifier uses animal attributes such as feathers, fins, and number of
+ legs to predict one of seven dataset classes. The model is intentionally
+ small so the lesson can focus on Artifact paths, version selection, and lineage.
- To achieve this, we will train a neural network to identify animal classes (mammal, amphibian, reptile, and so forth) based on features such as weather or not they ahve feathers, fins, and so on.
+ Before publishing, check that the target registries exist and that your
+ account can link versions to their collections. Registry access also determines
+ which colleagues can retrieve the linked files.
""")
return
@@ -74,6 +79,62 @@ def _():
return fetch_ucirepo, nn, pd, torch, train_test_split, wandb
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Authenticate with Weights & Biases
+
+ Paste your [Forge API key](https://docs.coreweave.com/products/wandb/platform/hosting/iam/api-keys) into the password field below. Keep your key out of the notebook’s source code.
+
+ If this runtime already has credentials in `WANDB_API_KEY` or your `.netrc` file, leave the API key field blank.
+
+ ## Set your team and project
+
+ Enter your Weights & Biases team name and choose a project for this lesson. The default project is `zoo_experiment`. Then select **Connect to Weights & Biases**.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ wandb_setup_form = mo.md("{api_key}\n\n{entity}\n\n{project}").batch(
+ api_key=mo.ui.text(kind="password", label="Forge API key", full_width=True),
+ entity=mo.ui.text(label="Team", placeholder="Your team name", full_width=True),
+ project=mo.ui.text(value="zoo_experiment", label="Project", full_width=True),
+ ).form(submit_button_label="Connect to Weights & Biases", bordered=True)
+ wandb_setup_form
+ return (wandb_setup_form,)
+
+
+@app.cell(hide_code=True)
+def _(mo, wandb, wandb_setup_form):
+ mo.stop(wandb_setup_form.value is None)
+
+ _submission = wandb_setup_form.value
+ _api_key = _submission["api_key"].strip()
+ _entity = _submission["entity"].strip()
+ _project = _submission["project"].strip()
+ mo.stop(
+ not _entity or not _project,
+ mo.callout("Enter your team and project, then connect again.", kind="warn"),
+ )
+
+ try:
+ _login_ok = wandb.login(key=_api_key or None, relogin=bool(_api_key))
+ except (wandb.errors.Error, ValueError):
+ _login_ok = False
+
+ mo.stop(
+ not _login_ok,
+ mo.callout("Check your API key or saved credentials, then connect again.", kind="danger"),
+ )
+
+ TEAM_ENTITY = _entity
+ PROJECT = _project
+ mo.md("✓ Connected to Weights & Biases.")
+ return PROJECT, TEAM_ENTITY
+
+
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
@@ -123,9 +184,14 @@ def _(X):
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- ### Process data
+ ### Convert the data to tensors
- For training let's convert our dataset from a [pandas `DataFrame`](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) to [a tensor with PyTorch](https://pytorch.org/docs/stable/generated/torch.tensor.html#torch.tensor), convert the data type of our input tensor(float64 to float32) to match the data type of the `nn.Linear module`, and convert our label tensor to index from 0-6:
+ Convert feature values to `float32`, matching the linear layers’ weights.
+ Convert labels to integer tensors and subtract one so the original class
+ IDs 1–7 become 0–6, as required by the seven-output classifier.
+
+ Keep the numeric class IDs consistent through training and evaluation.
+ Human-readable labels are added only when displaying predictions.
""")
return
@@ -161,24 +227,22 @@ def _(dataset, labels, torch):
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- ## Track and publish dataset
-
- Within the Dataset registry we will create a collection called "zoo-dataset-tensors". A *collection* is a set of linked artifact versions in a registry.
+ ## Log and link the dataset
- To create a collection we need to do two things:
- 1. Specify the collection and registry we want to link our artifact version to. To do this, we specify a "target path" for our artifact version.
- 2. Use the `wandb.run.link_artifact` method and pass our artifact object and the target path.
+ An Artifact stores a versioned set of files in a project. A Registry collection
+ contains links to Artifact versions so authorized users can discover and reuse
+ them. Linking creates a reference to the source version.
- #### Define target path of the collection
+ The target path identifies the registry and collection. The registry must
+ already exist. This example uses a `dataset` registry and a collection named
+ `zoo-dataset-tensors`. Match the registry name to the one in your organization.
- The target path of a collection consists of two parts:
- * The name of the registry
- * The name of the collection within the registry
-
- If you know these two fields, you can create the full name yourself with string concatenation, f-strings, and so forth:
```python
target_path = f"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}"
```
+
+ See [Link an artifact version to a collection](https://docs.coreweave.com/products/registry/link_version)
+ for the required permissions and supported paths.
""")
return
@@ -186,17 +250,15 @@ def _(mo):
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- ### Publish dataset to registry
+ ### Publish the dataset to Registry
- Let's publish our dataset to the Dataset registry in a collection called "zoo-dataset-tensors". To do this, we will
+ Define the target path, start a run, and create a dataset Artifact. Add the
+ feature and label tensor files, then call `run.link_artifact()` with the
+ Artifact and target path. For this draft Artifact, the call logs the files
+ as a run output and links the resulting version to the collection.
- 1. Get or create the target path. (For this notebook we will programmatically create the target path).
- 1. Initialize a run
- 1. Create an artifact object
- 2. Add each split dataset as individual files to the artifact object
- 3. Link the artifact object to the collection with `run.link_artifact()`. Here we specify the target path and the artifact we want to link.
-
- First, let's create the target path:
+ These first files contain the full dataset. The following section creates
+ and publishes separate training and test files.
""")
return
@@ -214,15 +276,15 @@ def _():
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
- Now that we have the target path, let's publish the dataset to the "Dataset" registry. In the following code cell, ensure to replace the values enclosed in `<>` with your team's entity:
+ The next cell uses the team and project from the setup form. It adds
+ `zoo_dataset.pt` and `zoo_labels.pt` to a dataset Artifact, links the version
+ to `dataset_target_path`, and finishes the run.
""")
return
@app.cell
-def _(dataset_target_path, wandb):
- TEAM_ENTITY = "