diff --git a/colabs/intro/run_quickstart.py b/colabs/intro/run_quickstart.py
new file mode 100644
index 00000000..1918ca36
--- /dev/null
+++ b/colabs/intro/run_quickstart.py
@@ -0,0 +1,226 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+# "marimo",
+# "wandb",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.23.16"
+app = marimo.App()
+
+with app.setup(hide_code=True):
+ import marimo as mo
+ import wandb
+ import random
+
+
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ # W&B Quickstart
+
+ Use W&B to track, visualize, and manage machine learning experiments of any size.
+
+ ## Create a machine learning training experiment
+
+ The following example simulates a simple training experiment and logs metrics to W&B.
+
+ First, define the W&B project name and a `config` dictionary. The config stores the input values for the experiment, such as the number of epochs and the learning rate. In this notebook, the form below collects the information for the `config`.
+
+ Next, initialize a W&B run with [`wandb.init()`](https://docs.wandb.ai/models/ref/python/functions/init). The run records the config, metrics, and other information from the training script.
+
+ 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.
+
+ ### Authentication
+
+ To save your experiment in W&B, you need to authenticate.
+
+ 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).
+ """)
+ 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)",
+ )
+
+ # 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**
+
+ {epochs}
+
+ {lr}
+
+ **W&B run.**
+
+ {api_key}
+
+ {project}
+
+ {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
+def _(form):
+ run_path = None
+
+ 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."
+ ),
+ )
+
+ # Get the data from the form
+ cfg = form.value
+
+ # Dictionary with hyperparameters
+ config = {
+ "epochs": cfg["epochs"],
+ "lr": cfg["lr"],
+ }
+
+ # 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}"
+ return (run_path,)
+
+
+@app.cell(hide_code=True)
+def _():
+ mo.md(r"""
+ ## Get the logged run results from W&B
+ """)
+ 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()
+ 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
+def _(history, remote_run):
+ mo.stop(
+ not history,
+ mo.md("The run finished, but no metric history is available yet."),
+ )
+
+ 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),
+ ]
+ )
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/audiocraft-audiocraft/audiocraft_audiocraft.py b/marimo/convert/audiocraft-audiocraft/audiocraft_audiocraft.py
new file mode 100644
index 00000000..8dc4c965
--- /dev/null
+++ b/marimo/convert/audiocraft-audiocraft/audiocraft_audiocraft.py
@@ -0,0 +1,433 @@
+# /// script
+# dependencies = ["audiocraft @ git+https://git@github.com/facebookresearch/audiocraft", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🎸 Generating Music using [Audiocraft](https://github.com/facebookresearch/audiocraft) and W&B 🐝
+
+
+ In this notebook we demonstrate how you can generate music and other types of audio from text prompts or generate new music from existing music using SoTA models such as [MusicGen](https://github.com/facebookresearch/audiocraft/blob/main/docs/MUSICGEN.md) and [AudioGen](https://github.com/facebookresearch/audiocraft/blob/main/docs/AUDIOGEN.md) from [Audiocraft](https://github.com/facebookresearch/audiocraft) and play and visualize them using [Weights & Biases](https://wandb.ai/site).
+
+ If you want to know more about the underlying architectures for MusicGen and AudioGen and explore some cool audio samples generated by these models, you can check out [this W&B report](http://wandb.me/audiocraft_2mp).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # @title Install AudioCraft + WandB
+ # packages added via marimo's package management: git+https://git@github.com/facebookresearch/audiocraft !pip install -U git+https://git@github.com/facebookresearch/audiocraft
+ #egg=audiocraft
+ # packages added via marimo's package management: wandb !pip install -qq -U wandb
+ return
+
+
+@app.cell
+def _():
+ # @title
+ import os
+ import random
+ from tempfile import TemporaryDirectory
+
+ from scipy import signal
+ from scipy.io import wavfile
+
+ import torchaudio
+ from audiocraft.models import AudioGen, MusicGen, MultiBandDiffusion
+ from audiocraft.data.audio import audio_write
+
+ import wandb
+ import numpy as np
+ from tqdm.auto import tqdm
+ from google.colab import files
+ import matplotlib.pyplot as plt
+
+ return (
+ AudioGen,
+ MultiBandDiffusion,
+ MusicGen,
+ TemporaryDirectory,
+ audio_write,
+ files,
+ np,
+ os,
+ plt,
+ random,
+ signal,
+ torchaudio,
+ wandb,
+ wavfile,
+ )
+
+
+@app.cell
+def _(files, random, torchaudio, wandb):
+ # @title ## Audio Generation Configs
+
+ # @markdown In this section, you can interact with the user interface to chose the models you want to use to generate audio, prompts and other configs. Once you execute this cell, it initializes a [wandb run](https://docs.wandb.ai/guides/runs) which will be used to automatically log all the generated audio along with all the prompts and configs, to ensure your AI-generated music is never lost and your experiments are always reproducible and easy to share.
+
+ # @markdown **Note:** If you have provided prompts, you will be prompted to provide an audio file in addition to the prompts to condition the model. If you don't want to provide a file as an additional condition to the model, just press on the `cancel` button.
+
+ # @markdown ---
+ # @markdown WandB Project Name
+ project_name = "audiocraft" # @param {type:"string"}
+
+ wandb.init(project=project_name, job_type="musicgen/inference")
+
+ config = wandb.config
+
+ # @markdown Select the Model for audio generation supported by [AudioCraft](https://github.com/facebookresearch/audiocraft). You can select either the MusicGen model variants (great for generating music) or the AudioGen model variants (great for generating non-musical audio). Also note that you can run all variants of MusicGen except the `large` one on the free-tier Colab GPU.
+ model_name = "musicgen-small" # @param ["musicgen-small", "musicgen-medium", "musicgen-large", "musicgen-melody", "audiogen-medium"]
+ config.model_name = "facebook/" + model_name if model_name == "audiogen-medium" else model_name
+
+ # @markdown Whether to enable [MultiBand Diffusion](https://github.com/facebookresearch/audiocraft/blob/main/docs/MBD.md) or not. MultiBand diffusion is a collection of 4 models that can decode tokens from EnCodec tokenizer into waveform audio. Note that enabling this increases the time required to generate the audio.
+ enable_multi_band_diffusion = True # @param {type:"boolean"}
+ # config.enable_multi_band_diffusion = enable_multi_band_diffusion
+
+ if "musicgen" not in model_name:
+ wandb.termwarn("Multi-band Diffusion is only available for Musicgen")
+ config.enable_multi_band_diffusion = False
+ else:
+ config.enable_multi_band_diffusion = enable_multi_band_diffusion
+
+ # @markdown ---
+ # @markdown ## Conditional Generation Configs
+
+ # @markdown The prompt for generating audio. You can give multiple prompts separated by `|` in the input. You can also leave it blank for unconditional generation.
+ config.prompts = "happy rock | energetic EDM | sad jazz" # @param {type:"string"}
+
+ descriptions = [prompt.strip() for prompt in config.prompts.split("|")]
+ config.is_unconditional = config.prompts.strip() == ""
+
+ input_audio, input_sampling_rate, wandb_input_audio = None, None, None
+ if not config.is_unconditional:
+ input_audio_file = files.upload()
+ if input_audio_file != {}:
+ if config.model_name == "facebook/audiogen-medium":
+ error = f"{config.model_name} does not support audio-based conditioning"
+ raise ValueError(error)
+ wandb_input_audio = wandb.Audio(list(input_audio_file.keys())[0])
+ input_audio, input_sampling_rate = torchaudio.load(
+ list(input_audio_file.keys())[0]
+ )
+ config.input_audio_available = True
+ else:
+ config.input_audio_available = False
+ else:
+ if config.model_name == "facebook/audiogen-medium":
+ error = f"{config.model_name} does not support unconditional generration"
+ raise ValueError(error)
+
+ # @markdown Number of audio samples generated, this is relevant only for unconditional generation, i.e, if `config.prompts` is left blank.
+ config.num_samples = 4 # @param {type:"slider", min:1, max:10, step:1}
+
+ # @markdown Specify the random seed
+ seed = None # @param {type:"raw"}
+
+ max_seed = int(1024 * 1024 * 1024)
+ if not isinstance(seed, int):
+ seed = random.randint(1, max_seed)
+ if seed < 0:
+ seed = - seed
+ seed = seed % max_seed
+ config.seed = seed
+
+ # @markdown ---
+ # @markdown ## Generation Parameters
+ # @markdown Use sampling if True, else do argmax decoding
+ config.use_sampling = True # @param {type:"boolean"}
+
+ # @markdown `top_k` used for sampling; limits us to `k` number of of the top tokens to consider.
+ config.top_k = 250 # @param {type:"slider", min:0, max:1000, step:1}
+
+ # @markdown `top_p` used for sampling; limits us to the top tokens within a probability mass `p`
+ config.top_p = 0.0 # @param {type:"slider", min:0, max:1.0, step:0.01}
+
+ # @markdown Softmax temperature parameter
+ config.temperature = 1.0 # @param {type:"slider", min:0, max:1.0, step:0.01}
+
+ # @markdown Duration of the generated waveform
+ config.duration = 10 # @param {type:"slider", min:1, max:30, step:1}
+
+ # @markdown Coefficient used for classifier free guidance
+ config.cfg_coef = 3 # @param {type:"slider", min:1, max:100, step:1}
+
+ # @markdown Whether to perform 2 forward for Classifier Free Guidance instead of batching together the two. This has some impact on how things are padded but seems to have little impact in practice.
+ config.two_step_cfg = False # @param {type:"boolean"}
+
+ # @markdown When doing extended generation (i.e. more than 30 seconds), by how much should we extend the audio each time. Larger values will mean less context is preserved, and shorter value will require extra computations.
+ config.extend_stride = 0 # @param {type:"slider", min:0, max:30, step:1}
+ return (
+ config,
+ descriptions,
+ input_audio,
+ input_sampling_rate,
+ model_name,
+ wandb_input_audio,
+ )
+
+
+@app.cell
+def _(
+ AudioGen,
+ MultiBandDiffusion,
+ MusicGen,
+ config,
+ descriptions,
+ input_audio,
+ input_sampling_rate,
+):
+ # @title Generate Audio using MusicGen
+
+ # @markdown In this section, the audio is generated using the configs, specified in the aforementioned section. If you wish to peek behind the curtain and checkout the code, click on the `Show Code` button. In order to know about the different APIs for audio generation, visit the [official audiocraft documentations](https://facebookresearch.github.io/audiocraft/api_docs/audiocraft/index.html).
+
+ model = None
+ if config.model_name == "facebook/audiogen-medium":
+ model = AudioGen.get_pretrained(config.model_name)
+ elif "musicgen" in config.model_name:
+ model = MusicGen.get_pretrained(config.model_name.split("-")[-1])
+
+ multi_band_diffusion = None
+ if config.enable_multi_band_diffusion:
+ multi_band_diffusion = MultiBandDiffusion.get_mbd_musicgen()
+
+ model.set_generation_params(
+ use_sampling=config.use_sampling,
+ top_k=config.top_k,
+ top_p=config.top_p,
+ temperature=config.temperature,
+ duration=config.duration,
+ cfg_coef=config.cfg_coef,
+ two_step_cfg=config.two_step_cfg,
+ extend_stride=config.extend_stride
+ )
+
+ generated_wav, tokens = None, None
+ if config.is_unconditional:
+ if input_audio is None:
+ if "musicgen" in config.model_name:
+ generated_wav, tokens = model.generate_unconditional(
+ num_samples=config.num_samples,
+ progress=True,
+ return_tokens=True
+ )
+ else:
+ generated_wav = model.generate_unconditional(
+ num_samples=config.num_samples,
+ progress=True,
+ )
+ else:
+ if "musicgen" in config.model_name:
+ generated_wav, tokens = model.generate_with_chroma(
+ descriptions,
+ input_audio[None].expand(3, -1, -1),
+ input_sampling_rate,
+ return_tokens=True
+ )
+ else:
+ generated_wav = model.generate_with_chroma(
+ descriptions,
+ input_audio[None].expand(3, -1, -1),
+ input_sampling_rate,
+ )
+ else:
+ if "musicgen" in config.model_name:
+ generated_wav, tokens = model.generate(
+ descriptions,
+ progress=True,
+ return_tokens=True
+ )
+ else:
+ generated_wav = model.generate(
+ descriptions,
+ progress=True,
+ )
+
+ generated_wav_diffusion = None
+ if config.enable_multi_band_diffusion:
+ generated_wav_diffusion = multi_band_diffusion.tokens_to_wav(tokens)
+ return generated_wav, generated_wav_diffusion, model
+
+
+@app.cell
+def _(
+ TemporaryDirectory,
+ audio_write,
+ config,
+ descriptions,
+ generated_wav,
+ generated_wav_diffusion,
+ input_audio,
+ model,
+ model_name,
+ np,
+ os,
+ plt,
+ signal,
+ wandb,
+ wandb_input_audio,
+ wavfile,
+):
+ # @title Log Audio to Weights & Biases Dashboard
+
+ # @markdown In this section, we log the generated audio to Weights & Biases where you can listen and visualize them using an interactive audio player and waveform visualizer. Also, shoutout to [Atanu Sarkar](https://github.com/mratanusarkar) for building the spectrogram viusalization function which lets you visualize the spectrogram of the generated audio inside a [`wandb.Table`](https://docs.wandb.ai/guides/tables/tables-walkthrough).
+
+ def get_spectrogram(audio_file, output_file):
+ sample_rate, samples = wavfile.read(audio_file)
+ frequencies, times, Sxx = signal.spectrogram(samples, sample_rate)
+
+ log_Sxx = 10 * np.log10(Sxx + 1e-10)
+ vmin = np.percentile(log_Sxx, 5)
+ vmax = np.percentile(log_Sxx, 95)
+
+ mean_spectrum = np.mean(log_Sxx, axis=1)
+ threshold_low = np.percentile(mean_spectrum, 5)
+ threshold_high = np.percentile(mean_spectrum, 95)
+
+ freq_indices = np.where(mean_spectrum > threshold_low)
+ freq_min = 20
+ freq_max = frequencies[freq_indices].max()
+
+ fig, ax = plt.subplots()
+ cmap = plt.get_cmap('magma')
+
+ ax.pcolormesh(
+ times,
+ frequencies,
+ log_Sxx,
+ shading='gouraud',
+ cmap=cmap,
+ vmin=vmin,
+ vmax=vmax
+ )
+ ax.axis('off')
+ ax.set_ylim([freq_min, freq_max])
+
+ plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+ plt.savefig(
+ output_file, format='png', bbox_inches='tight', pad_inches=0
+ )
+ plt.close()
+
+ return wandb.Image(output_file)
+
+
+ temp_dir = TemporaryDirectory()
+ columns = ["Model", "Prompt", "Generated-Audio", "Spectrogram", "Seed"]
+ if input_audio is not None:
+ columns.insert(2, "Input-Audio")
+ if config.enable_multi_band_diffusion:
+ columns.insert(4, "Generated-Audio-Diffusion")
+ columns.insert(5, "Spectrogram-Diffusion")
+ wandb_table = wandb.Table(columns=columns)
+
+ for idx, wav in enumerate(generated_wav):
+
+ file_name = os.path.join(temp_dir.name, str(idx))
+ audio_write(
+ file_name,
+ wav.cpu(),
+ model.sample_rate,
+ strategy="loudness",
+ loudness_compressor=True,
+ )
+ wandb_audio = wandb.Audio(file_name + ".wav")
+ wandb.log({"Generated-Audio": wandb_audio}, commit=False)
+
+ file_name_diffusion, wandb_diffusion_audio = None, None
+ if config.enable_multi_band_diffusion:
+ file_name_diffusion = os.path.join(
+ temp_dir.name, str(idx) + "_diffusion"
+ )
+ audio_write(
+ file_name_diffusion,
+ generated_wav_diffusion[idx].cpu(),
+ model.sample_rate,
+ strategy="loudness",
+ loudness_compressor=True,
+ )
+ wandb_diffusion_audio = wandb.Audio(file_name_diffusion + ".wav")
+ wandb.log(
+ {"Generated-Audio-Diffusion": wandb_diffusion_audio},
+ commit=False
+ )
+
+ wandb.log({}, commit=True)
+
+ desc = descriptions[idx] if len(descriptions) > 1 else config.prompts
+ wandb_table_row = [
+ model_name,
+ desc,
+ wandb_audio,
+ get_spectrogram(
+ audio_file=file_name + ".wav",
+ output_file=os.path.join(temp_dir.name, str(idx) + ".png")
+ ),
+ config.seed
+ ]
+ if input_audio is not None:
+ wandb_table_row.insert(2, wandb_input_audio)
+ if config.enable_multi_band_diffusion:
+ wandb_table_row.insert(4, wandb_diffusion_audio)
+ wandb_table_row.insert(
+ 5,
+ get_spectrogram(
+ audio_file=file_name_diffusion + ".wav",
+ output_file=os.path.join(
+ temp_dir.name, str(idx) + "_diffusion.png"
+ )
+ )
+ )
+ wandb_table.add_data(*wandb_table_row)
+
+ wandb.log({"Generated-Audio-Table": wandb_table})
+
+ wandb.finish()
+ temp_dir.cleanup()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This is how the W&B Table looks like with the interactive audio player, waveform visualizer and spectrogram visualization along with the prompts and other configs. Note that the notebook automatically sets the seed if you leave it blank, so your experiments are always reproducible.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If you want to know more about the underlying architectures for MusicGen and AudioGen and explore some cool audio samples generated by these models, you can check out [this W&B report](http://wandb.me/audiocraft_2mp).
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/azure-azure-gpt-medical-notes/azure_azure_gpt_medical_notes.py b/marimo/convert/azure-azure-gpt-medical-notes/azure_azure_gpt_medical_notes.py
new file mode 100644
index 00000000..1aab0b5f
--- /dev/null
+++ b/marimo/convert/azure-azure-gpt-medical-notes/azure_azure_gpt_medical_notes.py
@@ -0,0 +1,312 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ from typing import Dict, List, Literal, Optional, Tuple
+
+ import instructor
+ import openai
+ import pandas as pd
+ import weave
+ from pydantic import BaseModel, Field
+ from set_env import set_env
+ import json
+ import asyncio
+
+ return Dict, Tuple, asyncio, json, openai, pd, set_env, weave
+
+
+@app.cell
+def _(set_env):
+ set_env("OPENAI_API_KEY")
+ set_env("WANDB_API_KEY")
+ set_env("AZURE_OPENAI_ENDPOINT")
+ set_env("AZURE_OPENAI_API_KEY")
+ print("Env set")
+ return
+
+
+@app.cell
+def _():
+ from utils.config import ENTITY, WEAVE_PROJECT
+
+ return ENTITY, WEAVE_PROJECT
+
+
+@app.cell
+def _(ENTITY, WEAVE_PROJECT, weave):
+ weave.init(f"{ENTITY}/{WEAVE_PROJECT}")
+ return
+
+
+@app.cell
+def _():
+ N_SAMPLES = 67
+ return (N_SAMPLES,)
+
+
+@app.cell
+def _(openai):
+ client = openai.OpenAI()
+ return (client,)
+
+
+@app.cell
+def _(N_SAMPLES, Tuple, pd):
+ def load_medical_data(url: str, num_samples: int = N_SAMPLES) -> Tuple[pd.DataFrame, pd.DataFrame]:
+ """
+ Load medical data and split into train and test sets
+
+ Args:
+ url: URL of the CSV file
+ num_samples: Number of samples to load
+
+ Returns:
+ Tuple of (train_df, test_df)
+ """
+ df = pd.read_csv(url)
+ df = df.sample(n=num_samples, random_state=42) # Sample and shuffle data
+
+ # Split into 80% train, 20% test
+ train_size = int(0.8 * len(df))
+ train_df = df[:train_size]
+ test_df = df[train_size:]
+
+ return train_df, test_df
+
+ return (load_medical_data,)
+
+
+@app.cell
+def _():
+ medical_dataset_url = "https://raw.githubusercontent.com/wyim/aci-bench/main/data/challenge_data/train.csv"
+ return (medical_dataset_url,)
+
+
+@app.cell
+def _(load_medical_data, medical_dataset_url):
+ train_df, test_df = load_medical_data(medical_dataset_url)
+ train_samples = train_df.to_dict("records")
+ test_samples = test_df.to_dict("records")
+ return test_df, test_samples, train_df, train_samples
+
+
+@app.cell
+def _(train_samples):
+ train_samples[0]
+ return
+
+
+@app.cell
+def _(test_samples):
+ test_samples[0]
+ return
+
+
+@app.cell
+def _(json, pd):
+ def convert_to_jsonl(df: pd.DataFrame, output_file: str = "medical_conversations.jsonl"):
+ """
+ Convert medical dataset to JSONL format with conversation structure
+
+ Args:
+ df: DataFrame to convert
+ output_file: Output JSONL filename
+ """
+
+ with open(output_file, 'w', encoding='utf-8') as f:
+ for _, row in df.iterrows():
+ # Create the conversation structure
+ conversation = {
+ "messages": [
+ {
+ "role": "system",
+ "content": "You are a medical scribe assistant. Your task is to accurately document medical conversations between doctors and patients, creating detailed medical notes that capture all relevant clinical information."
+ },
+ {
+ "role": "user",
+ "content": row['dialogue']
+ },
+ {
+ "role": "assistant",
+ "content": row['note']
+ }
+ ]
+ }
+
+ # Write as JSON line
+ json_line = json.dumps(conversation, ensure_ascii=False)
+ f.write(json_line + '\n')
+
+ print(f"Converted {len(df)} records to {output_file}")
+
+ return (convert_to_jsonl,)
+
+
+@app.cell
+def _(convert_to_jsonl, test_df, train_df):
+ convert_to_jsonl(train_df, "medical_conversations_train.jsonl")
+ convert_to_jsonl(test_df, "medical_conversations_test.jsonl")
+ return
+
+
+@app.cell
+def _():
+ from utils.prompts import medical_task, medical_system_prompt
+
+ return medical_system_prompt, medical_task
+
+
+@app.cell
+def _(Dict, client, medical_system_prompt, medical_task, weave):
+ def format_dialogue(dialogue: str):
+ dialogue = dialogue.replace("\n", " ")
+ transcript = f"Dialogue: {dialogue}"
+ return transcript
+
+
+ @weave.op()
+ def process_medical_record(dialogue: str) -> Dict:
+ transcript = format_dialogue(dialogue)
+ prompt = medical_task.format(transcript=transcript)
+
+ response = client.chat.completions.create(
+ model="gpt-3.5-turbo",
+ messages=[
+ {"role": "system", "content": medical_system_prompt},
+ {"role": "user", "content": prompt},
+ ],
+ )
+
+ extracted_info = response.choices[0].message.content
+
+ return {
+ "input": transcript,
+ "output": extracted_info,
+ }
+
+ return (process_medical_record,)
+
+
+@app.cell
+def _(client, json, weave):
+ # Define the LLM scoring function
+ @weave.op()
+ async def medical_note_accuracy(note: str, output: dict) -> dict:
+ scoring_prompt = """Compare the generated medical note with the ground truth note and evaluate accuracy.
+ Score as 1 if the generated note captures the key medical information accurately, 0 if not.
+ Output in valid JSON format with just a "score" field.
+
+ Ground Truth Note:
+ {ground_truth}
+
+ Generated Note:
+ {generated}"""
+
+ prompt = scoring_prompt.format(
+ ground_truth=note,
+ generated=output['output']
+ )
+
+ response = client.chat.completions.create(
+ model="gpt-4o",
+ messages=[{"role": "user", "content": prompt}],
+ response_format={ "type": "json_object" }
+ )
+ return json.loads(response.choices[0].message.content)
+
+ return (medical_note_accuracy,)
+
+
+@app.cell
+def _(medical_note_accuracy, test_samples, weave):
+ # Create evaluation for test samples
+ test_evaluation = weave.Evaluation(
+ name='medical_record_extraction_test',
+ dataset=test_samples,
+ scorers=[medical_note_accuracy]
+ )
+ return (test_evaluation,)
+
+
+@app.cell
+def _():
+ try:
+ in_jupyter = True
+ except ImportError:
+ in_jupyter = False
+ if in_jupyter:
+ import nest_asyncio
+
+ nest_asyncio.apply()
+ return
+
+
+@app.cell
+def _(asyncio, process_medical_record, test_evaluation):
+ test_results = asyncio.run(test_evaluation.evaluate(process_medical_record))
+ print(f"Completed test evaluation")
+ return
+
+
+@app.cell
+def _(Dict, weave):
+ import os
+ from openai import AzureOpenAI
+
+ # Initialize Azure client
+ azure_client = AzureOpenAI(
+ azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
+ api_key=os.getenv("AZURE_OPENAI_API_KEY"),
+ api_version="2024-02-01"
+ )
+
+ @weave.op()
+ def process_medical_record_azure(dialogue: str) -> Dict:
+
+ response = azure_client.chat.completions.create(
+ model="gpt-35-turbo-0125-ft-d30b3aee14864c29acd9ac54eb92457f",
+ messages=[
+ {"role": "system", "content": "You are a medical scribe assistant. Your task is to accurately document medical conversations between doctors and patients, creating detailed medical notes that capture all relevant clinical information."},
+ {"role": "user", "content": dialogue},
+ ],
+ )
+
+ extracted_info = response.choices[0].message.content
+
+ return {
+ "input": dialogue,
+ "output": extracted_info,
+ }
+
+ return (process_medical_record_azure,)
+
+
+@app.cell
+def _(asyncio, process_medical_record_azure, test_evaluation):
+ test_results_azure = asyncio.run(test_evaluation.evaluate(process_medical_record_azure))
+ print(f"Completed test evaluation")
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/azure-azure-weave-cookbook-colab/azure_azure_weave_cookbook_colab.py b/marimo/convert/azure-azure-weave-cookbook-colab/azure_azure_weave_cookbook_colab.py
new file mode 100644
index 00000000..fbcd6718
--- /dev/null
+++ b/marimo/convert/azure-azure-weave-cookbook-colab/azure_azure_weave_cookbook_colab.py
@@ -0,0 +1,429 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install weave openai
+ return
+
+
+@app.cell
+def _():
+ model_id = "gpt-4-turbo" # @param {type:"string"}
+ model_id = "mistral-7b-instruct-weave"
+ azure_model_option = "openai" # @param ["openai", "ai_studio"]
+ return azure_model_option, model_id
+
+
+@app.cell
+def _():
+ wandb_entity = "a-sh0ts" # @param {type:"string"}
+ weave_project = "azure-weave-cookbook" # @param {type:"string"}
+ eval_dataset_name = "customer_service_inquiries" # @param {type:"string"}
+ publish_eval_data = True # @param {type:"boolean"}
+ return eval_dataset_name, publish_eval_data, wandb_entity, weave_project
+
+
+@app.cell
+def _(azure_model_option):
+ from google.colab import userdata
+ import os
+ from openai import AzureOpenAI, OpenAI
+
+ os.environ["WANDB_API_KEY"] = userdata.get('WANDB_API_KEY')
+
+ if azure_model_option == "openai":
+ os.environ["AZURE_OPENAI_ENDPOINT"] = userdata.get('AZURE_OPENAI_ENDPOINT')
+ os.environ["AZURE_OPENAI_API_KEY"] = userdata.get('AZURE_OPENAI_API_KEY')
+ client = AzureOpenAI(
+ api_key=os.getenv("AZURE_API_KEY"),
+ api_version="2024-02-01",
+ azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+ )
+ elif azure_model_option == "ai_studio":
+ os.environ["AZURE_AI_STUDIO_API_ENDPOINT"] = userdata.get('AZURE_AI_STUDIO_API_ENDPOINT')
+ os.environ["AZURE_AI_STUDIO_API_KEY"] = userdata.get('AZURE_AI_STUDIO_API_KEY')
+
+ api_version = "v1"
+ client = OpenAI(
+ base_url=f"{os.getenv('AZURE_AI_STUDIO_API_ENDPOINT')}/v1",
+ api_key=os.getenv('AZURE_AI_STUDIO_API_KEY')
+ )
+ else:
+ print("Please us one of the above options")
+ return (client,)
+
+
+@app.cell
+def _(weave_project):
+ import weave
+ weave.init(weave_project)
+ return (weave,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Calling Azure directly
+ """)
+ return
+
+
+@app.cell
+def _(client, weave):
+ @weave.op()
+ def call_azure_chat(model_id: str, messages: list, max_tokens: int = 1000, temperature: float = 0.5):
+ response = client.chat.completions.create(
+ model=model_id,
+ messages=messages,
+ max_tokens=max_tokens,
+ temperature=temperature
+ )
+ return {"status": "success", "response": response.choices[0].message.content}
+
+ return (call_azure_chat,)
+
+
+@app.cell
+def _(weave):
+ @weave.op()
+ def format_messages_for_mistral(messages: list):
+ system_message = messages[0]["content"]
+ formatted_messages = []
+
+ for message in messages[1:]:
+ if message["role"] == "user":
+ formatted_message = {
+ "role": "user",
+ "content": f"[INST]\n{system_message}\n{message['content']}\n[/INST]"
+ }
+ else:
+ formatted_message = {
+ "role": message["role"],
+ "content": message["content"]
+ }
+ formatted_messages.append(formatted_message)
+
+ return formatted_messages
+
+ return (format_messages_for_mistral,)
+
+
+@app.cell
+def _(call_azure_chat, format_messages_for_mistral, model_id):
+ messages = [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Create a snack recipe for a dish called the Azure Weav-e-ohs"}
+ ]
+ if "mistral" in model_id.lower():
+ messages = format_messages_for_mistral(messages)
+ result = call_azure_chat(model_id, messages)
+ print(result)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Creating Functional LLM Apps
+ """)
+ return
+
+
+@app.cell
+def _(weave):
+ @weave.op()
+ def format_prompt(prompt: str):
+ "A formatting function for OpenAI models"
+ system_prompt_formatted = "You are a helpful assistant."
+
+ human_prompt = """
+ {prompt}
+ """
+
+ human_prompt_formatted = human_prompt.format(prompt=prompt)
+ messages = [{"role":"system", "content":system_prompt_formatted}, {"role":"user", "content":human_prompt_formatted}]
+ return messages
+
+ return (format_prompt,)
+
+
+@app.cell
+def _(
+ call_azure_chat,
+ format_messages_for_mistral,
+ format_prompt,
+ model,
+ weave,
+):
+ @weave.op()
+ def run_chat(model_id: str, prompt: str):
+ formatted_messages = format_prompt(prompt=prompt)
+ if "mistral" in model.lower():
+ formatted_messages = format_messages_for_mistral(formatted_messages)
+ result = call_azure_chat(model_id, formatted_messages, max_tokens=1000)
+ return result
+
+ return (run_chat,)
+
+
+@app.cell
+def _():
+ prompt = "Give a full recipe for a Weights & Biases inspired cocktail. Ensure you provide a list of ingredients, tools, and step by step instructions"
+ return (prompt,)
+
+
+@app.cell
+def _(model_id, prompt, run_chat):
+ result_1 = run_chat(model_id, prompt)
+ return (result_1,)
+
+
+@app.cell
+def _(result_1):
+ result_1['response']
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Create LLM Model Classes to iterate over hyperparameters
+ """)
+ return
+
+
+@app.cell
+def _(weave):
+ from dataclasses import dataclass
+
+ @dataclass
+ class PromptTemplate:
+ system_prompt: str
+ human_prompt: str
+
+ @weave.op()
+ def format_prompt(self, email_content: str):
+ "A formatting function for OpenAI models"
+ system_prompt_formatted = self.system_prompt.format()
+ human_prompt_formatted = self.human_prompt.format(email_content=email_content)
+ messages = [{"role":"system", "content":system_prompt_formatted}, {"role":"user", "content":human_prompt_formatted}]
+ return messages
+
+ return (PromptTemplate,)
+
+
+@app.cell
+def _(
+ PromptTemplate,
+ call_azure_chat,
+ format_messages_for_mistral,
+ model_id,
+ weave,
+):
+ from weave import Model
+ from typing import Tuple
+
+ class AzureEmailAssistant(Model):
+ model_id: str = model_id
+ prompt_template: PromptTemplate
+ max_tokens: int = 2048
+ temperature: float = 0.0
+
+ @weave.op()
+ def format_doc(self, doc: str) -> list:
+ "Read and format the document"
+ messages = self.prompt_template.format_prompt(doc)
+ return messages
+
+ @weave.op()
+ def respond(self, doc: str) -> dict:
+ "Generate a response to the email inquiry"
+ messages = self.format_doc(doc)
+ if "mistral" in self.model_id.lower():
+ messages = format_messages_for_mistral(messages)
+ output = call_azure_chat(
+ self.model_id,
+ messages=messages,
+ max_tokens=self.max_tokens,
+ temperature=self.temperature)
+ return output
+
+ @weave.op()
+ async def predict(self, email_content: str) -> str:
+ return self.respond(email_content)["response"]
+
+ return (AzureEmailAssistant,)
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%writefile customer_inquiry.txt
+ # Subject: Inquiry about Order Delay
+ #
+ # Hello,
+ #
+ # I placed an order last week for the new UltraGlow Skin Serum, but I have not received a shipping update yet. My order number is 12345. Could you please update me on the status of my shipment?
+ #
+ # Thank you,
+ # Jane Doe
+ return
+
+
+@app.cell
+def _():
+ system_prompt = """
+ # Instructions
+ You are a customer service response assistant. Our goal is to provide clear, concise, and polite responses to customer inquiries about products, shipping, and any issues they may have encountered. Some rules to remember:
+ - Always be courteous and respectful.
+ - Provide accurate and helpful information.
+ - Responses should be concise and to the point.
+ - Use formal language suitable for professional communication.
+ ## Formatting Rules
+ Maintain a formal greeting and closing in each response. Do not use slang or overly casual language. Ensure all provided information is correct and double-check for typographical errors.
+ """
+
+ human_prompt = """
+ Here is a customer inquiry received via email. Craft a suitable response based on the guidelines provided:
+
+ {email_content}
+
+ """
+ return human_prompt, system_prompt
+
+
+@app.cell
+def _(PromptTemplate, human_prompt, system_prompt):
+ prompt_template = PromptTemplate(
+ system_prompt=system_prompt,
+ human_prompt=human_prompt)
+ return (prompt_template,)
+
+
+@app.cell
+def _(weave, weave_project):
+ weave.init(weave_project) # Colab specific
+ return
+
+
+@app.cell
+def _():
+ from pathlib import Path
+
+ return (Path,)
+
+
+@app.cell
+def _(AzureEmailAssistant, Path, model_id, prompt_template):
+ doc = Path('customer_inquiry.txt').read_text()
+ model = AzureEmailAssistant(model_id=model_id, prompt_template=prompt_template)
+ response = model.respond(doc)
+ return model, response
+
+
+@app.cell
+def _(response):
+ print(response["response"])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## [Optional] Publish synthetically generated Evaluation data to Weave
+ """)
+ return
+
+
+@app.cell
+def _(eval_dataset_name, publish_eval_data, weave):
+ if publish_eval_data:
+ from weave import Dataset
+ dataset = Dataset(name=eval_dataset_name, rows=[
+ {'id': '1', 'email_content': 'Subject: Inquiry about Order Delay\n\nHello,\n\nI placed an order last week for the new UltraGlow Skin Serum, but I have not received a shipping update yet. My order number is 12345. Could you please update me on the status of my shipment?\n\nThank you,\nJane Doe'},
+ {'id': '2', 'email_content': 'Subject: Damaged Item Received\n\nHello,\n\nI received my order yesterday, but one of the items, a glass vase, was broken. My order number is 67890. How can I get a replacement or a refund?\n\nBest regards,\nJohn Smith'},
+ {'id': '3', 'email_content': 'Subject: Wrong Item Delivered\n\nHi,\n\nI ordered a pair of blue sneakers, but I received a black pair instead. My order number is 54321. Could you please assist me with this issue?\n\nThank you,\nEmily Johnson'},
+ {'id': '4', 'email_content': 'Subject: Request for Return Instructions\n\nDear Customer Service,\n\nI would like to return a dress I purchased last week as it does not fit well. My order number is 11223. Could you please provide the return instructions?\n\nSincerely,\nLaura Davis'},
+ {'id': '5', 'email_content': 'Subject: Missing Items in Order\n\nHello,\n\nI just received my order, but two items are missing. My order number is 33445. Could you please help me resolve this?\n\nKind regards,\nMichael Brown'},
+ {'id': '6', 'email_content': 'Subject: Delay in Order Confirmation\n\nDear Support Team,\n\nI placed an order two days ago but have not received a confirmation email yet. My order number is 99887. Can you confirm if my order was processed?\n\nThank you,\nSarah Wilson'},
+ {'id': '7', 'email_content': 'Subject: Inquiry About Product Availability\n\nHi,\n\nI\'m interested in purchasing the Professional Chef Knife Set, but it appears to be out of stock. Can you let me know when it will be available again?\n\nBest regards,\nDavid Martinez'},
+ {'id': '8', 'email_content': 'Subject: Request for Invoice\n\nDear Customer Service,\n\nCould you please send me an invoice for my recent purchase? My order number is 55667. I need it for my records.\n\nThank you,\nJessica Taylor'},
+ {'id': '9', 'email_content': 'Subject: Issue with Discount Code\n\nHello,\n\nI tried using the discount code SAVE20 during checkout, but it did not apply. My order number is 77654. Could you please assist me?\n\nSincerely,\nRobert Anderson'},
+ {'id': '10', 'email_content': 'Subject: Request for Expedited Shipping\n\nHi,\n\nI need my order delivered urgently. Is it possible to upgrade to expedited shipping? My order number is 44556.\n\nThank you,\nLinda Thompson'},
+ {'id': '11', 'email_content': 'Subject: Order Cancellation Request\n\nDear Support Team,\n\nI would like to cancel my recent order as I made a mistake while ordering. My order number is 33221. Can you please process the cancellation?\n\nBest regards,\nWilliam Clark'}
+ ])
+ # Publish the dataset
+ weave.publish(dataset)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Run Evaluation whilst logging results to Weave
+ """)
+ return
+
+
+@app.cell
+def _(eval_dataset_name, wandb_entity, weave, weave_project):
+ dataset_uri = f'weave:///{wandb_entity}/{weave_project}/object/{eval_dataset_name}:latest'
+ dataset_1 = weave.ref(dataset_uri).get()
+ return (dataset_1,)
+
+
+@app.cell
+def _(weave):
+ # Scoring function checking length of summary
+ @weave.op()
+ def check_conciseness(model_output: str) -> dict:
+ result = len(model_output.split()) < 300
+ return {'conciseness': result}
+
+ return (check_conciseness,)
+
+
+@app.cell
+def _(check_conciseness, dataset_1, weave):
+ evaluation = weave.Evaluation(dataset=dataset_1, scorers=[check_conciseness])
+ return (evaluation,)
+
+
+@app.cell
+async def _(evaluation, model):
+ await evaluation.evaluate(model)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/boosting-credit-scorecards-with-xgboost-and-w-b/boosting_credit_scorecards_with_xgboost_and_w_b.py b/marimo/convert/boosting-credit-scorecards-with-xgboost-and-w-b/boosting_credit_scorecards_with_xgboost_and_w_b.py
new file mode 100644
index 00000000..3ddd7663
--- /dev/null
+++ b/marimo/convert/boosting-credit-scorecards-with-xgboost-and-w-b/boosting_credit_scorecards_with_xgboost_and_w_b.py
@@ -0,0 +1,855 @@
+# /// script
+# dependencies = ["dill", "scikit-learn", "wandb", "xgboost"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Vehicle Loan Default Prediction with XGBoost
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this notebook we'll train a XGBoost model to classify whether submitted loan applications will default or not. Using boosting algorithms such as XGBoost increases the performance of a loan assesment, whilst retaining interpretability for internal Risk Management functions as well as external regulators.
+
+ This notebook is based on a talk from Nvidia GTC21 by Paul Edwards at ScotiaBank who [presented](https://www.nvidia.com/en-us/on-demand/session/gtcspring21-s31327/) how XGBoost can be used to construct more performant credit scorecards that remain interpretable. They also kindly [shared sample code](https://github.com/rapidsai-community/showcase/tree/main/event_notebooks/GTC_2021/credit_scorecard) which we will use throughout this notebook, credit to [Stephen Denton](stephen.denton@scotiabank.com) from Scotiabank for sharing this code publicly.
+
+ ### [Click here](https://wandb.ai/morgan/credit_scorecard) to view and interact with a live W&B Dashboard built with this notebook
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # In this notebook
+
+ In this colab we'll cover how Weights and Biases enables regulated entities to
+ - **Track and version** their data ETL pipelines (locally or in cloud services such as S3 and GCS)
+ - **Track experiment results** and store trained models
+ - **Visually inspect** multiple evaluation metrics
+ - **Optimize performance** with hyperparameter sweeps
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Track Experiments and Results**
+
+ We will track all of the training hyperparameters and output metrics in order to generate an Experiments Dashboard like the one below:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Run a Hyperparameter Sweep to Find the Best HyperParameters**
+
+ Weights and Biases also enables you to do hyperparameter sweeps, either with our own [Sweeps functionality](https://docs.wandb.ai/guides/sweeps) or with our [Ray Tune integration](https://docs.wandb.ai/guides/sweeps/advanced-sweeps/ray-tune). See our docs for a full guide of how to use more advanced hyperparameter sweeps options.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb>=0.13.10 dill !pip install -qq "wandb>=0.13.10" dill
+ # packages added via marimo's package management: xgboost>=2.0.0 scikit-learn>=1.2.1 !pip install -qq "xgboost>=2.0.0" "scikit-learn>=1.2.1"
+ return
+
+
+@app.cell
+def _():
+ import ast
+ import sys
+ import json
+ from pathlib import Path
+ from dill.source import getsource
+ from dill import detect
+
+ import pandas as pd
+ import numpy as np
+ import plotly
+ import matplotlib.pyplot as plt
+
+ from scipy.stats import ks_2samp
+ from sklearn import metrics
+ from sklearn import model_selection
+ import xgboost as xgb
+
+ pd.set_option('display.max_columns', None)
+ return (
+ Path,
+ detect,
+ getsource,
+ ks_2samp,
+ metrics,
+ model_selection,
+ np,
+ pd,
+ sys,
+ xgb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Data
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## AWS S3, Google Cloud Storage and W&B Artifacts
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Weights and Biases **Artifacts** enable you to log end-to-end training pipelines to ensure your experiments are always reproducible.
+
+ Data privacy is critical to Weights & Biases and so we support the creation of Artifacts from reference locations such as your own private cloud such as AWS S3 or Google Cloud Storage. Local, on-premises of W&B are also available upon request.
+
+ 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.
+
+ ## Artifacts Reference Example
+ **Create an artifact with the S3/GCS metadata**
+
+ 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 [Artifact References](https://docs.wandb.ai/guides/artifacts/references) for more on how to use Artifacts by reference, credentials setup etc.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Login to W&B
+ Login to Weights and Biases
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ WANDB_PROJECT ='vehicle_loan_default'
+ return WANDB_PROJECT, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Vehicle Loan Dataset
+
+ We will be using a simplified version of the [Vehicle Loan Default Prediction dataset](https://www.kaggle.com/sneharshinde/ltfs-av-data) from L&T which has been stored in W&B Artifacts.
+ """)
+ return
+
+
+@app.cell
+def _(Path, sys):
+ # specify a folder to save the data, a new folder will be created if it doesn't exist
+ data_dir = Path(sys.path[0]) # get this notebook path
+ model_dir = data_dir / 'models'
+ model_dir.mkdir(exist_ok=True)
+
+ id_vars = ['UniqueID']
+ targ_var = 'loan_default'
+ return data_dir, id_vars, targ_var
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Create function to pickle functions
+ """)
+ return
+
+
+@app.cell
+def _(detect, getsource):
+ def function_to_string(fn):
+ return getsource(detect.code(fn))
+
+ return (function_to_string,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Download Data from W&B Artifacts
+
+ We will download our dataset from W&B Artifacts. First we need to create a W&B run object, which we will use to download the data. Once the data is downloaded it will be one-hot encoded. This processed data will then be logged to the same W&B as a new Artifact. By logging to the W&B that downloaded the data, we tie this new Artifact to the raw dataset Artifact
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, wandb):
+ run = wandb.init(project=WANDB_PROJECT, job_type='preprocess-data')
+ return (run,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Download the subset of the vehicle loan default data from W&B, this contains `train.csv` and `val.csv` files as well as some utils files.
+ """)
+ return
+
+
+@app.cell
+def _(data_dir, run):
+ ARTIFACT_PATH = 'morgan/credit_scorecard/vehicle_loan_defaults:latest'
+ _dataset_art = run.use_artifact(ARTIFACT_PATH, type='dataset')
+ _dataset_dir = _dataset_art.download(data_dir)
+ return
+
+
+@app.cell
+def _():
+ from data_utils import (
+ describe_data_g_targ,
+ one_hot_encode_data,
+ )
+
+ return describe_data_g_targ, one_hot_encode_data
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### One-Hot Encode the Data
+ """)
+ return
+
+
+@app.cell
+def _(data_dir, id_vars, one_hot_encode_data, pd, targ_var):
+ # Load data into Dataframe
+ _dataset = pd.read_csv(data_dir / 'vehicle_loans_subset.csv')
+ _dataset, p_vars = one_hot_encode_data(_dataset, id_vars, targ_var)
+ # One Hot Encode Data
+ processed_data_path = data_dir / 'proc_ds.csv'
+ # Save Preprocessed data
+ _dataset.to_csv(processed_data_path, index=False)
+ return p_vars, processed_data_path
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Log Processed Data to W&B Artifacts
+ """)
+ return
+
+
+@app.cell
+def _(
+ function_to_string,
+ one_hot_encode_data,
+ processed_data_path,
+ run,
+ wandb,
+):
+ # Create a new artifact for the processed data, including the function that created it, to Artifacts
+ processed_ds_art = wandb.Artifact(name='vehicle_defaults_processed',
+ type='processed_dataset',
+ description='One-hot encoded dataset',
+ metadata={'preprocessing_fn': function_to_string(one_hot_encode_data)}
+ )
+
+ # Attach our processed data to the Artifact
+ processed_ds_art.add_file(processed_data_path)
+
+ # Log this Artifact to the current wandb run
+ run.log_artifact(processed_ds_art);
+
+ run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Get Train/Validation Split
+
+ Here we show an alternative pattern for how to create a wandb run object. In the cell below, the code to split the dataset is wrapped with a call to `wandb.init() as run`.
+
+ Here we will:
+
+ - Start a wandb run
+ - Download our one-hot-encoded dataset from Artifacts
+ - Do the Train/Val split and log the params used in the split
+ - Log the new `trndat` and `valdat` datasets to Artifacts
+ - Finish the wandb run automatically
+ """)
+ return
+
+
+@app.cell
+def _(
+ WANDB_PROJECT,
+ data_dir,
+ model_selection,
+ pd,
+ processed_data_path,
+ targ_var,
+ wandb,
+):
+ with wandb.init(project=WANDB_PROJECT, job_type='train-val-split') as run_1:
+ _dataset_art = run_1.use_artifact('vehicle_defaults_processed:latest', type='processed_dataset')
+ _dataset_dir = _dataset_art.download(data_dir)
+ _dataset = pd.read_csv(processed_data_path)
+ test_size = 0.25
+ random_state = 42
+ run_1.config.update({'test_size': test_size, 'random_state': random_state})
+ trndat, valdat = model_selection.train_test_split(_dataset, test_size=test_size, random_state=random_state, stratify=_dataset[[targ_var]])
+ print(f'Train dataset size: {trndat[targ_var].value_counts()} \n')
+ print(f'Validation dataset sizeL {valdat[targ_var].value_counts()}')
+ train_path = data_dir / 'train.csv'
+ val_path = data_dir / 'val.csv'
+ trndat.to_csv(train_path, index=False)
+ valdat.to_csv(val_path, index=False)
+ split_ds_art = wandb.Artifact(name='vehicle_defaults_split', type='train-val-dataset', description='Processed dataset split into train and valiation', metadata={'test_size': test_size, 'random_state': random_state})
+ split_ds_art.add_file(train_path)
+ split_ds_art.add_file(val_path)
+ run_1.log_artifact(split_ds_art)
+ return trndat, valdat
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Inspect Training Dataset
+ Get an overview of the training dataset
+ """)
+ return
+
+
+@app.cell
+def _(describe_data_g_targ, targ_var, trndat):
+ trndict = describe_data_g_targ(trndat, targ_var)
+ trndat.head()
+ return (trndict,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Log Dataset with W&B Tables
+
+ With W&B Tables you can log, query, and analyze tabular data that contains rich media such as images, video, audio and more. With it you can understand your datasets, visualize model predictions, and share insights, for more see more in our [W&B Tables Guide](https://docs.wandb.ai/guides/data-vis)
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, trndat, wandb):
+ # Create a wandb run, with an optional "log-dataset" job type to keep things tidy
+ run_2 = wandb.init(project=WANDB_PROJECT, job_type='log-dataset') # config is optional here
+ table = wandb.Table(dataframe=trndat.sample(1000))
+ # Create a W&B Table and log 1000 random rows of the dataset to explore
+ wandb.log({'processed_dataset': table})
+ # Log the Table to your W&B workspace
+ # Close the wandb run
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Modelling
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Fit the XGBoost Model
+
+ We will now fit an XGBoost model to classify whether a vehicle loan application will result in a default or not
+
+ ### Training on GPU
+ If you'd like to train your XGBoost model on your GPU, simply change set the following in the parameters you pass to XGBoost:
+
+ ```
+ 'tree_method': 'gpu_hist'
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### 1) Initialise a W&B Run
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, wandb):
+ run_3 = wandb.init(project=WANDB_PROJECT, job_type='train-model')
+ return (run_3,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### 2) Setup and Log the Model Parameters
+ """)
+ return
+
+
+@app.cell
+def _(trndict):
+ base_rate = round(trndict['base_rate'], 6)
+ early_stopping_rounds = 40
+ return base_rate, early_stopping_rounds
+
+
+@app.cell
+def _(base_rate):
+ bst_params = {
+ 'objective': 'binary:logistic'
+ , 'base_score': base_rate
+ , 'gamma': 1 ## def: 0
+ , 'learning_rate': 0.1 ## def: 0.1
+ , 'max_depth': 3
+ , 'min_child_weight': 100 ## def: 1
+ , 'n_estimators': 25
+ , 'nthread': 24
+ , 'random_state': 42
+ , 'reg_alpha': 0
+ , 'reg_lambda': 0 ## def: 1
+ , 'eval_metric': ['auc', 'logloss']
+ , 'tree_method': 'hist' # use `gpu_hist` to train on GPU
+ }
+ return (bst_params,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Log the xgboost training parameters to the W&B run config
+ """)
+ return
+
+
+@app.cell
+def _(bst_params, early_stopping_rounds, run_3):
+ run_3.config.update(dict(bst_params))
+ run_3.config.update({'early_stopping_rounds': early_stopping_rounds})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### 3) Let's select the data for train/validation
+ """)
+ return
+
+
+@app.cell
+def _(data_dir):
+ data_dir
+ return
+
+
+@app.cell
+def _(valdat):
+ valdat
+ return
+
+
+@app.cell
+def _(targ_var, trndat, valdat):
+ ## Extract target column as a series
+ y_trn = trndat.loc[:,targ_var].astype(int)
+ y_val = valdat.loc[:,targ_var].astype(int)
+ return y_trn, y_val
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### 4) Fit the model, log results to W&B and save model to W&B Artifacts
+
+ To log all our xgboost model parameters we used the `WandbCallback`. This will . See the [W&B docs](https://docs.wandb.ai/guides/integrations), including documentation for other libraries that have integrated W&B including LightGBM and more.
+ """)
+ return
+
+
+@app.cell
+def _(bst_params, p_vars, run_3, trndat, valdat, xgb, y_trn, y_val):
+ from wandb.integration.xgboost import WandbCallback
+ xgbmodel = xgb.XGBClassifier(**bst_params, callbacks=[WandbCallback(log_model=True)], early_stopping_rounds=run_3.config['early_stopping_rounds'])
+ # Initialize the XGBoostClassifier with the WandbCallback
+ # Train the model
+ xgbmodel.fit(trndat[p_vars], y_trn, eval_set=[(valdat[p_vars], y_val)])
+ return WandbCallback, xgbmodel
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### 5) Log Additional Train and Evaluation Metrics to W&B
+ """)
+ return
+
+
+@app.cell
+def _(
+ ks_2samp,
+ metrics,
+ np,
+ p_vars,
+ run_3,
+ trndat,
+ valdat,
+ xgbmodel,
+ y_trn,
+ y_val,
+):
+ bstr = xgbmodel.get_booster()
+ trnYpreds = xgbmodel.predict_proba(trndat[p_vars])[:, 1]
+ # Get train and validation predictions
+ valYpreds = xgbmodel.predict_proba(valdat[p_vars])[:, 1]
+ false_positive_rate, true_positive_rate, thresholds = metrics.roc_curve(y_trn, trnYpreds)
+ run_3.summary['train_ks_stat'] = max(true_positive_rate - false_positive_rate)
+ # Log additional Train metrics
+ run_3.summary['train_auc'] = metrics.auc(false_positive_rate, true_positive_rate)
+ run_3.summary['train_log_loss'] = -(y_trn * np.log(trnYpreds) + (1 - y_trn) * np.log(1 - trnYpreds)).sum() / len(y_trn)
+ ks_stat, ks_pval = ks_2samp(valYpreds[y_val == 1], valYpreds[y_val == 0])
+ run_3.summary['val_ks_2samp'] = ks_stat
+ run_3.summary['val_ks_pval'] = ks_pval
+ # Log additional Validation metrics
+ run_3.summary['val_auc'] = metrics.roc_auc_score(y_val, valYpreds)
+ run_3.summary['val_acc_0.5'] = metrics.accuracy_score(y_val, np.where(valYpreds >= 0.5, 1, 0))
+ run_3.summary['val_log_loss'] = -(y_val * np.log(valYpreds) + (1 - y_val) * np.log(1 - valYpreds)).sum() / len(y_val)
+ return (valYpreds,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### 6) Log the ROC Curve To W&B
+ """)
+ return
+
+
+@app.cell
+def _(np, run_3, valYpreds, wandb, y_val):
+ # Log the ROC curve to W&B
+ valYpreds_2d = np.array([1 - valYpreds, valYpreds]) # W&B expects a 2d array
+ y_val_arr = y_val.values
+ d = 0
+ while len(valYpreds_2d.T) > 10000:
+ d = d + 1
+ valYpreds_2d = valYpreds_2d[::1, ::d]
+ y_val_arr = y_val_arr[::d]
+ run_3.log({'ROC_Curve': wandb.plot.roc_curve(y_val_arr, valYpreds_2d.T, labels=['no_default', 'loan_default'], classes_to_plot=[1])})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Finish the W&B Run
+ """)
+ return
+
+
+@app.cell
+def _(run_3):
+ run_3.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now that we've trained a single model, lets try and optimize its performance by running a Hyperparameter Sweep.
+
+ # HyperParameter Sweep
+
+ Weights and Biases also enables you to do hyperparameter sweeps, either with our own [Sweeps functionality](https://docs.wandb.ai/guides/sweeps/python-api) or with our [Ray Tune integration](https://docs.wandb.ai/guides/sweeps/advanced-sweeps/ray-tune). See [our docs](https://docs.wandb.ai/guides/sweeps/python-api) for a full guide of how to use more advanced hyperparameter sweeps options.
+
+ **[Click Here](https://wandb.ai/morgan/credit_score_sweeps/sweeps/iuppbs45)** to check out the results of a 1000 run sweep generated using this notebook
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Define the Sweep Config
+ First we define the hyperparameters to sweep over as well as the type of sweep to use, we'll do a random search over the learning_rate, gamma, min_child_weights and easrly_stopping_rounds
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, wandb):
+ sweep_config = {
+ "method" : "random",
+ "parameters" : {
+ "learning_rate" :{
+ "min": 0.001,
+ "max": 1.0
+ },
+ "gamma" :{
+ "min": 0.001,
+ "max": 1.0
+ },
+ "min_child_weight" :{
+ "min": 1,
+ "max": 150
+ },
+ "early_stopping_rounds" :{
+ "values" : [10, 20, 30, 40]
+ },
+ }
+ }
+
+ sweep_id = wandb.sweep(sweep_config, project=WANDB_PROJECT)
+ return (sweep_id,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Define the Training Function
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Then we define the function that will train our model using these hyperparameters. Note that `job_type='sweep'` when initialising the run, so that we can easily filter out these runs from our main workspace if we need to
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbCallback,
+ base_rate,
+ ks_2samp,
+ metrics,
+ np,
+ p_vars,
+ trndat,
+ valdat,
+ wandb,
+ xgb,
+ y_trn,
+ y_val,
+):
+ def train():
+ with wandb.init(job_type="sweep") as run:
+
+ bst_params = {
+ 'objective': 'binary:logistic'
+ , 'base_score': base_rate
+ , 'gamma': run.config['gamma']
+ , 'learning_rate': run.config['learning_rate']
+ , 'max_depth': 3
+ , 'min_child_weight': run.config['min_child_weight']
+ , 'n_estimators': 25
+ , 'nthread': 24
+ , 'random_state': 42
+ , 'reg_alpha': 0
+ , 'reg_lambda': 0 ## def: 1
+ , 'eval_metric': ['auc', 'logloss']
+ , 'tree_method': 'hist'
+ }
+
+ # Initialize the XGBoostClassifier with the WandbCallback
+ xgbmodel = xgb.XGBClassifier(**bst_params,
+ callbacks=[WandbCallback()],
+ early_stopping_rounds=run.config['early_stopping_rounds'])
+
+ # Train the model
+ xgbmodel.fit(trndat[p_vars], y_trn,
+ eval_set=[(valdat[p_vars], y_val)])
+
+ bstr = xgbmodel.get_booster()
+
+ # Log booster metrics
+ run.summary["best_iteration"] = bstr.best_iteration
+
+ # Get train and validation predictions
+ trnYpreds = xgbmodel.predict_proba(trndat[p_vars])[:,1]
+ valYpreds = xgbmodel.predict_proba(valdat[p_vars])[:,1]
+
+ # Log additional Train metrics
+ false_positive_rate, true_positive_rate, thresholds = metrics.roc_curve(y_trn, trnYpreds)
+ run.summary['train_ks_stat'] = max(true_positive_rate - false_positive_rate)
+ run.summary['train_auc'] = metrics.auc(false_positive_rate, true_positive_rate)
+ run.summary['train_log_loss'] = -(y_trn * np.log(trnYpreds) + (1-y_trn) * np.log(1-trnYpreds)).sum() / len(y_trn)
+
+ # Log additional Validation metrics
+ ks_stat, ks_pval = ks_2samp(valYpreds[y_val==1], valYpreds[y_val==0])
+ run.summary["val_ks_2samp"] = ks_stat
+ run.summary["val_ks_pval"] = ks_pval
+ run.summary["val_auc"] = metrics.roc_auc_score(y_val, valYpreds)
+ run.summary["val_acc_0.5"] = metrics.accuracy_score(y_val, np.where(valYpreds >= 0.5, 1, 0))
+ run.summary["val_log_loss"] = -(y_val * np.log(valYpreds)
+ + (1-y_val) * np.log(1-valYpreds)).sum() / len(y_val)
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Run the Sweeps Agent
+ """)
+ return
+
+
+@app.cell
+def _(sweep_id, train, wandb):
+ count = 5 # number of runs to execute
+ wandb.agent(sweep_id, function=train, count=count)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## W&B already in your favorite ML library
+
+ Weights and Biases has integrations in all of your favourite ML and Deep Learning libraries such as:
+
+ - Pytorch Lightning
+ - Keras
+ - Hugging Face
+ - JAX
+ - Fastai
+ - XGBoost
+ - Sci-Kit Learn
+ - LightGBM
+
+ **See [W&B integrations for details](https://docs.wandb.ai/guides/integrations)**
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/boosting-simple-lightgbm-integration/boosting_simple_lightgbm_integration.py b/marimo/convert/boosting-simple-lightgbm-integration/boosting_simple_lightgbm_integration.py
new file mode 100644
index 00000000..36ea5e94
--- /dev/null
+++ b/marimo/convert/boosting-simple-lightgbm-integration/boosting_simple_lightgbm_integration.py
@@ -0,0 +1,362 @@
+# /// script
+# dependencies = ["lightgbm", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B + LightGBM
+ Gradient boosting decision trees are the state of the art when it comes to building predictive models for structured data.
+
+ [LigthGBM](https://github.com/microsoft/LightGBM), a gradient boosting framework by Microsoft, has dethroned xgboost and become the go to GBDT algorithm (along with catboost). It outperforms xgboost in training speeds, memory usage and the size of datasets it can handle. LightGBM does so by using histogram-based algorithms to bucket continuous features into discrete bins during training.
+
+ You can find the **[W&B + LightGBM documentation here](https://docs.wandb.ai/guides/integrations/boosting)**
+
+ ## What this notebook covers
+ * Easy integration of Weights and Biases with LightGBM.
+ * `wandb_callback()` callback for metrics logging
+ * `log_summary()` function to log a feature importance plot and enable model saving to W&B
+
+ We want to make it incredible easy for people to look under the hood of their models, so we built a callback that helps you visualize your LightGBM’s performance in just one line of code.
+
+ **Note**: Sections starting with _Step_ is all you need to integrate W&B.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install, Import, and Log in
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## The Usual Suspects
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: lightgbm>=4.0.0 !pip install -Uq 'lightgbm>=4.0.0'
+ return
+
+
+@app.cell
+def _():
+ import pandas as pd
+ import lightgbm as lgb
+ from sklearn.metrics import mean_squared_error
+
+ return lgb, mean_squared_error, pd
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Step 0: Install W&B
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Step 1: Import W&B and Login
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.integration.lightgbm import wandb_callback, log_summary
+
+ return log_summary, wandb, wandb_callback
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download and Prepare Dataset
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! wget https://raw.githubusercontent.com/microsoft/LightGBM/master/examples/regression/regression.train -qq
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/microsoft/LightGBM/master/examples/regression/regression.train', '-qq'])
+ #! wget https://raw.githubusercontent.com/microsoft/LightGBM/master/examples/regression/regression.test -qq
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/microsoft/LightGBM/master/examples/regression/regression.test', '-qq'])
+ return
+
+
+@app.cell
+def _(lgb, pd):
+ # load or create your dataset
+ df_train = pd.read_csv('regression.train', header=None, sep='\t')
+ df_test = pd.read_csv('regression.test', header=None, sep='\t')
+
+ y_train = df_train[0]
+ y_test = df_test[0]
+ X_train = df_train.drop(0, axis=1)
+ X_test = df_test.drop(0, axis=1)
+
+ # create dataset for lightgbm
+ lgb_train = lgb.Dataset(X_train, y_train)
+ lgb_eval = lgb.Dataset(X_test, y_test, reference=lgb_train)
+ return X_test, lgb_eval, lgb_train, y_test
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 2: Initialize your wandb run.
+
+ Using `wandb.init()` initialize your W&B run. You can also pass a dictionary of configs. [Check out the official documentation here $\rightarrow$](https://docs.wandb.com/library/init)
+
+ You can't deny the importance of configs in your ML/DL workflow. W&B makes sure that you have access to the right config to reproduce your model.
+
+ [Learn more about configs in this colab notebook $\rightarrow$](http://wandb.me/config-colab)
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # specify your configurations as a dict
+ params = {
+ 'boosting_type': 'gbdt',
+ 'objective': 'regression',
+ 'metric': ['rmse', 'l2', 'l1', 'huber'],
+ 'num_leaves': 31,
+ 'learning_rate': 0.05,
+ 'feature_fraction': 0.9,
+ 'bagging_fraction': 0.8,
+ 'bagging_freq': 5,
+ 'verbosity': 0,
+ 'early_stopping_rounds': 5,
+ }
+
+ wandb.init(project='my-lightgbm-project', config=params)
+ return (params,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > Once you have trained your model come back and click on the **Project page**.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 3: Train with `wandb_callback`
+ """)
+ return
+
+
+@app.cell
+def _(lgb, lgb_eval, lgb_train, params, wandb_callback):
+ # train
+ # add lightgbm callback
+ gbm = lgb.train(params,
+ lgb_train,
+ num_boost_round=30,
+ valid_sets=lgb_eval,
+ valid_names=('validation'),
+ callbacks=[wandb_callback()],
+ )
+ return (gbm,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 4: Log Feature Importance and Upload Model with `log_summary`
+ `log_summary` will upload calculate and upload the feature importance import and (optionally) upload your trained model to W&B Artifacts so you can use it later
+ """)
+ return
+
+
+@app.cell
+def _(gbm, log_summary):
+ log_summary(gbm, save_model_checkpoint=True)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Evaluate
+ """)
+ return
+
+
+@app.cell
+def _(X_test, gbm, mean_squared_error, wandb, y_test):
+ # predict
+ y_pred = gbm.predict(X_test, num_iteration=gbm.best_iteration)
+
+ # eval
+ print('The rmse of prediction is:', mean_squared_error(y_test, y_pred) ** 0.5)
+ wandb.log({'rmse_prediction': mean_squared_error(y_test, y_pred) ** 0.5})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ When you are finished logging for a particular W&B run its a good idea to call `wandb.finish()` to tidy up the wandb process (only necessary when using notebooks/colabs)
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Visualize Results
+
+ Click on the **project page** link above to see your results automatically visualized.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Sweep 101
+
+ Use Weights & Biases Sweeps to automate hyperparameter optimization and explore the space of possible models.
+
+ ## [Check out Hyperparameter Optimization with XGBoost using W&B Sweep $\rightarrow$](http://wandb.me/xgb-colab)
+
+ 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.com/library/sweeps/configuration) that specifies the parameters to search through, the search strategy, the optimization metric et all.
+
+ 2. **Initialize the sweep:**
+ `sweep_id = wandb.sweep(sweep_config)`
+
+ 3. **Run the sweep agent:**
+ `wandb.agent(sweep_id, function=train)`
+
+ And voila! That's all there is to running a hyperparameter sweep!
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Example Gallery
+
+ See examples of projects tracked and visualized with W&B in our [Gallery →](https://app.wandb.ai/gallery)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Basic Setup
+ 1. **Projects**: Log multiple runs to a project to compare them. `wandb.init(project="project-name")`
+ 2. **Groups**: For multiple processes or cross validation folds, log each process as a runs and group them together. `wandb.init(group='experiment-1')`
+ 3. **Tags**: Add tags to track your current baseline or production model.
+ 4. **Notes**: Type notes in the table to track the changes between runs.
+ 5. **Reports**: Take quick notes on progress to share with colleagues and make dashboards and snapshots of your ML projects.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Advanced Setup
+ 1. [Environment variables](https://docs.wandb.com/library/environment-variables): Set API keys in environment variables so you can run training on a managed cluster.
+ 2. [Offline mode](https://docs.wandb.com/library/technical-faq#can-i-run-wandb-offline): Use `dryrun` mode to train offline and sync results later.
+ 3. [On-prem](https://docs.wandb.com/self-hosted): 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.com/sweeps): Set up hyperparameter search quickly with our lightweight tool for tuning.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/boosting-using-w-b-sweeps-with-xgboost/boosting_using_w_b_sweeps_with_xgboost.py b/marimo/convert/boosting-using-w-b-sweeps-with-xgboost/boosting_using_w_b_sweeps_with_xgboost.py
new file mode 100644
index 00000000..e5dc2414
--- /dev/null
+++ b/marimo/convert/boosting-using-w-b-sweeps-with-xgboost/boosting_using_w_b_sweeps_with_xgboost.py
@@ -0,0 +1,385 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using W&B Sweeps with XGBoost
+
+ 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?
+
+ 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.
+
+ In this tutorial we'll see how you can run sophisticated hyperparameter sweeps on XGBoost models in 3 easy steps using Weights and Biases.
+
+ For a teaser, check out the plots below:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 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-like object that specifies the sweep: which parameters to search through, which search strategy to use, which metric to optimize.
+
+ 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 w`andb.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)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 1. 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).
+
+ 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.
+
+ You can find a list of all configuration options [here](https://docs.wandb.com/library/sweeps/configuration).
+ """)
+ return
+
+
+@app.cell
+def _():
+ sweep_config = {
+ "method": "random", # try grid or random
+ "metric": {
+ "name": "accuracy",
+ "goal": "maximize"
+ },
+ "parameters": {
+ "booster": {
+ "values": ["gbtree","gblinear"]
+ },
+ "max_depth": {
+ "values": [3, 6, 9, 12]
+ },
+ "learning_rate": {
+ "values": [0.1, 0.05, 0.2]
+ },
+ "subsample": {
+ "values": [1, 0.5, 0.3]
+ }
+ }
+ }
+ return (sweep_config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 2. Initialize the Sweep
+ """)
+ return
+
+
+@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.
+ """)
+ return
+
+
+@app.cell
+def _(sweep_config, wandb):
+ sweep_id = wandb.sweep(sweep_config, project="XGBoost-sweeps")
+ return (sweep_id,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We also need to download the data:
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! wget https://raw.githubusercontent.com/jbrownlee/Datasets/master/pima-indians-diabetes.data.csv
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/jbrownlee/Datasets/master/pima-indians-diabetes.data.csv'])
+ return
+
+
+@app.cell
+def _(wandb):
+ # XGBoost model for Pima Indians dataset
+ from numpy import loadtxt
+ from xgboost import XGBClassifier
+ from sklearn.model_selection import train_test_split
+ from sklearn.metrics import accuracy_score
+
+ # load data
+ def train():
+ config_defaults = {
+ "booster": "gbtree",
+ "max_depth": 3,
+ "learning_rate": 0.1,
+ "subsample": 1,
+ "seed": 117,
+ "test_size": 0.33,
+ }
+
+ wandb.init(config=config_defaults) # defaults are over-ridden during the sweep
+ config = wandb.config
+
+ # load data and split into predictors and targets
+ dataset = loadtxt("pima-indians-diabetes.data.csv", 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 model on train
+ 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 test
+ y_pred = model.predict(X_test)
+ predictions = [round(value) for value in y_pred]
+
+ # evaluate predictions
+ accuracy = accuracy_score(y_test, predictions)
+ print(f"Accuracy: {accuracy:.0%}")
+ wandb.log({"accuracy": accuracy})
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 3. Run the Sweep with an agent
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we call `wandb.agent` to start up our sweep.
+
+ 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
+
+ and that machine will join the sweep!
+
+ > _Note_: a `random` sweep will by defauly 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.
+ """)
+ return
+
+
+@app.cell
+def _(sweep_id, train, wandb):
+ wandb.agent(sweep_id, train, count=25)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+
+ Weights & Biases will generate a number of useful plots for you automatically.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Parallel coordinates plot
+
+ 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.
+
+ 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):
+ 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
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/catalyst-catalyst-x-wandb/catalyst_catalyst_x_wandb.py b/marimo/convert/catalyst-catalyst-x-wandb/catalyst_catalyst_x_wandb.py
new file mode 100644
index 00000000..135bb466
--- /dev/null
+++ b/marimo/convert/catalyst-catalyst-x-wandb/catalyst_catalyst_x_wandb.py
@@ -0,0 +1,231 @@
+# /// script
+# dependencies = ["catalyst", "opencv-python-headless", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Rapid Experimentation with Catalyst and W&B
+ 
+ Catalyst is a PyTorch framework for deep learning and R&D. It focuses on reproducibility, rapid experimentation, and codebase reuse so you can create something new rather than write yet another train loop.
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [](https://colab.research.google.com/drive/1woAYD9hot7mbknGhbxtix7x7u1fvIZJx?usp=sharing)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: catalyst==21.10 !pip install -q catalyst==21.10
+ # packages added via marimo's package management: catalyst[ml]==21.10 !pip install -q catalyst[ml]==21.10
+ # packages added via marimo's package management: catalyst[cv]==21.10 !pip install -q catalyst[cv]==21.10
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ # packages added via marimo's package management: opencv-python-headless==4.1.2.30 !pip install -q opencv-python-headless==4.1.2.30
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Restart run time after the above setup steps!
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ from tempfile import TemporaryDirectory
+
+ from torch import nn, optim
+ from torch.utils.data import DataLoader
+
+ from catalyst import dl, utils
+ from catalyst.contrib.datasets import MNIST
+ from catalyst.data.transforms import ToTensor
+ from catalyst.settings import SETTINGS
+
+ return (
+ DataLoader,
+ MNIST,
+ SETTINGS,
+ TemporaryDirectory,
+ ToTensor,
+ dl,
+ nn,
+ optim,
+ os,
+ utils,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Example - MNIST Classification
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, MNIST, SETTINGS, ToTensor, dl, nn, optim, os, utils):
+ class CustomRunner(dl.IRunner):
+ def __init__(self, logdir, device, engine):
+ super().__init__()
+ self._logdir = logdir
+ self._device = device
+ self._engine = engine
+ self._name = "mnist"
+
+ def get_engine(self):
+ return self._engine or dl.DeviceEngine(self._device)
+
+ def get_loggers(self):
+ loggers = {}
+ if SETTINGS.wandb_required:
+ loggers["wandb"] = dl.WandbLogger(project="catalyst_wandb", name=self._name)
+ return loggers
+
+ @property
+ def stages(self):
+ return ["train_freezed", "train_unfreezed"]
+
+ def get_stage_len(self, stage: str) -> int:
+ return 1
+
+ def get_loaders(self, stage: str):
+ loaders = {
+ "train": DataLoader(
+ MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()),
+ batch_size=32,
+ ),
+ "valid": DataLoader(
+ MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()),
+ batch_size=32,
+ ),
+ }
+ return loaders
+
+ def get_model(self, stage: str):
+ model = (
+ utils.get_nn_from_ddp_module(self.model)
+ if self.model is not None
+ else nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
+ )
+ if stage == "train_freezed":
+ # freeze layer
+ utils.set_requires_grad(model[1], False)
+ else:
+ utils.set_requires_grad(model, True)
+ return model
+
+ def get_criterion(self, stage: str):
+ return nn.CrossEntropyLoss()
+
+ def get_optimizer(self, stage: str, model):
+ if stage == "train_freezed":
+ return optim.Adam(model.parameters(), lr=1e-3)
+ else:
+ return optim.SGD(model.parameters(), lr=1e-1)
+
+ def get_scheduler(self, stage: str, optimizer):
+ return None
+
+ def get_callbacks(self, stage: str):
+ callbacks = {
+ "criterion": dl.CriterionCallback(
+ metric_key="loss", input_key="logits", target_key="targets"
+ ),
+ "optimizer": dl.OptimizerCallback(
+ metric_key="loss",
+ grad_clip_fn=nn.utils.clip_grad_norm_,
+ grad_clip_params={"max_norm": 1.0},
+ ),
+ # "scheduler": dl.SchedulerCallback(loader_key="valid", metric_key="loss"),
+ "accuracy": dl.AccuracyCallback(
+ input_key="logits", target_key="targets", topk_args=(1, 3, 5)
+ ),
+ "classification": dl.PrecisionRecallF1SupportCallback(
+ input_key="logits", target_key="targets", num_classes=10
+ ),
+ "checkpoint": dl.CheckpointCallback(
+ self._logdir, loader_key="valid", metric_key="loss", minimize=True, save_n_best=3
+ ),
+ }
+ if SETTINGS.ml_required:
+ callbacks["confusion_matrix"] = dl.ConfusionMatrixCallback(
+ input_key="logits", target_key="targets", num_classes=10
+ )
+ return callbacks
+
+ def handle_batch(self, batch):
+ x, y = batch
+ logits = self.model(x)
+
+ self.batch = {
+ "features": x,
+ "targets": y,
+ "logits": logits,
+ }
+
+ return (CustomRunner,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## W&B integration
+ Catalyst comes with `WandbLogger` based on its logger API which can be used to stream your experiment metrics, media and experiment configurations.
+ 
+ """)
+ return
+
+
+@app.cell
+def _(CustomRunner, TemporaryDirectory):
+ def train_experiment(device, engine=None):
+ with TemporaryDirectory() as logdir:
+ runner = CustomRunner(logdir, device, engine)
+ runner.run()
+
+ return (train_experiment,)
+
+
+@app.cell
+def _(train_experiment):
+ train_experiment("cuda:0") # Train on GPU
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/convnext-finetune-convnext-on-cifar10-using-w-b/convnext_finetune_convnext_on_cifar10_using_w_b.py b/marimo/convert/convnext-finetune-convnext-on-cifar10-using-w-b/convnext_finetune_convnext_on_cifar10_using_w_b.py
new file mode 100644
index 00000000..5b42c2b7
--- /dev/null
+++ b/marimo/convert/convnext-finetune-convnext-on-cifar10-using-w-b/convnext_finetune_convnext_on_cifar10_using_w_b.py
@@ -0,0 +1,166 @@
+# /// script
+# dependencies = ["https://download-pytorch-org/whl/torch-stable-html", "six", "tensorboardx", "timm", "torch", "torchvision", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use this notebook to finetune a ConvNeXt-tiny model on CIFAR 10 dataset. The [official ConvNeXt repository](https://github.com/facebookresearch/ConvNeXt) is instrumented with [Weights and Biases](https://wandb.ai/site). You can now easily log your train/test metrics and version control your model checkpoints to Weigths and Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ⚽️ Installation and Setup
+
+ The following installation instruction is based on [INSTALL.md](https://github.com/facebookresearch/ConvNeXt/blob/main/INSTALL.md) provided by the official ConvNeXt repository.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: torch==1.8.0+cu111 torchvision==0.9.0+cu111 https://download.pytorch.org/whl/torch_stable.html !pip install -qq torch==1.8.0+cu111 torchvision==0.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
+ # packages added via marimo's package management: wandb timm==0.3.2 six tensorboardX !pip install -qq wandb timm==0.3.2 six tensorboardX
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Download the official ConvNeXt respository.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! git clone --depth 1 https://github.com/facebookresearch/ConvNeXt
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/facebookresearch/ConvNeXt'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏀 Download the Dataset
+
+ We will be finetuning on CIFAR-10 dataset. To use any custom dataset (CIFAR-10 here) the format of the dataset should be as shown below:
+
+ ```
+ /path/to/dataset/
+ train/
+ class1/
+ img1.jpeg
+ class2/
+ img2.jpeg
+ val/
+ class1/
+ img3.jpeg
+ class2/
+ img4.jpeg
+ ```
+
+ I have used this [repository](https://github.com/YoongiKim/CIFAR-10-images) that has the CIFAR-10 images in the required format.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! git clone --depth 1 https://github.com/YoongiKim/CIFAR-10-images
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/YoongiKim/CIFAR-10-images'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏈 Download Pretrained Weights
+
+ We will be finetuning the ConvNeXt Tiny model pretrained on ImageNet 1K dataset.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ import os
+ os.chdir('ConvNeXt/')
+ #! wget https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth
+ subprocess.call(['wget', 'https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🎾 Train with Weights and Biases
+
+ If you want to log the train and evaluation metrics using Weights and Biases pass `--enable_wandb true`.
+
+ You can also save the finetuned checkpoints as version controlled W&B [Artifacts](https://docs.wandb.ai/guides/artifacts) if you pass `--wandb_ckpt true`.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python main.py --epochs 10 --model convnext_tiny --data_set image_folder --data_path ../CIFAR-10-images/train --eval_data_path ../CIFAR-10-images/test --nb_classes 10 --num_workers 8 --warmup_epochs 0 --save_ckpt true --output_dir model_ckpt --finetune convnext_tiny_1k_224_ema.pth --cutmix 0 --mixup 0 --lr 4e-4 --enable_wandb true --wandb_ckpt true
+ subprocess.call(['python', 'main.py', '--epochs', '10', '--model', 'convnext_tiny', '--data_set', 'image_folder', '--data_path', '../CIFAR-10-images/train', '--eval_data_path', '../CIFAR-10-images/test', '--nb_classes', '10', '--num_workers', '8', '--warmup_epochs', '0', '--save_ckpt', 'true', '--output_dir', 'model_ckpt', '--finetune', 'convnext_tiny_1k_224_ema.pth', '--cutmix', '0', '--mixup', '0', '--lr', '4e-4', '--enable_wandb', 'true', '--wandb_ckpt', 'true'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏐 Conclusion
+
+ * **The above setting gives a top-1 accuracy of ~95%.**
+ * The ConvNeXt repository comes with modern training regimes and is easy to finetune on any dataset.
+ * The finetune model achieves competitive results.
+
+ * By passing two arguments you get the following:
+
+ * Repository of all your experiments (train and test metrics) as a [W&B Project](https://docs.wandb.ai/ref/app/pages/project-page). You can easily compare experiments to find the best performing model.
+ * Hyperparameters (Configs) used to train individual models.
+ * System (CPU/GPU/Disk) metrics.
+ * Model checkpoints saved as W&B Artifacts. They are versioned and easy to share.
+
+ Check out the associated [W&B run page](https://wandb.ai/ayut/convnext/runs/16vi9e31). $→$
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/cross-attention-control-cross-attention-control-wandb/cross_attention_control_cross_attention_control_wandb.py b/marimo/convert/cross-attention-control-cross-attention-control-wandb/cross_attention_control_cross_attention_control_wandb.py
new file mode 100644
index 00000000..9e1429c9
--- /dev/null
+++ b/marimo/convert/cross-attention-control-cross-attention-control-wandb/cross_attention_control_cross_attention_control_wandb.py
@@ -0,0 +1,746 @@
+# /// script
+# dependencies = ["diffusers", "ftfy", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Cross-Attention Control with Stable Diffusion + WandB Playground 🪄🐝
+
+
+
+ An implementation of Prompt-to-Prompt Image Editing
+ with Cross Attention Control using [Stable Diffusion](https://github.com/CompVis/stable-diffusion), [HuggingFace Diffusers](https://github.com/huggingface/diffusers) and [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 1: Setup required libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ #@title
+
+ # packages added via marimo's package management: diffusers transformers ftfy wandb !pip install -q diffusers transformers ftfy wandb
+ return
+
+
+@app.cell
+def _():
+ #@title
+
+ import io
+ import wandb
+ import random
+ import numpy as np
+ from PIL import Image
+ from tqdm.auto import tqdm
+ from difflib import SequenceMatcher
+ from google.colab import files as colab_files
+
+ import torch
+ from torch import autocast
+
+ from transformers import CLIPModel, CLIPTextModel, CLIPTokenizer
+ from diffusers import (
+ AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler
+ )
+
+ return (
+ AutoencoderKL,
+ CLIPModel,
+ CLIPTokenizer,
+ Image,
+ LMSDiscreteScheduler,
+ SequenceMatcher,
+ UNet2DConditionModel,
+ autocast,
+ np,
+ random,
+ torch,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 2: Set up Models and Weights & Biases Run
+
+ - `wandb_project`: Weights & Biases project.
+ - `wandb_project`: Weights & Biases entity.
+ - `huggingface_access_token`: HuggingFace Access Token. Check out this page from the official HuggingFace docs as to how to generate your own access token.
+ - `config.device`: Accelerator device. Choose `mps` if you're running the code on an M1 Mac.
+ - `config.model_path_clip`: Alias for pre-trained CLIP Model.
+ - `config.model_path_diffusion`: Alias for pre-trained Stable Diffusion Model.
+ """)
+ return
+
+
+@app.cell
+def _(
+ AutoencoderKL,
+ CLIPModel,
+ CLIPTokenizer,
+ UNet2DConditionModel,
+ torch,
+ wandb,
+):
+ wandb_project = "cross-attention-control" #@param {"type": "string"}
+ wandb_entity = "wandb" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, entity=wandb_entity, job_type="generate")
+ config = wandb.config
+
+ huggingface_access_token = "" #@param {"type": "string"}
+ torch_dtype = torch.float16
+
+ config.model_precision_type = "fp16"
+ config.device = "cuda" #@param['cuda', 'cpu', 'mps']
+ config.model_path_clip = "openai/clip-vit-large-patch14" #@param['openai/clip-vit-large-patch14']
+ config.model_path_diffusion = "CompVis/stable-diffusion-v1-4" #@param['CompVis/stable-diffusion-v1-4']
+
+
+ clip_tokenizer = CLIPTokenizer.from_pretrained(config.model_path_clip)
+ clip_model = CLIPModel.from_pretrained(
+ config.model_path_clip,
+ torch_dtype=torch_dtype
+ )
+ clip = clip_model.text_model
+
+
+ model_path_diffusion = "CompVis/stable-diffusion-v1-4"
+ unet = UNet2DConditionModel.from_pretrained(
+ model_path_diffusion,
+ subfolder="unet",
+ use_auth_token=huggingface_access_token,
+ revision=config.model_precision_type,
+ torch_dtype=torch.float16
+ )
+ vae = AutoencoderKL.from_pretrained(
+ model_path_diffusion,
+ subfolder="vae",
+ use_auth_token=huggingface_access_token,
+ revision=config.model_precision_type,
+ torch_dtype=torch.float16
+ )
+
+
+ unet.to(config.device)
+ vae.to(config.device)
+ clip.to(config.device)
+ return clip, clip_tokenizer, config, unet, vae
+
+
+@app.cell
+def _(Image, SequenceMatcher, clip_tokenizer, config, torch, unet):
+ #@title
+
+ def init_attention_weights(weight_tuples):
+ tokens_length = clip_tokenizer.model_max_length
+ weights = torch.ones(tokens_length)
+
+ for i, w in weight_tuples:
+ if i < tokens_length and i >= 0:
+ weights[i] = w
+
+
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention" and "attn2" in name:
+ module.last_attn_slice_weights = weights.to(config.device)
+ if module_name == "CrossAttention" and "attn1" in name:
+ module.last_attn_slice_weights = None
+
+
+ def init_attention_edit(tokens, tokens_edit):
+ tokens_length = clip_tokenizer.model_max_length
+ mask = torch.zeros(tokens_length)
+ indices_target = torch.arange(tokens_length, dtype=torch.long)
+ indices = torch.zeros(tokens_length, dtype=torch.long)
+
+ tokens = tokens.input_ids.numpy()[0]
+ tokens_edit = tokens_edit.input_ids.numpy()[0]
+
+ for name, a0, a1, b0, b1 in SequenceMatcher(
+ None, tokens, tokens_edit
+ ).get_opcodes():
+ if b0 < tokens_length:
+ if name == "equal" or (name == "replace" and a1-a0 == b1-b0):
+ mask[b0:b1] = 1
+ indices[b0:b1] = indices_target[a0:a1]
+
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention" and "attn2" in name:
+ module.last_attn_slice_mask = mask.to(config.device)
+ module.last_attn_slice_indices = indices.to(config.device)
+ if module_name == "CrossAttention" and "attn1" in name:
+ module.last_attn_slice_mask = None
+ module.last_attn_slice_indices = None
+
+
+ def init_attention_func():
+ def new_attention(self, query, key, value, sequence_length, dim):
+ batch_size_attention = query.shape[0]
+ hidden_states = torch.zeros(
+ (batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype
+ )
+ slice_size = self._slice_size if self._slice_size is not None else hidden_states.shape[0]
+ for i in range(hidden_states.shape[0] // slice_size):
+ start_idx = i * slice_size
+ end_idx = (i + 1) * slice_size
+ attn_slice = (
+ torch.einsum("b i d, b j d -> b i j", query[start_idx:end_idx], key[start_idx:end_idx]) * self.scale
+ )
+ attn_slice = attn_slice.softmax(dim=-1)
+
+ if self.use_last_attn_slice:
+ if self.last_attn_slice_mask is not None:
+ new_attn_slice = torch.index_select(self.last_attn_slice, -1, self.last_attn_slice_indices)
+ attn_slice = attn_slice * (1 - self.last_attn_slice_mask) + new_attn_slice * self.last_attn_slice_mask
+ else:
+ attn_slice = self.last_attn_slice
+
+ self.use_last_attn_slice = False
+
+ if self.save_last_attn_slice:
+ self.last_attn_slice = attn_slice
+ self.save_last_attn_slice = False
+
+ if self.use_last_attn_weights and self.last_attn_slice_weights is not None:
+ attn_slice = attn_slice * self.last_attn_slice_weights
+ self.use_last_attn_weights = False
+
+ attn_slice = torch.einsum("b i j, b j d -> b i d", attn_slice, value[start_idx:end_idx])
+
+ hidden_states[start_idx:end_idx] = attn_slice
+
+ # reshape hidden_states
+ hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
+ return hidden_states
+
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention":
+ module.last_attn_slice = None
+ module.use_last_attn_slice = False
+ module.use_last_attn_weights = False
+ module.save_last_attn_slice = False
+ module._attention = new_attention.__get__(module, type(module))
+
+ def use_last_tokens_attention(use=True):
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention" and "attn2" in name:
+ module.use_last_attn_slice = use
+
+ def use_last_tokens_attention_weights(use=True):
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention" and "attn2" in name:
+ module.use_last_attn_weights = use
+
+ def use_last_self_attention(use=True):
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention" and "attn1" in name:
+ module.use_last_attn_slice = use
+
+ def save_last_tokens_attention(save=True):
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention" and "attn2" in name:
+ module.save_last_attn_slice = save
+
+ def save_last_self_attention(save=True):
+ for name, module in unet.named_modules():
+ module_name = type(module).__name__
+ if module_name == "CrossAttention" and "attn1" in name:
+ module.save_last_attn_slice = save
+
+
+ def postprocess(image):
+ image = (image / 2 + 0.5).clamp(0, 1)
+ image = image.cpu().permute(0, 2, 3, 1).numpy()
+ image = (image[0] * 255).round().astype("uint8")
+ return Image.fromarray(image)
+
+ return (
+ init_attention_edit,
+ init_attention_func,
+ init_attention_weights,
+ postprocess,
+ save_last_self_attention,
+ save_last_tokens_attention,
+ use_last_self_attention,
+ use_last_tokens_attention,
+ use_last_tokens_attention_weights,
+ )
+
+
+@app.cell
+def _(
+ Image,
+ LMSDiscreteScheduler,
+ autocast,
+ clip,
+ clip_tokenizer,
+ config,
+ init_attention_edit,
+ init_attention_func,
+ init_attention_weights,
+ np,
+ postprocess,
+ random,
+ save_last_self_attention,
+ save_last_tokens_attention,
+ torch,
+ tqdm,
+ unet,
+ use_last_self_attention,
+ use_last_tokens_attention,
+ use_last_tokens_attention_weights,
+ vae,
+ wandb,
+):
+ #@title
+
+ @torch.no_grad()
+ def stablediffusion(
+ prompt="",
+ prompt_edit="",
+ prompt_edit_token_weights=[],
+ prompt_edit_tokens_start=0.0,
+ prompt_edit_tokens_end=1.0,
+ prompt_edit_spatial_start=0.0,
+ prompt_edit_spatial_end=1.0,
+ guidance_scale=7.5,
+ steps=50,
+ seed=None,
+ width=512,
+ height=512,
+ init_image=None,
+ init_image_strength=0.5,
+ ):
+ log_key = (
+ "Generated Image without Promp Edit"
+ if prompt_edit == ""
+ else "Generated Image with Promp Edit"
+ )
+ print(log_key)
+
+ # Change size to multiple of 64 to prevent size mismatches inside model
+ width = width - width % 64
+ height = height - height % 64
+
+ #If seed is None, randomly select seed from 0 to 2^32-1
+ if seed is None: seed = random.randrange(2**32 - 1)
+ generator = torch.cuda.manual_seed(seed)
+
+ # Set inference timesteps to scheduler
+ scheduler = LMSDiscreteScheduler(
+ beta_start=0.00085,
+ beta_end=0.012,
+ beta_schedule="scaled_linear",
+ num_train_timesteps=1000
+ )
+ scheduler.set_timesteps(steps)
+
+ # Preprocess image if it exists (img2img)
+ if init_image is not None:
+ #Resize and transpose for numpy b h w c -> torch b c h w
+ init_image = init_image.resize(
+ (width, height), resample=Image.LANCZOS
+ )
+ init_image = np.array(
+ init_image
+ ).astype(np.float32) / 255.0 * 2.0 - 1.0
+ init_image = torch.from_numpy(
+ init_image[np.newaxis, ...].transpose(0, 3, 1, 2)
+ )
+
+ # If there is alpha channel, composite alpha for white,
+ # as the diffusion model does not support alpha channel
+ if init_image.shape[1] > 3:
+ init_image = init_image[:, :3] * init_image[:, 3:] + (
+ 1 - init_image[:, 3:]
+ )
+
+ #Move image to GPU
+ init_image = init_image.to(config.device)
+
+ #Encode image
+ with autocast(config.device):
+ init_latent = vae.encode(
+ init_image
+ ).latent_dist.sample(generator=generator) * 0.18215
+
+ t_start = steps - int(steps * init_image_strength)
+
+ else:
+ init_latent = torch.zeros(
+ (1, unet.in_channels, height // 8, width // 8),
+ device=config.device
+ )
+ t_start = 0
+
+ # Generate random normal noise
+ noise = torch.randn(
+ init_latent.shape, generator=generator, device=config.device
+ )
+ latent = scheduler.add_noise(
+ init_latent, noise, t_start
+ ).to(config.device)
+
+ # Process clip
+ with autocast(config.device):
+ tokens_unconditional = clip_tokenizer(
+ "",
+ padding="max_length",
+ max_length=clip_tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ return_overflowing_tokens=True
+ )
+ embedding_unconditional = clip(
+ tokens_unconditional.input_ids.to(config.device)
+ ).last_hidden_state
+
+ tokens_conditional = clip_tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=clip_tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ return_overflowing_tokens=True
+ )
+ embedding_conditional = clip(
+ tokens_conditional.input_ids.to(config.device)
+ ).last_hidden_state
+
+ # Process prompt editing
+ if prompt_edit != "":
+ tokens_conditional_edit = clip_tokenizer(
+ prompt_edit,
+ padding="max_length",
+ max_length=clip_tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ return_overflowing_tokens=True
+ )
+ embedding_conditional_edit = clip(
+ tokens_conditional_edit.input_ids.to(config.device)
+ ).last_hidden_state
+
+ init_attention_edit(
+ tokens_conditional, tokens_conditional_edit
+ )
+
+ init_attention_func()
+ init_attention_weights(prompt_edit_token_weights)
+
+ timesteps = scheduler.timesteps[t_start:]
+
+ for i, t in tqdm(enumerate(timesteps), total=len(timesteps)):
+ t_index = t_start + i
+
+ sigma = scheduler.sigmas[t_index]
+ latent_model_input = latent
+ latent_model_input = (
+ latent_model_input / ((sigma**2 + 1) ** 0.5)
+ ).to(unet.dtype)
+
+ # Predict the unconditional noise residual
+ noise_pred_uncond = unet(
+ latent_model_input,
+ t,
+ encoder_hidden_states=embedding_unconditional
+ ).sample
+
+ # Prepare the Cross-Attention layers
+ if prompt_edit is not None:
+ save_last_tokens_attention()
+ save_last_self_attention()
+ else:
+ #Use weights on non-edited prompt when edit is None
+ use_last_tokens_attention_weights()
+
+ # Predict the conditional noise residual and save
+ # the cross-attention layer activations
+ noise_pred_cond = unet(
+ latent_model_input,
+ t,
+ encoder_hidden_states=embedding_conditional
+ ).sample
+
+ # Edit the Cross-Attention layer activations
+ if prompt_edit != "":
+ t_scale = t / scheduler.num_train_timesteps
+ if t_scale >= prompt_edit_tokens_start and t_scale <= prompt_edit_tokens_end:
+ use_last_tokens_attention()
+ if t_scale >= prompt_edit_spatial_start and t_scale <= prompt_edit_spatial_end:
+ use_last_self_attention()
+
+ # Use weights on edited prompt
+ use_last_tokens_attention_weights()
+
+ # Predict the edited conditional noise residual
+ # using the cross-attention masks
+ noise_pred_cond = unet(
+ latent_model_input,
+ t,
+ encoder_hidden_states=embedding_conditional_edit
+ ).sample
+
+ #Perform guidance
+ noise_pred = noise_pred_uncond + guidance_scale * (
+ noise_pred_cond - noise_pred_uncond
+ )
+
+ latent = scheduler.step(noise_pred, t_index, latent).prev_sample
+
+ wandb.log({
+ log_key: wandb.Image(
+ postprocess(
+ vae.decode((latent / 0.18215).to(vae.dtype)).sample
+ )
+ )
+ }, step=i)
+
+ # scale and decode the image latents with vae
+ latent = latent / 0.18215
+ image = vae.decode(latent.to(vae.dtype)).sample
+
+ return postprocess(image)
+
+ return (stablediffusion,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 3: Enter Prompts and Additional Configs
+
+ - `config.prompt`: The prompt as a string.
+ - `config.prompt_edit`: The second prompt as a string, used to edit the first prompt using cross attention, set `"\"` to disable.
+ - `config.prompt_edit_token_weights`: Values to scale the importance of the tokens in cross attention layers, as a list of tuples representing `(token id, strength)`, this is used to increase or decrease the importance of a word in the prompt, it is applied to prompt_edit when possible (if `prompt_edit` is `"\"`, weights are applied to prompt).
+ - `config.prompt_edit_tokens_start`: How strict is the generation with respect to the initial prompt, increasing this will let the network be more creative for smaller details/textures, should be smaller than `prompt_edit_tokens_end`.
+ - `config.prompt_edit_tokens_end`: How strict is the generation with respect to the initial prompt, decreasing this will let the network be more creative for larger features/general scene composition, should be bigger than `prompt_edit_tokens_start`.
+ - `config.prompt_edit_spatial_start`: How strict is the generation with respect to the initial image (generated from the first prompt, not from img2img), increasing this will let the network be more creative for smaller details/textures, should be smaller than `prompt_edit_spatial_end`.
+ - `config.prompt_edit_spatial_end`: How strict is the generation with respect to the initial image (generated from the first prompt, not from img2img), decreasing this will let the network be more creative for larger features/general scene composition, should be bigger than `prompt_edit_spatial_start`.
+ - `config.guidance_scale`: Standard classifier-free guidance strength for stable diffusion.
+ - `config.steps`: Number of diffusion steps as an integer, higher usually produces better images but is slower.
+ - `config.seed`: Random Seed.
+ - `config.image_width`: Width of generated image.
+ - `config.image_height`: Height of generated image.
+ """)
+ return
+
+
+@app.cell
+def _(clip_tokenizer, config):
+ def display_prompt_tokens(prompt):
+ tokens = clip_tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=clip_tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ return_overflowing_tokens=True
+ ).input_ids[0]
+ for idx, token in enumerate(tokens):
+ decoded_token = clip_tokenizer.decode(token)
+ if decoded_token == "<|startoftext|>":
+ continue
+ elif decoded_token == "<|endoftext|>":
+ break
+ else:
+ print(idx, "->", decoded_token)
+
+
+ # the prompt as a string
+ config.prompt = "A photo of a Person with flower headpiece and elegant jewels" #@param {"type": "string"}
+
+ # the second prompt as a string, used to edit the first prompt
+ # using cross attention, set "" to disable
+ config.prompt_edit = "A photo of a Person with butterfly headpiece and elegant jewels" #@param {"type": "string"}
+
+ display_prompt_tokens(config.prompt_edit)
+ return
+
+
+@app.cell
+def _(config):
+ # values to scale the importance of the tokens in
+ # cross attention layers, as a list of tuples representing
+ # (token id, strength), this is used to increase or decrease
+ # the importance of a word in the prompt, it is applied to prompt_edit when possible (if prompt_edit is None, weights are applied to prompt)
+ config.prompt_edit_token_weights = [(7, 4)] #@param {type:"raw"}
+
+ # how strict is the generation with respect to the initial prompt,
+ # increasing this will let the network be more creative for smaller
+ # details/textures, should be smaller than prompt_edit_tokens_end
+ config.prompt_edit_tokens_start = 0.0 #@param {type:"slider", min:0, max:1, step:0.1}
+
+ # how strict is the generation with respect to the initial prompt,
+ # decreasing this will let the network be more creative for larger
+ # features/general scene composition, should be bigger than
+ # prompt_edit_tokens_start
+ config.prompt_edit_tokens_end = 1.0 #@param {type:"slider", min:0, max:1, step:0.1}
+
+ # how strict is the generation with respect to the initial image
+ # (generated from the first prompt, not from img2img), increasing
+ # this will let the network be more creative for smaller
+ # details/textures, should be smaller than prompt_edit_spatial_end
+ config.prompt_edit_spatial_start = 0.0 #@param {type:"slider", min:0, max:1, step:0.1}
+
+ # how strict is the generation with respect to the initial image
+ # (generated from the first prompt, not from img2img), decreasing
+ # this will let the network be more creative for larger
+ # features/general scene composition, should be bigger than
+ # prompt_edit_spatial_start
+ config.prompt_edit_spatial_end = 0.8 #@param {type:"slider", min:0, max:1, step:0.1}
+
+ # standard classifier-free guidance strength for stable diffusion
+ config.guidance_scale = 7.5 #@param {type:"slider", min:0, max:100, step:0.1}
+
+ # number of diffusion steps as an integer, higher usually produces
+ # better images but is slower
+ config.steps = 50 #@param {type:"slider", min:0, max:1000, step:1}
+
+ # random seed as an integer
+ config.seed = 98374234 #@param {type:"number"}
+
+ # image width and heigh
+ config.image_width = 768 #@param {type:"slider", min:512, max:1024, step:1}
+ config.image_height = 512 #@param {type:"slider", min:512, max:1024, step:1}
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 4: Generate Images with Prompt and Prompt Edits.
+
+ Image genetaed will be automatically logged to the respective **Weights & Biases** workspace as an interactive [**Table**](https://docs.wandb.ai/guides/data-vis) with all configs.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _(config, stablediffusion, wandb):
+ #@title
+
+ generated_image_with_prompt_edit = stablediffusion(
+ prompt=config.prompt,
+ prompt_edit=config.prompt_edit,
+ prompt_edit_token_weights=config.prompt_edit_token_weights,
+ prompt_edit_tokens_start=config.prompt_edit_tokens_start,
+ prompt_edit_tokens_end=config.prompt_edit_tokens_end,
+ prompt_edit_spatial_start=config.prompt_edit_spatial_start,
+ prompt_edit_spatial_end=config.prompt_edit_spatial_end,
+ guidance_scale=config.guidance_scale,
+ steps=config.steps,
+ seed=config.seed,
+ width=config.image_width,
+ height=config.image_height,
+ init_image=None,
+ init_image_strength=0.5
+ )
+
+ if config.prompt_edit != "":
+ generated_image_without_prompt_edit = stablediffusion(
+ prompt=config.prompt,
+ prompt_edit="",
+ prompt_edit_token_weights=config.prompt_edit_token_weights,
+ prompt_edit_tokens_start=config.prompt_edit_tokens_start,
+ prompt_edit_tokens_end=config.prompt_edit_tokens_end,
+ prompt_edit_spatial_start=config.prompt_edit_spatial_start,
+ prompt_edit_spatial_end=config.prompt_edit_spatial_end,
+ guidance_scale=config.guidance_scale ,
+ steps=config.steps,
+ seed=config.seed,
+ width=config.image_width,
+ height=config.image_height,
+ init_image=None,
+ init_image_strength=0.5
+ )
+ table = wandb.Table(
+ columns=[
+ "Seed",
+ "Guidance Scale",
+ "Image Height",
+ "Image Width",
+ "Number of Steps",
+ "Prompt",
+ "Image Generated With Prompt",
+ "Prompt Edit",
+ "Edit Token Weights",
+ "Image Generated With Prompt Edit"
+ ]
+ )
+ table.add_data(
+ config.seed,
+ config.guidance_scale,
+ config.image_height,
+ config.image_width,
+ config.steps,
+ config.prompt,
+ wandb.Image(generated_image_without_prompt_edit),
+ config.prompt_edit,
+ config.prompt_edit_token_weights,
+ wandb.Image(generated_image_with_prompt_edit)
+ )
+ wandb.log({
+ "Image Editing with Cross Attention Control": table
+ })
+
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **References:**
+ - https://arxiv.org/abs/2208.01626
+ - https://github.com/bloc97/CrossAttentionControl
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/datasets-predictions-image-classification-with-tables/datasets_predictions_image_classification_with_tables.py b/marimo/convert/datasets-predictions-image-classification-with-tables/datasets_predictions_image_classification_with_tables.py
new file mode 100644
index 00000000..05c19d1d
--- /dev/null
+++ b/marimo/convert/datasets-predictions-image-classification-with-tables/datasets_predictions_image_classification_with_tables.py
@@ -0,0 +1,507 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Image Classification with W&B Tables
+
+ This is a walkthrough of [Tables for visualization](https://docs.wandb.ai/guides/data-vis/tables) and [Artifacts for versioning](https://docs.wandb.com/artifacts) deep learning models in Weights & Biases. As an example, I finetune a convnet in Keras on photos from [iNaturalist 2017](https://github.com/visipedia/inat_comp/tree/master/2017) to identify 10 classes of living things (plants, insects, birds, etc).
+
+ 
+
+ ## [Explore more examples in this W&B Report](https://wandb.ai/stacey/mendeleev/reports/DSViz-for-Image-Classification--VmlldzozNjE3NjA)
+
+ ## Sign up or login
+
+ [Sign up or login](https://wandb.ai/login) to W&B to see and interact with your experiments in the browser.
+
+ In this example we're using Google Colab as a convenient hosted environment, but you can run your own training scripts from anywhere and visualize metrics with W&B's experiment tracking tool.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download sample data: Choose 1 of 4 sizes
+
+ Choose one of the three dataset size options below to run the rest of the demo. With fewer images, you'll run through the demo much faster and use less storage space. With more images, you'll get more realistic model training and more interesting results and examples to explore.
+
+ Note: **for the largest dataset, this stage might take a few minutes**. If you end up needing to rerun a cell, comment out the first capture line (change ```%%capture``` to ```#%%capture``` ) so you can respond to the prompt about re-downloading the dataset (and see the progress bar).
+
+ Each zipped directory contains randomly sampled images from the [iNaturalist dataset](https://github.com/visipedia/inat_comp), evenly distributed across 10 classes of living things like birds, insects, plants, and mammals (names given in Latin—so Aves, Insecta, Plantae, etc :).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # set SIZE to "TINY", "SMALL", "MEDIUM", or "LARGE"
+ # to select one of these three datasets
+ # TINY dataset: 100 images, 30MB
+ # SMALL dataset: 1000 images, 312MB
+ # MEDIUM dataset: 5000 images, 1.5GB
+ # LARGE dataset: 12,000 images, 3.6GB
+
+ SIZE = "SMALL"
+ return (SIZE,)
+
+
+@app.cell
+def _(SIZE):
+ if SIZE == "TINY":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_100.zip"
+ src_zip = "nature_100.zip"
+ DATA_SRC = "nature_100"
+ IMAGES_PER_LABEL = 10
+ BALANCED_SPLITS = {"train" : 8, "val" : 1, "test": 1}
+ elif SIZE == "SMALL":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_1K.zip"
+ src_zip = "nature_1K.zip"
+ DATA_SRC = "nature_1K"
+ IMAGES_PER_LABEL = 100
+ BALANCED_SPLITS = {"train" : 80, "val" : 10, "test": 10}
+ elif SIZE == "MEDIUM":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_12K.zip"
+ src_zip = "nature_12K.zip"
+ DATA_SRC = "inaturalist_12K/train" # (technically a subset of only 10K images)
+ IMAGES_PER_LABEL = 500
+ BALANCED_SPLITS = {"train" : 400, "val" : 50, "test": 50}
+ elif SIZE == "LARGE":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_12K.zip"
+ src_zip = "nature_12K.zip"
+ DATA_SRC = "inaturalist_12K/train" # (technically a subset of only 10K images)
+ IMAGES_PER_LABEL = 1000
+ BALANCED_SPLITS = {"train" : 800, "val" : 100, "test": 100}
+ return BALANCED_SPLITS, DATA_SRC, IMAGES_PER_LABEL
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !curl -SL $src_url > $src_zip
+ # !unzip $src_zip
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 0: Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Start out by installing the experiment tracking library and setting up your free W&B account:
+
+ * **pip install wandb** – Install the W&B library
+ * **import wandb** – Import the wandb library
+ * **wandb login** – Login to your W&B account so you can log all your metrics in one place
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(DATA_SRC, IMAGES_PER_LABEL):
+ import os
+ from random import shuffle
+ import numpy as np
+
+ # source directory for all raw data
+ SRC = DATA_SRC
+ PREFIX = "inat" # convenient for tracking local data
+ PROJECT_NAME = "nature_photos"
+
+ # number of images per class label
+ # the total number of images is 10X this (10 classes)
+ TOTAL_IMAGES = IMAGES_PER_LABEL * 10
+ return PREFIX, PROJECT_NAME, SRC, TOTAL_IMAGES, np, os, shuffle
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 1: Upload raw data
+ """)
+ return
+
+
+@app.cell
+def _(
+ IMAGES_PER_LABEL,
+ PREFIX,
+ PROJECT_NAME,
+ SRC,
+ TOTAL_IMAGES,
+ os,
+ shuffle,
+ wandb,
+):
+ # if this is a substantially new dataset, give it a new name
+ # this will create a whole new placeholder (Artifact) for this dataset
+ # instead of just incrementing a version of the old dataset
+ RAW_DATA_AT = '_'.join([PREFIX, 'raw_data', str(TOTAL_IMAGES)])
+ _run = wandb.init(project=PROJECT_NAME, job_type='upload')
+ # create an artifact for all the raw data
+ raw_data_at = wandb.Artifact(RAW_DATA_AT, type='raw_data')
+ _labels = os.listdir(SRC)
+ # SRC_DIR contains 10 folders, one for each of 10 class labels
+ # each folder contains images of the corresponding class
+ for _l in _labels:
+ _imgs_per_label = os.path.join(SRC, _l)
+ if os.path.isdir(_imgs_per_label):
+ _imgs = [i for i in os.listdir(_imgs_per_label) if not i.startswith('.DS')]
+ shuffle(_imgs) # filter out "DS_Store"
+ img_file_ids = _imgs[:IMAGES_PER_LABEL]
+ for f in img_file_ids: # randomize the order
+ file_path = os.path.join(SRC, _l, f)
+ raw_data_at.add_file(file_path, name=_l + '/' + f)
+ _run.log_artifact(raw_data_at)
+ # save artifact to W&B
+ _run.finish() # add file to artifact by full path
+ return (RAW_DATA_AT,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 2: Split raw data to prepare for training
+ """)
+ return
+
+
+@app.cell
+def _(
+ BALANCED_SPLITS,
+ PREFIX,
+ PROJECT_NAME,
+ RAW_DATA_AT,
+ SIZE,
+ TOTAL_IMAGES,
+ os,
+ shuffle,
+ wandb,
+):
+ # if this is a substantially different dataset, give it a new name
+ # this will create a whole new placeholder (Artifact) for this split
+ # instead of just incrementing a version of the old data split
+ SPLIT_DATA_AT = '_'.join([PREFIX, '80-10-10', str(TOTAL_IMAGES)])
+ _run = wandb.init(project=PROJECT_NAME, job_type='data_split')
+ SPLIT_COUNTS = BALANCED_SPLITS
+ # create balanced train, val, test splits
+ # each count is the number of images per label
+ data_at = _run.use_artifact(RAW_DATA_AT + ':latest')
+ data_dir = data_at.download()
+ # find the most recent ("latest") version of the full raw data
+ # you can of course pass around programmatic aliases and not string literals
+ # note: RAW_DATA_AT is defined in the previous cell—if you're running
+ # just this step, you may need to hardcode it
+ data_split_at = wandb.Artifact(SPLIT_DATA_AT, type='balanced_data')
+ # download it locally (for illustration purposes/across hardware; you can
+ # also sync/version artifacts by reference)
+ preview_dt = wandb.Table(columns=['id', 'image', 'label', 'split'])
+ _labels = os.listdir(data_dir)
+ for _l in _labels:
+ if _l.startswith('.'):
+ # create a table with columns we want to track/compare
+ continue
+ _imgs_per_label = os.listdir(os.path.join(data_dir, _l))
+ shuffle(_imgs_per_label)
+ start_id = 0
+ for split, count in SPLIT_COUNTS.items(): # skip non-label file
+ split_imgs = _imgs_per_label[start_id:start_id + count]
+ for img_file in split_imgs:
+ f_id = img_file.split('.')[0]
+ full_path = os.path.join(data_dir, _l, img_file)
+ data_split_at.add_file(full_path, name=os.path.join(split, _l, img_file))
+ if SIZE == 'LARGE': # take a subset
+ continue
+ if split != 'test':
+ preview_dt.add_data(f_id, wandb.Image(full_path), _l, split)
+ else:
+ preview_dt.add_data(f_id, wandb.Image(full_path), 'unknown', split) # add file to artifact by full path
+ start_id += count # note: pass the label to the name parameter to retain it in
+ data_split_at.add(preview_dt, 'data_split') # the data structure
+ _run.log_artifact(data_split_at)
+ # log artifact to W&B
+ _run.finish() # add a preview of the image # skip for the largest dataset for efficiency # pretend we have unlabeled test data # (replace "unknown" with l if you'd like to keep the labels :)
+ return data_split_at, preview_dt
+
+
+@app.cell
+def _(data_split_at, preview_dt):
+ # NOTE: if this Colab is running out of RAM, try running this cell
+ del data_split_at
+ del preview_dt
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 3: Train with artifacts and save model
+ """)
+ return
+
+
+@app.cell
+def _(BALANCED_SPLITS, PREFIX, PROJECT_NAME, TOTAL_IMAGES, np, os, wandb):
+ # EXPERIMENT CONFIG
+ #------------------------
+ # Core globals to modify
+ NUM_EPOCHS = 1 # set low for demo purposes, try 3, or 5, or as many as you like
+ _RUN_NAME = ''
+ VAL_TABLE_NAME = 'predictions'
+ # optional globals to modify
+ # set to a custom name to help keep your experiments organized
+ NUM_TRAIN = BALANCED_SPLITS['train'] * 10
+ # change this if you'd like start a new set of comparable Tables
+ # (only Tables logged to the same key can be compared)
+ NUM_VAL = BALANCED_SPLITS['val'] * 10
+ NUM_LOG_BATCHES = 16
+ # hyperparams set low for demo/training speed
+ # if you set these higher, be mindful of how many items are in
+ # the dataset artifacts you chose by setting the SIZE at the top
+ TRAIN_DATA_AT = PREFIX + '_80-10-10_' + str(TOTAL_IMAGES)
+ _MODEL_NAME = 'iv3_finetuned'
+ SAVE_MODEL_DIR = 'finetune_iv3_keras'
+ # enforced max for this is ceil(NUM_VAL/batch_size)
+ from tensorflow.keras.applications.inception_v3 import InceptionV3
+ from tensorflow.keras.callbacks import Callback
+ # ARTIFACTS CONFIG
+ from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
+ # training data artifact to load
+ from tensorflow.keras.models import Model
+ from tensorflow.keras.preprocessing.image import ImageDataGenerator
+ # model name
+ # if you want to train a sufficiently different model, give this a new name
+ # to start a new lineage for the model, instead of just incrementing the
+ # version of the old model
+ from wandb.keras import WandbCallback
+ CFG = {'num_train': NUM_TRAIN, 'num_val': NUM_VAL, 'num_classes': 10, 'fc_size': 1024, 'epochs': NUM_EPOCHS, 'batch_size': 32, 'img_width': 299, 'img_height': 299}
+ # folder in which to save the final, trained model
+ max_log_batches = int(np.ceil(float(CFG['num_val']) / float(CFG['batch_size'])))
+ CFG['num_log_batches'] = min(max_log_batches, NUM_LOG_BATCHES)
+
+ def finetune_inception_model(fc_size, num_classes):
+ """Load InceptionV3 with ImageNet weights, freeze it,
+ and attach a finetuning top for this classification task"""
+ base = InceptionV3(weights='imagenet', include_top='False')
+ for layer in base.layers:
+ layer.trainable = False
+ x = base.get_layer('mixed10').output
+ x = GlobalAveragePooling2D()(x)
+ x = Dense(fc_size, activation='relu')(x)
+ # experiment configuration saved to W&B
+ guesses = Dense(num_classes, activation='softmax')(x)
+ model = Model(inputs=base.input, outputs=guesses)
+ model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
+ return model
+
+ def train():
+ """ Main training loop which freezes the InceptionV3 layers of the model
+ and only trains the new top layers on the new data. A subsequent training
+ phase might unfreeze all the layers and finetune the whole model on the new data""" # inceptionV3 settings
+ _run = wandb.init(project=PROJECT_NAME, name=_RUN_NAME, job_type='train', config=CFG)
+ cfg = wandb.config
+ data_at = TRAIN_DATA_AT + ':latest'
+ data = _run.use_artifact(data_at, type='balanced_data')
+ # number of validation data batches to log/use when computing metrics
+ # at the end of each epoch
+ data_dir = data.download()
+ # change this min to max to log ALL the available images to a Table
+ train_dir = os.path.join(data_dir, 'train')
+ val_dir = os.path.join(data_dir, 'val')
+ train_datagen = ImageDataGenerator(rescale=1.0 / 255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True)
+ val_datagen = ImageDataGenerator(rescale=1.0 / 255)
+ train_generator = train_datagen.flow_from_directory(train_dir, target_size=(cfg.img_width, cfg.img_height), batch_size=cfg.batch_size, class_mode='categorical')
+ val_generator = val_datagen.flow_from_directory(val_dir, target_size=(cfg.img_width, cfg.img_height), batch_size=cfg.batch_size, class_mode='categorical', shuffle=False) # load InceptionV3 as base
+ model = finetune_inception_model(cfg.fc_size, cfg.num_classes)
+ callbacks = [WandbCallback(), ValLog(val_generator, cfg.num_log_batches)] # freeze base layers
+ model.fit(train_generator, steps_per_epoch=cfg.num_train // cfg.batch_size, epochs=cfg.epochs, validation_data=val_generator, callbacks=callbacks, validation_steps=cfg.num_val // cfg.batch_size)
+ trained_model_artifact = wandb.Artifact(_MODEL_NAME, type='model', description='finetuned inception v3', metadata=dict(cfg))
+ model.save(SAVE_MODEL_DIR)
+ trained_model_artifact.add_dir(SAVE_MODEL_DIR)
+ _run.log_artifact(trained_model_artifact) # attach a fine-tuning layer
+ _run.finish()
+
+ class ValLog(Callback):
+ """ Custom callback to log validation images
+ at the end of each training epoch"""
+
+ def __init__(self, generator=None, num_log_batches=1):
+ self.generator = generator
+ self.num_batches = num_log_batches
+ self.flat_class_names = [k for k, v in generator.class_indices.items()]
+
+ def on_epoch_end(self, epoch, logs={}):
+ val_data, val_labels = zip(*(self.generator[i] for i in range(self.num_batches)))
+ val_data, val_labels = (np.vstack(val_data), np.vstack(val_labels))
+ val_preds = self.model.predict(val_data)
+ true_ids = val_labels.argmax(axis=1) # locate and download training and validation data
+ max_preds = val_preds.argmax(axis=1)
+ columns = ['id', 'image', 'guess', 'truth']
+ for a in self.flat_class_names:
+ columns.append('score_' + a)
+ predictions_table = wandb.Table(columns=columns)
+ for filepath, img, top_guess, scores, truth in zip(self.generator.filenames, val_data, max_preds, val_preds, true_ids):
+ img_id = filepath.split('/')[-1].split('.')[0] # create train and validation data generators
+ row = [img_id, wandb.Image(img), self.flat_class_names[top_guess], self.flat_class_names[truth]]
+ for s in scores.tolist():
+ row.append(np.round(s, 4))
+ predictions_table.add_data(*row)
+ wandb.run.log({VAL_TABLE_NAME: predictions_table}) # instantiate model and callbacks # train! # save trained model as artifact # store full names of classes # collect validation data and ground truth labels from generator # use the trained model to generate predictions for the given number # of validation data batches (num_batches) # log validation predictions alongside the run # log image, predicted and actual labels, and all scores
+
+ return (train,)
+
+
+@app.cell
+def _(train):
+ train()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 4: Load model for inference
+ """)
+ return
+
+
+@app.cell
+def _(PREFIX, PROJECT_NAME, TOTAL_IMAGES, np, os, wandb):
+ _RUN_NAME = ''
+ TEST_TABLE_NAME = 'test_results'
+ from tensorflow import keras
+ from tensorflow.keras.preprocessing import image
+ _MODEL_NAME = 'iv3_finetuned'
+ TEST_DATA_AT = '_'.join([PREFIX, '80-10-10', str(TOTAL_IMAGES)])
+ _run = wandb.init(project=PROJECT_NAME, job_type='inference', name=_RUN_NAME)
+ model_at = _run.use_artifact(_MODEL_NAME + ':latest')
+ model_dir = model_at.download()
+ print('model: ', model_dir)
+ model = keras.models.load_model(model_dir)
+ model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
+ test_data_at = _run.use_artifact(TEST_DATA_AT + ':latest')
+ test_dir = test_data_at.download()
+ test_dir += '/test/'
+ class_names = ['Animalia', 'Amphibia', 'Arachnida', 'Aves', 'Fungi', 'Insecta', 'Mammalia', 'Mollusca', 'Plantae', 'Reptilia']
+ _imgs = []
+ filenames = []
+ class_labels = os.listdir(test_dir)
+ truth = []
+ for _l in class_labels:
+ if _l.startswith('.'):
+ continue
+ imgs_per_class = os.listdir(os.path.join(test_dir, _l))
+ for img in imgs_per_class:
+ filenames.append(img.split('.')[0])
+ truth.append(_l)
+ img_path = os.path.join(test_dir, _l, img)
+ img = image.load_img(img_path, target_size=(299, 299))
+ img = image.img_to_array(img)
+ img = np.expand_dims(img / 255.0, axis=0)
+ _imgs.append(img)
+ preds = {}
+ _imgs = np.vstack(_imgs)
+ classes = model.predict(_imgs, batch_size=32)
+ for c in classes:
+ class_id = np.argmax(c)
+ if class_id in preds:
+ preds[class_id] += 1
+ else:
+ preds[class_id] = 1
+ columns = ['id', 'image', 'guess', 'truth']
+ for a in class_names:
+ columns.append('score_' + a)
+ test_dt = wandb.Table(columns=columns)
+ for img_id, i, t, c in zip(filenames, _imgs, truth, classes):
+ guess = class_names[np.argmax(c)]
+ row = [img_id, wandb.Image(i), guess, t]
+ for c_i in c.tolist():
+ row.append(np.round(c_i, 4))
+ test_dt.add_data(*row)
+ _run.log({TEST_TABLE_NAME: test_dt})
+ print('Quick distribution of predicted classes: ')
+ print(preds)
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # More about Weights & Biases
+ We're always free for academics and open source projects. Email carey@wandb.com with any questions or feature suggestions. Here are some more resources:
+
+ 1. [Documentation](http://docs.wandb.com) - Python docs
+ 2. [Gallery](https://app.wandb.ai/gallery) - example reports in W&B
+ 3. [Articles](https://www.wandb.com/articles) - blog posts and tutorials
+ 4. [Community](wandb.me/slack) - join our Slack community forum
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/datasets-predictions-logging-timbre-transfer-with-w-b/datasets_predictions_logging_timbre_transfer_with_w_b.py b/marimo/convert/datasets-predictions-logging-timbre-transfer-with-w-b/datasets_predictions_logging_timbre_transfer_with_w_b.py
new file mode 100644
index 00000000..77443155
--- /dev/null
+++ b/marimo/convert/datasets-predictions-logging-timbre-transfer-with-w-b/datasets_predictions_logging_timbre_transfer_with_w_b.py
@@ -0,0 +1,587 @@
+# /// script
+# dependencies = ["ddsp", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Log timbre transfer audio experiments to W&B
+
+ Given some input audio (a microphone recording or a file upload), resynthesize the melody of the audio as if it were played on a violin, flute, trumpet, or tenor sax. Log all your experiments to an interactive W&B Table for easy exploration and tuning.
+
+ ### Source Colab
+
+ This notebook is a Weights & Biases integration and wrapper around the amazing [Timbre Transfer Demo with DDSP (Differentiable Digital Signal Processing) from Tensorflow Magenta](https://colab.research.google.com/github/magenta/ddsp/blob/master/ddsp/colab/demos/timbre_transfer.ipynb)
+
+ # Timbre Transfer with Interactive Visualization
+
+ The notebook processes audio input with timbre transfer, resynthesizing the melody using a model pretrained for various instruments (violin, flute, trumpet, etc).
+
+ ### [Explore an example with whale songs on W&B](https://wandb.ai/stacey/cshanty/reports/Whale2Song-W-B-Tables-for-Audio--Vmlldzo4NDI3NzM)
+
+ 
+
+ This notebook extracts features from input audio:
+ * uploaded files
+ * microphone recordings (to use this option, make sure to allow microphone access in your browser)
+ * URLs to sound files (hardcoded for this demo, feel free to edit the variable SONG_URL)
+
+ The available models are trained to generate audio conditioned on a time series of fundamental frequency and loudness. The input audio, synthesized song, and visualizations of the signal will be uploaded to an interactive W&B Table. You can experiment with different recordings, instruments, and various audio settings (using sliders) in this notebook. All of this configuraion will be organized alongside the song versions in one W&B project.
+
+ 
+
+ ## Additional Resources
+ * Full W&B Example: [Visualizing Audio Data with W&B Tables](https://wandb.ai/stacey/cshanty/reports/Whale2Song-W-B-Tables-for-Audio--Vmlldzo4NDI3NzM)
+ * [DDSP ICLR paper](https://openreview.net/forum?id=B1x1ma4tDr)
+ * [Audio Examples](http://goo.gl/magenta/ddsp-examples)
+ * marine mammal recordings from [Watkins Marine Mammal Sound Database](https://cis.whoi.edu/science/B/whalesounds/index.cfm), Woods Hole Oceanographic Institution
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 0.Dependencies and helper functions
+
+ Install dependencies and wandb, and download the model. The DDSP part transfers a lot of data and _should take a minute or two according to the source colab_. Also define helper functions to process audio data.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %tensorflow_version 2.x
+ print('Installing from pip package...')
+ # packages added via marimo's package management: ddsp==1.9.0 !pip install -qU ddsp==1.9.0
+ # packages added via marimo's package management: wandb !pip install -qqq wandb
+ import warnings
+ # Ignore a bunch of deprecation warnings
+ warnings.filterwarnings('ignore')
+ import copy
+ import os
+ import time
+ import crepe
+ import ddsp
+ import ddsp.training
+ from ddsp.colab import colab_utils
+ from ddsp.colab.colab_utils import audio_bytes_to_np, auto_tune, get_tuning_factor, download, play, record, specplot, upload, DEFAULT_SAMPLE_RATE
+ from ddsp.training.postprocessing import detect_notes, fit_quantile_transform
+ from ddsp import core
+ from ddsp import spectral_ops
+ import gin
+ from google.colab import files
+ import librosa
+ import matplotlib.pyplot as plt
+ from matplotlib import gridspec
+ import numpy as np
+ import pickle
+ import tensorflow.compat.v2 as tf
+ import tensorflow_datasets as tfds
+ import wandb
+ from urllib.request import urlretrieve
+ TRIM = -15
+ DEFAULT_SAMPLE_RATE = spectral_ops.CREPE_SAMPLE_RATE
+ print('Done!') # 16000
+ return (
+ DEFAULT_SAMPLE_RATE,
+ TRIM,
+ audio_bytes_to_np,
+ auto_tune,
+ core,
+ ddsp,
+ detect_notes,
+ files,
+ fit_quantile_transform,
+ get_tuning_factor,
+ gin,
+ librosa,
+ np,
+ os,
+ pickle,
+ play,
+ plt,
+ record,
+ spectral_ops,
+ tf,
+ time,
+ upload,
+ urlretrieve,
+ wandb,
+ )
+
+
+@app.cell
+def _(core, ddsp, librosa, np, plt, spectral_ops, time):
+ def process_song(audio, song_id, save_fig='_wave.png'):
+ ddsp.spectral_ops.reset_crepe() # Setup the session.
+ _start_time = time.time()
+ audio_features = ddsp.training.metrics.compute_audio_features(audio)
+ audio_features['loudness_db'] = audio_features['loudness_db'].astype(np.float32) # Compute features.
+ audio_features_mod = None
+ print('Audio features took %.1f seconds' % (time.time() - _start_time))
+ TRIM = -15
+ fig, ax = plt.subplots(nrows=3, ncols=1, sharex=True, figsize=(6, 8))
+ ax[0].plot(audio_features['loudness_db'][:TRIM])
+ ax[0].set_ylabel('loudness_db')
+ ax[1].plot(librosa.hz_to_midi(audio_features['f0_hz'][:TRIM]))
+ ax[1].set_ylabel('f0 [midi]') # Plot Features.
+ ax[2].plot(audio_features['f0_confidence'][:TRIM])
+ ax[2].set_ylabel('f0 confidence')
+ _ = ax[2].set_xlabel('Time step [frame]')
+ save_fig_path = song_id + save_fig
+ fig.savefig(save_fig_path)
+ return (audio_features, save_fig_path)
+
+ def specplot_local(audio, song_id, save_fig='_spec.png', vmin=-5, vmax=1, rotate=True, size=512 + 256, **matshow_kwargs):
+ """Plot the log magnitude spectrogram of audio."""
+ if len(audio.shape) == 2:
+ audio = audio[0]
+ logmag = spectral_ops.compute_logmag(core.tf_float32(audio), size=size)
+ if rotate:
+ logmag = np.rot90(logmag)
+ save_fig_path = song_id + save_fig
+ plt.imsave(save_fig_path, logmag, vmin=vmin, vmax=vmax, cmap=plt.cm.magma)
+ return save_fig_path # If batched, take first element. #plt.xticks([]) #plt.yticks([]) #plt.xlabel('Time') #plt.ylabel('Frequency')
+
+ return process_song, specplot_local
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 1.Initialize and login to W&B
+ """)
+ return
+
+
+@app.cell
+def _():
+ WANDB_PROJECT = "timbre_demo"
+ return (WANDB_PROJECT,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 2.Song setup (run for every new song)
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, np, wandb):
+ wandb.init(project=WANDB_PROJECT)
+
+ # generate one random song id, feel free to replace
+ SONG_ID = str(np.random.choice(1000, 1)[0])
+
+ # hardcoded to a favorite marine mammal melody, feel free to replace
+ SONG_URL = "https://whoicf2.whoi.edu/science/B/whalesounds/WhaleSounds/6301900Y.wav"
+ return SONG_ID, SONG_URL
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 3.Audio input: Record, upload, or URL
+
+ You have several options for audio input:
+ 1. **Record** audio from your microphone (NOTE: allow microphone access in your browser to do this)
+ 2. **Upload** audio from a file (.mp3 or .wav)
+ 3. **Download a URL** (this is hardcoded for the demo, and you can change SONG_URL in Step 2 to edit this)
+
+ Additional notes:
+ * Audio should be monophonic (single instrument / voice)
+ * Extracts fundmanetal frequency (f0) and loudness features.
+ """)
+ return
+
+
+@app.cell
+def _(SONG_URL, audio_bytes_to_np, np, record, upload, urlretrieve):
+ record_or_upload = "URL" #@param ["Record", "Upload (.mp3 or .wav)", "URL"]
+
+ record_seconds = 5#@param {type:"number", min:1, max:10, step:1}
+
+ if record_or_upload == "Record":
+ audio = record(seconds=record_seconds)
+ elif record_or_upload == "URL":
+ filename = SONG_URL.strip().split('/')[-1]
+ urlretrieve(SONG_URL, filename)
+ wav_bytes = open(filename, "rb").read()
+ audio = audio_bytes_to_np(wav_bytes)
+ else:
+ # Load audio sample here (.mp3 or .wav3 file)
+ # Just use the first file.
+ filenames, audios = upload()
+ audio = audios[0]
+
+ audio = audio[np.newaxis, :]
+ return (audio,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Upload sample song to W&B
+
+ You will see a URL to your W&B run, which will show a playable version of the song and some audio visualizations in a new Table.
+ """)
+ return
+
+
+@app.cell
+def _(
+ DEFAULT_SAMPLE_RATE,
+ SONG_ID,
+ audio,
+ np,
+ process_song,
+ specplot_local,
+ wandb,
+):
+ columns = ['id', 'orig_song', 'orig_plot', 'orig_spec']
+ audio_features, orig_waveplot = process_song(audio, SONG_ID, '_orig_plot.png')
+ orig_specplot = specplot_local(audio, SONG_ID, '_orig_spec.png')
+ orig_song = wandb.Audio(np.squeeze(audio), sample_rate=DEFAULT_SAMPLE_RATE)
+ _data = [[SONG_ID, orig_song, wandb.Image(orig_waveplot), wandb.Image(orig_specplot)]]
+ _table = wandb.Table(data=_data, columns=columns)
+ wandb.run.log({'sample_song': _table})
+ wandb.run.finish()
+ return (audio_features,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 4.Synthetic output (run for every new song)
+ """)
+ return
+
+
+@app.cell
+def _(
+ audio,
+ audio_features,
+ ddsp,
+ files,
+ gin,
+ os,
+ pickle,
+ subprocess,
+ tf,
+ time,
+):
+ #@title Choose an instrument (load a model)
+ #@markdown Run for every new audio input
+ model = 'Tenor_Saxophone' #@param ['Violin', 'Flute', 'Flute2', 'Trumpet', 'Tenor_Saxophone', 'Upload your own (checkpoint folder as .zip)']
+ MODEL = model
+
+ def find_model_dir(dir_name):
+ for root, dirs, filenames in os.walk(dir_name): # Iterate through directories until model directory is found
+ for filename in filenames:
+ if filename.endswith('.gin') and (not filename.startswith('.')):
+ model_dir = root
+ break
+ return model_dir
+ if model in ('Violin', 'Flute', 'Flute2', 'Trumpet', 'Tenor_Saxophone'):
+ PRETRAINED_DIR = '/content/pretrained'
+ subprocess.call(['rm', '-r', '$PRETRAINED_DIR', '&>', '/dev/null'])
+ subprocess.call(['mkdir', '$PRETRAINED_DIR', '&>', '/dev/null']) # Pretrained models.
+ GCS_CKPT_DIR = 'gs://ddsp/models/timbre_transfer_colab/2021-07-08'
+ model_dir = os.path.join(GCS_CKPT_DIR, 'solo_%s_ckpt' % model.lower()) # Copy over from gs:// for faster loading.
+ subprocess.call(['gsutil', 'cp', '$model_dir/*', '$PRETRAINED_DIR', '&>', '/dev/null']) #! rm -r $PRETRAINED_DIR &> /dev/null
+ model_dir = PRETRAINED_DIR
+ gin_file = os.path.join(model_dir, 'operative_config-0.gin') #! mkdir $PRETRAINED_DIR &> /dev/null
+ else:
+ UPLOAD_DIR = '/content/uploaded'
+ subprocess.call(['mkdir', '$UPLOAD_DIR'])
+ uploaded_files = files.upload()
+ for fnames in uploaded_files.keys(): #! gsutil cp $model_dir/* $PRETRAINED_DIR &> /dev/null
+ print('Unzipping... {}'.format(fnames))
+ subprocess.call(['unzip', '-o', '/content/$fnames', '-d', '$UPLOAD_DIR', '&>', '/dev/null'])
+ model_dir = find_model_dir(UPLOAD_DIR)
+ gin_file = os.path.join(model_dir, 'operative_config-0.gin')
+ DATASET_STATS = None
+ dataset_stats_file = os.path.join(model_dir, 'dataset_statistics.pkl') # User models.
+ print(f'Loading dataset statistics from {dataset_stats_file}')
+ try: #! mkdir $UPLOAD_DIR
+ if tf.io.gfile.exists(dataset_stats_file):
+ with tf.io.gfile.GFile(dataset_stats_file, 'rb') as f:
+ DATASET_STATS = pickle.load(f)
+ except Exception as err:
+ print('Loading dataset statistics from pickle failed: {}.'.format(err))
+ with gin.unlock_config(): #! unzip -o "/content/$fnames" -d $UPLOAD_DIR &> /dev/null
+ gin.parse_config_file(gin_file, skip_unknown=True)
+ ckpt_files = [f for f in tf.io.gfile.listdir(model_dir) if 'ckpt' in f]
+ ckpt_name = ckpt_files[0].split('.')[0]
+ ckpt = os.path.join(model_dir, ckpt_name)
+ time_steps_train = gin.query_parameter('F0LoudnessPreprocessor.time_steps')
+ # Load the dataset statistics.
+ n_samples_train = gin.query_parameter('Harmonic.n_samples')
+ hop_size = int(n_samples_train / time_steps_train)
+ time_steps = int(audio.shape[1] / hop_size)
+ n_samples = time_steps * hop_size
+ gin_params = ['Harmonic.n_samples = {}'.format(n_samples), 'FilteredNoise.n_samples = {}'.format(n_samples), 'F0LoudnessPreprocessor.time_steps = {}'.format(time_steps), 'oscillator_bank.use_angular_cumsum = True']
+ with gin.unlock_config():
+ gin.parse_config(gin_params)
+ for key in ['f0_hz', 'f0_confidence', 'loudness_db']:
+ audio_features[key] = audio_features[key][:time_steps]
+ audio_features['audio'] = audio_features['audio'][:, :n_samples]
+ model = ddsp.training.models.Autoencoder()
+ # Parse gin config,
+ model.restore(ckpt)
+ _start_time = time.time()
+ _ = model(audio_features, training=False)
+ # Assumes only one checkpoint in the folder, 'ckpt-[iter]`.
+ # Ensure dimensions and sampling rates are equal
+ # print("===Trained model===")
+ # print("Time Steps", time_steps_train)
+ # print("Samples", n_samples_train)
+ # print("Hop Size", hop_size)
+ # print("\n===Resynthesis===")
+ # print("Time Steps", time_steps)
+ # print("Samples", n_samples)
+ # print('')
+ # Trim all input vectors to correct lengths
+ # Set up the model just to predict audio given new conditioning
+ # Build model by running a batch through it.
+ print('Restoring model took %.1f seconds' % (time.time() - _start_time)) # Avoids cumsum accumulation errors.
+ return DATASET_STATS, MODEL, model
+
+
+@app.cell
+def _(
+ DATASET_STATS,
+ MODEL,
+ SONG_ID,
+ TRIM,
+ audio_features,
+ auto_tune,
+ ddsp,
+ detect_notes,
+ fit_quantile_transform,
+ get_tuning_factor,
+ librosa,
+ np,
+ plt,
+ wandb,
+):
+ #@title Modify conditioning
+ threshold = 1
+ #@markdown These models were not explicitly trained to perform timbre transfer, so they may sound unnatural if the incoming loudness and frequencies are very different then the training data (which will always be somewhat true).
+ ADJUST = True
+ quiet = 20
+ #@markdown ## Note Detection
+ autotune = 0
+ #@markdown You can leave this at 1.0 for most cases
+ pitch_shift = 0 #@param {type:"slider", min: 0.0, max:2.0, step:0.01}
+ loudness_shift = 0
+ SONG_CFG = {'threshold': threshold, 'adjust': ADJUST, 'quiet': quiet, 'autotune': autotune, 'pitch_shift': pitch_shift, 'loudness_shift': loudness_shift}
+ #@markdown ## Automatic
+ audio_features_mod = {k: v.copy() for k, v in audio_features.items()}
+ #@param{type:"boolean"}
+ def shift_ld(audio_features, ld_shift=0.0):
+ #@markdown Quiet parts without notes detected (dB)
+ """Shift loudness by a number of ocatves.""" #@param {type:"slider", min: 0, max:60, step:1}
+ audio_features['loudness_db'] = audio_features['loudness_db'] + ld_shift
+ #@markdown Force pitch to nearest note (amount)
+ return audio_features #@param {type:"slider", min: 0.0, max:1.0, step:0.1}
+
+ #@markdown ## Manual
+ def shift_f0(audio_features, pitch_shift=0.0):
+ """Shift f0 by a number of ocatves."""
+ #@markdown Shift the pitch (octaves)
+ audio_features['f0_hz'] = audio_features['f0_hz'] * 2.0 ** pitch_shift #@param {type:"slider", min:-2, max:2, step:1}
+ audio_features['f0_hz'] = np.clip(audio_features['f0_hz'], 0.0, librosa.midi_to_hz(110.0))
+ #@markdown Adjsut the overall loudness (dB)
+ return audio_features #@param {type:"slider", min:-20, max:20, step:1}
+ mask_on = None
+ # save settings
+ if ADJUST and DATASET_STATS is not None:
+ mask_on, note_on_value = detect_notes(audio_features['loudness_db'], audio_features['f0_confidence'], threshold)
+ if np.any(mask_on):
+ target_mean_pitch = DATASET_STATS['mean_pitch']
+ pitch = ddsp.core.hz_to_midi(audio_features['f0_hz'])
+ mean_pitch = np.mean(pitch[mask_on])
+ p_diff = target_mean_pitch - mean_pitch
+ p_diff_octave = p_diff / 12.0
+ round_fn = np.floor if p_diff_octave > 1.5 else np.ceil
+ p_diff_octave = round_fn(p_diff_octave)
+ audio_features_mod = shift_f0(audio_features_mod, p_diff_octave)
+ _, loudness_norm = fit_quantile_transform(audio_features['loudness_db'], mask_on, inv_quantile=DATASET_STATS['quantile_transform'])
+ ## Helper functions.
+ mask_off = np.logical_not(mask_on)
+ loudness_norm[mask_off] = loudness_norm[mask_off] - quiet * (1.0 - note_on_value[mask_off][:, np.newaxis])
+ loudness_norm = np.reshape(loudness_norm, audio_features['loudness_db'].shape)
+ audio_features_mod['loudness_db'] = loudness_norm
+ if autotune:
+ f0_midi = np.array(ddsp.core.hz_to_midi(audio_features_mod['f0_hz']))
+ tuning_factor = get_tuning_factor(f0_midi, audio_features_mod['f0_confidence'], mask_on)
+ f0_midi_at = auto_tune(f0_midi, tuning_factor, mask_on, amount=autotune)
+ audio_features_mod['f0_hz'] = ddsp.core.midi_to_hz(f0_midi_at)
+ else:
+ print('\nSkipping auto-adjust (no notes detected or ADJUST box empty).')
+ else:
+ print('\nSkipping auto-adujst (box not checked or no dataset statistics found).')
+ audio_features_mod = shift_ld(audio_features_mod, loudness_shift)
+ audio_features_mod = shift_f0(audio_features_mod, pitch_shift)
+ has_mask = int(mask_on is not None)
+ n_plots = 3 if has_mask else 2
+ fig, axes = plt.subplots(nrows=n_plots, ncols=1, sharex=True, figsize=(2 * n_plots, 8))
+ if has_mask: # Detect sections that are "on".
+ ax = axes[0]
+ ax.plot(np.ones_like(mask_on[:TRIM]) * threshold, 'k:')
+ ax.plot(note_on_value[:TRIM])
+ ax.plot(mask_on[:TRIM])
+ ax.set_ylabel('Note-on Mask')
+ ax.set_xlabel('Time step [frame]') # Shift the pitch register.
+ ax.legend(['Threshold', 'Likelihood', 'Mask'])
+ ax = axes[0 + has_mask]
+ ax.plot(audio_features['loudness_db'][:TRIM])
+ ax.plot(audio_features_mod['loudness_db'][:TRIM])
+ ax.set_ylabel('loudness_db')
+ ax.legend(['Original', 'Adjusted'])
+ ax = axes[1 + has_mask]
+ ax.plot(librosa.hz_to_midi(audio_features['f0_hz'][:TRIM]))
+ ax.plot(librosa.hz_to_midi(audio_features_mod['f0_hz'][:TRIM]))
+ ax.set_ylabel('f0 [midi]')
+ _ = ax.legend(['Original', 'Adjusted']) # Quantile shift the note_on parts.
+ final_wave_plot = SONG_ID + '_final.png'
+ fig.savefig(final_wave_plot)
+ SONG_CFG['final_wave_plot'] = wandb.Image(final_wave_plot)
+ # Manual Shifts.
+ # Plot Features.
+ SONG_CFG['instrument'] = MODEL # Turn down the note_off parts. # Auto-tune.
+ return SONG_CFG, audio_features_mod
+
+
+@app.cell
+def _(
+ DEFAULT_SAMPLE_RATE,
+ SONG_CFG,
+ SONG_ID,
+ audio,
+ audio_features,
+ audio_features_mod,
+ model,
+ np,
+ play,
+ specplot_local,
+ time,
+ wandb,
+):
+ af = audio_features if audio_features_mod is None else audio_features_mod
+ _start_time = time.time()
+ outputs = model(af, training=False)
+ audio_gen = model.get_audio_from_outputs(outputs)
+ print('Prediction took %.1f seconds' % (time.time() - _start_time))
+ print('Original')
+ play(audio)
+ orig_song_1 = wandb.Audio(np.squeeze(audio), sample_rate=DEFAULT_SAMPLE_RATE)
+ SONG_CFG['orig_song'] = orig_song_1
+ print('Resynthesis')
+ play(audio_gen)
+ synth_song = wandb.Audio(np.squeeze(audio_gen), sample_rate=DEFAULT_SAMPLE_RATE)
+ SONG_CFG['synth_song'] = synth_song
+ final_spec_plot = specplot_local(audio_gen, SONG_ID, fig_name='_final_spec.png')
+ SONG_CFG['final_spec'] = wandb.Image(final_spec_plot)
+ return orig_song_1, synth_song
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 5.Upload synthesized song to W&B
+ """)
+ return
+
+
+@app.cell
+def _(SONG_CFG, SONG_ID, WANDB_PROJECT, orig_song_1, synth_song, wandb):
+ wandb.init(project=WANDB_PROJECT)
+ output_columns = ['id', 'orig_song', 'synth_song', 'synth_waves', 'synth_spec', 'instrument', 'threshold', 'adjust', 'quiet', 'autotune', 'pitch_shift', 'loudness_shift']
+ s = SONG_CFG
+ _data = [[SONG_ID, orig_song_1, synth_song, s['final_wave_plot'], s['final_spec'], s['instrument'], s['threshold'], s['adjust'], s['quiet'], s['autotune'], s['pitch_shift'], s['loudness_shift']]]
+ _table = wandb.Table(data=_data, columns=output_columns)
+ wandb.run.log({'synth_song': _table})
+ wandb.run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Source colab: Timbre transfer demo from Magenta
+
+ This colab relies substantially on the following Timbre Transfer demo:
+
+ ##### Copyright 2021 Google LLC.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Copyright 2021 Google LLC. All Rights Reserved.
+
+ # Licensed under the Apache License, Version 2.0 (the "License");
+ # you may not use this file except in compliance with the License.
+ # You may obtain a copy of the License at
+
+ # http://www.apache.org/licenses/LICENSE-2.0
+
+ # Unless required by applicable law or agreed to in writing, software
+ # distributed under the License is distributed on an "AS IS" BASIS,
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ # See the License for the specific language governing permissions and
+ # limitations under the License.
+ # ==============================================================================
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/datasets-predictions-w-b-dataset-visualization/datasets_predictions_w_b_dataset_visualization.py b/marimo/convert/datasets-predictions-w-b-dataset-visualization/datasets_predictions_w_b_dataset_visualization.py
new file mode 100644
index 00000000..3937fd61
--- /dev/null
+++ b/marimo/convert/datasets-predictions-w-b-dataset-visualization/datasets_predictions_w_b_dataset_visualization.py
@@ -0,0 +1,363 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ _W&B Datasets & Predictions is currently in the early-access phase. You can use it in our production service at [wandb.ai](https://wandb.ai), with [some limitations](https://docs.wandb.com/datasets-and-predictions#current-limitations). APIs are subject to change. We'd love to hear questions, comments, and ideas! Drop us a line at feedback@wandb.com._
+
+ # WandB Dataset Visualization Demo
+
+ This notebook demonstrates WandB's dataset visualization features. In particular we will show how WandB [Artifacts](https://docs.wandb.com/artifacts) can be used to visualize datasets and predictions, with a focus on image data. We will track model and data lineage as well as perform interactive model analysis on the resulting datasets. The overall flow will be:
+
+ 1. Create a dataset
+ 2. Split the dataset into train and test
+ 3. Train a model to make predictions on the transformed dataet
+ 4. Log predications from the model against training and evaluation sets
+ 5. Analyze the model in WandB's UI
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 0: Setup
+
+ ## Install requirements & utils
+
+ For brevity, we put utility functions for working with the dataset in `util.py`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # # Install the python dependencies
+ # !pip install matplotlib numpy Pillow wandb
+ #
+ # # Download a util file of helper methods for this notebook
+ # !curl https://raw.githubusercontent.com/wandb/dsviz-demo/master/util.py --output util.py
+ return
+
+
+@app.cell
+def _():
+ # Colab sometimes has problems with Pillow. If you are facing this issue,
+ # uncomment the `exit()` line and run this cell. Then rerun the `!pip install`
+ # cell above.
+
+ # exit()
+ return
+
+
+@app.cell
+def _():
+ import util
+ import matplotlib.pyplot as plt
+ from PIL import Image
+ import os
+ import wandb
+ print(wandb.__version__)
+ return os, plt, util, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Login to wandb
+ """)
+ return
+
+
+@app.cell
+def _():
+ # default project name where results will be logged
+ WANDB_PROJECT = "dsviz-demo-colab"
+ NUM_EXAMPLES = 50
+ return NUM_EXAMPLES, WANDB_PROJECT
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Download the data
+
+ Before we get started, we will download an example dataset to our local machine. This is a big dataset, so please be patient if you are on a slow connection. For brevity, we put utility functions for working with the dataset in `util.py`. After the download is complete, we will show an example of the data.
+
+ **Note:** if you see the error "``AttributeError: module 'PIL.TiffTags' has no attribute 'IFD'``", this is likely a [Colab issue](https://github.com/facebookresearch/detectron2/issues/2231) which can be solved by restarting your runtime (header menu > Runtime > Restart runtime).
+ """)
+ return
+
+
+@app.cell
+def _(util):
+ # Download the data if not already present
+ util.download_data()
+ # Show an example training image
+ util.show_image(util.get_train_image_path(0))
+ # Show an example of color mask
+ util.show_image(util.get_color_label_image_path(0))
+
+ # Print the label types:
+ print("Class Mapping:")
+ print(list(zip(util.BDD_IDS, util.BDD_CLASSES)))
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 1: Build the dataset
+
+ First, let's build a dataset for use in the rest of this project. We will do this in the context of a `wandb.Run`. A `Run` is an isolated process which can optionally depend on upstream artifacts as well as optionally produce artifacts for later consumption. In this step, we will create a `wandb.Table` during our run and output it in an artifact. This table will contain all of our raw data for later use. Moreover, W&B offers rich tools to analyze and visualize such Tables in the interactive UI.
+ """)
+ return
+
+
+@app.cell
+def _(NUM_EXAMPLES, WANDB_PROJECT, util, wandb):
+ # Initialize the run
+ with wandb.init(project=WANDB_PROJECT, job_type='create_dataset', config={'num_examples': NUM_EXAMPLES, 'scale_factor': 2}) as _run:
+ class_set = wandb.Classes([{'name': name, 'id': id} for name, id in zip(util.BDD_CLASSES, util.BDD_IDS)]) # The project to register this Run to
+ table = wandb.Table(columns=['id', 'train_image', 'colored_image', 'label_mask', 'dominant_class']) # The type of this Run. Runs of the same type can be grouped together in the UI
+ for ndx in range(_run.config['num_examples']): # Custom configuration parameters which you might want to tune or adjust for the Run
+ example = wandb.Image(util.get_scaled_train_image(ndx, _run.config.scale_factor), classes=class_set, masks={'ground_truth': {'mask_data': util.get_scaled_mask_label(ndx, _run.config.scale_factor)}}, boxes={'ground_truth': {'box_data': util.get_scaled_bounding_boxes(ndx, _run.config.scale_factor)}}) # The number of raw samples to include.
+ color_label = wandb.Image(util.get_scaled_color_mask(ndx, _run.config.scale_factor)) # The scaling factor for the images
+ label_mask = wandb.Image(util.get_scaled_mask_label(ndx, _run.config.scale_factor))
+ table.add_data(util.train_ids[ndx], example, color_label, label_mask, util.get_dominant_class(label_mask))
+ _artifact = wandb.Artifact(name='raw_data', type='dataset') # Setup a WandB Classes object. This will give additional metadata for visuals
+ _artifact.add(table, 'raw_examples')
+ _run.log_artifact(_artifact)
+ print('Saving data to WandB...')
+ print('... Run Complete') # Setup a WandB Table object to hold our dataset # Fill up the table # First, we will build a wandb.Image to act as our raw example object # classes: the classes which map to masks and/or box metadata # masks: the mask metadata. In this case, we use a 2d array where each cell corresponds to the label (this comes directly from the dataset) # boxes: the bounding box metadata. For example sake, we create bounding boxes by looking at the mask data and creating boxes which fully enclose each class. # The data is an array of objects like: # "position": { # "minX": minX, # "maxX": maxX, # "minY": minY, # "maxY": maxY, # }, # "class_id" : id_num, # } # Next, we create two additional images which may be helpful during analysis. Notice that the additional metadata is optional. # Finally, we add a row of our newly constructed data. # Create an Artifact (versioned folder) # .add the table to the artifact # Finally, log the artifact
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Review the dataset in the Dashboard
+
+ Great, now if you click on the URL above, you should land on a run page. Since we did not log any metrics, there are no charts. Click the database icon (it looks like a stack of hockey pucks) on the left panel to see this run's artifacts. You should see something similar to the following:
+
+ 
+
+ Click on the "`raw_data`" row and navigate to the "Files" table. It should look like this:
+
+ 
+
+ Clicking on the "`raw_examples.table.json`" entry will launch an interactive data explorer to review the table we just built:
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 2: Splitting the data into train and test
+
+ Next, we will split the data into a train and a test dataset. Similar to before, we will launch a `Run` to perform this operation. Remember, this new execution could happen on a different machine as we will dynamically load the needed resources. In particular, we will lood in the raw dataset from the last run, and output 2 new datasets.
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, os, plt, wandb):
+ # This step should look familiar by now:
+ with wandb.init(project=WANDB_PROJECT, job_type='split_dataset', config={'train_pct': 0.7}) as _run:
+ _dataset_artifact = _run.use_artifact('raw_data:latest')
+ _data_table = _dataset_artifact.get('raw_examples')
+ print('\nExample Data row\n', _data_table.data[0])
+ print('\nExample Image\n')
+ plt.imshow(_data_table.data[0][1]._image)
+ plt.show()
+ print('\nArtifact Directory Contents: \n', os.listdir('artifacts'))
+ train_count = int(len(_data_table.data) * _run.config.train_pct) # Get the latest version of the artifact. Notice the name alias follows this convention: ":"
+ _train_table = wandb.Table(columns=_data_table.columns, data=_data_table.data[:train_count]) # When version is set to "latest", then the latest version will always be used.
+ _test_table = wandb.Table(columns=_data_table.columns, data=_data_table.data[train_count:]) # However, you can pin to a version by using an alias such as "raw_data:v0"
+ _train_artifact = wandb.Artifact('train_data', 'dataset')
+ _test_artifact = wandb.Artifact('test_data', 'dataset')
+ _train_artifact.add(_train_table, 'train_table') # Next, we .get the table by the same name that we saved it in the last run.
+ _test_artifact.add(_test_table, 'test_table')
+ _run.log_artifact(_train_artifact)
+ _run.log_artifact(_test_artifact) # Print a row
+ print('Saving data to WandB...')
+ print('... Run Complete') # Show an example image # Notice that a new directory was made: artifacts which is managed by wandb # Now we can build two separate artifacts for later use. We will first split the raw table into two parts, # then create two different artifacts, each of which will hold our new tables. We create two artifacts so that # in future runs, we can selectively decide which subsets of data to download. # Create the tables # Create the artifacts # Save the tables to the artifacts with .add # Log the artifacts out as outputs of the run
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Review the splits in the Dashboard
+ Notice, in this step, the raw_data `wandb.Table` was reinstatiated and the data, images, etc... came along for the ride. This makes it easy for ML practitioners on a team to share data and assets easily. To manage this, you can see that we created an artifacts directory to save local data.
+
+ Now we have two new datasets. Feel free to browse them similar to our last step. However, this time, click "Graph View" rather than "Files" to see the lineage of the artifact:
+
+ 
+
+ 
+
+ We will come back to this graph view later on!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 3: Model Training
+
+ Now we will train a model to predict bounding boxes. For the sake of simplicity, we will "train" a model which splits the image into it's grayscale quantiles and assigns labels to each patch. As you can imagine, the model performance can be improved dramatically.
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, util, wandb):
+ # Again, create a run.
+ with wandb.init(project=WANDB_PROJECT, job_type='model_train') as _run:
+ _train_artifact = _run.use_artifact('train_data:latest')
+ _train_table = _train_artifact.get('train_table') # Similar to before, we will load in the artifact and asset we need. In this case, the training data
+ train_data, _mask_data = util.make_datasets(_train_table, util.n_classes)
+ _model = util.ExampleSegmentationModel(util.n_classes)
+ _model.train(train_data, _mask_data)
+ _scores, _results = util.score_model(_model, train_data, _mask_data, util.n_classes) # Next, we split out the labels and train the model
+ results_table = wandb.Table(columns=['id', 'pred_mask', 'dominant_pred'] + util.BDD_CLASSES, data=[[_train_table.data[ndx][0], wandb.Image(_train_table.data[ndx][1], masks={'train_predicted_truth': {'mask_data': _results[ndx]}}, boxes={'ground_truth': {'box_data': util.mask_to_bounding(_results[ndx])}}), util.BDD_CLASSES[util.get_dominant_id_ndx(_results[ndx])]] + list(row) for ndx, row in enumerate(_scores)])
+ _results_artifact = wandb.Artifact('train_results', 'dataset')
+ _results_artifact.add(results_table, 'train_iou_score_table')
+ _run.log_artifact(_results_artifact)
+ _model.save('model.pkl') # Finally we score the model. Behind the scenes, we score each mask on its IOU score.
+ _model_artifact = wandb.Artifact('trained_model', 'model')
+ _model_artifact.add_file('model.pkl')
+ _run.log_artifact(_model_artifact) # Let's create a new table. Notice that we create many columns - an evaluation score for each class type.
+ print('Saving data to WandB...')
+ print('... Run Complete') # Data construction is similar to before, but we now use the predicted masks and bound boxes. # We create an artifact, add the table, and log it as part of the run. # Finally, let's save the model as a flat file and add that to its own artifact.
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 4: Model Evaluation
+
+ Now that we have a trained model, we want to score it on the test data which was held out in step 2. This code is very similar to the training step, with the execption of slightly different naming. The important difference is that we load the saved model from the artifact.
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, util, wandb):
+ with wandb.init(project=WANDB_PROJECT, job_type='model_eval') as _run:
+ _test_artifact = _run.use_artifact('test_data:latest')
+ _test_table = _test_artifact.get('test_table') # Retrieve the test data
+ test_data, _mask_data = util.make_datasets(_test_table, util.n_classes)
+ _model_artifact = _run.use_artifact('trained_model:latest')
+ path = _model_artifact.get_path('model.pkl').download()
+ _model = util.ExampleSegmentationModel.load(path)
+ _scores, _results = util.score_model(_model, test_data, _mask_data, util.n_classes) # Download the saved model file.
+ _results_artifact = wandb.Artifact('test_results', 'dataset')
+ data = [[_test_table.data[ndx][0], wandb.Image(_test_table.data[ndx][1], masks={'test_predicted_truth': {'mask_data': _results[ndx]}}, boxes={'ground_truth': {'box_data': util.mask_to_bounding(_results[ndx])}}), util.BDD_CLASSES[util.get_dominant_id_ndx(_results[ndx])]] + list(row) for ndx, row in enumerate(_scores)]
+ _results_artifact.add(wandb.Table(['id', 'pred_mask_test', 'dominant_pred_test'] + util.BDD_CLASSES, data=data), 'test_iou_score_table')
+ _run.log_artifact(_results_artifact) # Load the model from the file and score it
+ print('Saving data to WandB...')
+ print('... Run Complete') # Create a predicted score table similar to step 3. # And log out the results.
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 5: Model Analysis
+
+ This is where it all comes together. In this step, we join the train and test scoring results with the original dataset and output corresponding artifacts. The new idea introduced here is a `wandb.JoinedTable` which allows you to join two `Table`s for further analysis in the UI.
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, wandb):
+ with wandb.init(project=WANDB_PROJECT, job_type='model_result_analysis') as _run:
+ _dataset_artifact = _run.use_artifact('raw_data:latest')
+ _data_table = _dataset_artifact.get('raw_examples') # Retrieve the original raw dataset
+ _train_artifact = _run.use_artifact('train_results:latest')
+ _train_table = _train_artifact.get('train_iou_score_table')
+ _test_artifact = _run.use_artifact('test_results:latest')
+ _test_table = _test_artifact.get('test_iou_score_table') # Retrieve the train and test score tables
+ train_results = wandb.JoinedTable(_train_table, _data_table, 'id')
+ test_results = wandb.JoinedTable(_test_table, _data_table, 'id')
+ _artifact = wandb.Artifact('summary_results', 'dataset')
+ _artifact.add(train_results, 'train_results')
+ _artifact.add(test_results, 'test_results')
+ _run.log_artifact(_artifact)
+ print('Saving data to WandB...') # Join the tables on ID column and log them as outputs.
+ print('... Run Complete')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Review the model analysis in the Dashboard
+ Now, click on the above **Project** page (second link). This will look like the following:
+
+ 
+
+ Click on the database icon, as previously, to see the artifacts. This time, you are seeing the artifacts for the entire project, with counts of their versions:
+
+ 
+
+ Go ahead and click the "`model`" artifact type, "Files", and "`model.pkl`". The viewer will provide different renderings based on the file type. For a pickled class, you get the following image. For deep networks saved as `.h5` files, you can see all the layers and their attributes.
+
+ 
+
+ Next, head back to the artifact page, click Database type, expand `summary_results`, and select your most recent version. Click "Files" and select one of the join tables:
+
+ 
+
+ Exploring a bit, you can toggle the bounding boxes, masks, group, filter, and sort the data:
+
+ 
+
+ 
+
+ Finally, click graph view, and "explode". Now, you can visualize the entire process end-to-end:
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/datasets-predictions-w-b-tables-quickstart/datasets_predictions_w_b_tables_quickstart.py b/marimo/convert/datasets-predictions-w-b-tables-quickstart/datasets_predictions_w_b_tables_quickstart.py
new file mode 100644
index 00000000..8e97ba6d
--- /dev/null
+++ b/marimo/convert/datasets-predictions-w-b-tables-quickstart/datasets_predictions_w_b_tables_quickstart.py
@@ -0,0 +1,305 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # View & analyze model predictions during training
+
+ 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.
+ 
+
+ ## Compare model performance and find patterns
+
+ [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)
+
+ 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.
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Sign up or login
+
+ [Sign up or login](https://wandb.ai/login) to W&B to see and interact with your experiments in the browser.
+
+ In this example we're using Google Colab as a convenient hosted environment, but you can run your own training scripts from anywhere and visualize metrics with W&B's experiment tracking tool.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ log to your account
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ WANDB_PROJECT = "mnist-viz"
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 0. Setup
+
+ Install dependencies, download MNIST, and create train and test datasets using PyTorch.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import torch
+ import torch.nn as nn
+ import torchvision
+ import torchvision.transforms as T
+ import torch.nn.functional as F
+
+
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+ # create train and test dataloaders
+ def get_dataloader(is_train, batch_size, slice=5):
+ "Get a training dataloader"
+ 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=True, num_workers=2)
+ return loader
+
+ return F, get_dataloader, nn, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 1. Define the model and training schedule
+
+ * 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
+ """)
+ return
+
+
+@app.cell
+def _(get_dataloader, nn, torch):
+ # Number of epochs to run
+ # Each epoch includes a training step and a test step, so this sets
+ # the number of tables of test predictions to log
+ EPOCHS = 1
+ NUM_BATCHES_TO_LOG = 10
+ # Number of batches to log from the test data for each test step
+ # (default set low to simplify demo)
+ NUM_IMAGES_PER_BATCH = 32 #79
+ NUM_CLASSES = 10
+ # Number of images to log per test batch
+ BATCH_SIZE = 32
+ LEARNING_RATE = 0.001 #128
+ L1_SIZE = 32
+ # training configuration and hyperparameters
+ L2_SIZE = 64
+ CONV_KERNEL_SIZE = 5
+
+ class ConvNet(nn.Module):
+
+ # changing this may require changing the shape of adjacent layers
+ def __init__(self, num_classes=10):
+ super(ConvNet, self).__init__()
+ # define a two-layer convolutional neural network
+ self.layer1 = nn.Sequential(nn.Conv2d(1, L1_SIZE, CONV_KERNEL_SIZE, stride=1, padding=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.BatchNorm2d(L2_SIZE), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2))
+ self.fc = nn.Linear(7 * 7 * L2_SIZE, NUM_CLASSES)
+ self.softmax = nn.Softmax(NUM_CLASSES)
+
+ def forward(self, x):
+ out = self.layer1(x)
+ out = self.layer2(out)
+ out = out.reshape(out.size(0), -1)
+ out = self.fc(out)
+ return out
+ train_loader = get_dataloader(is_train=True, batch_size=BATCH_SIZE)
+ test_loader = get_dataloader(is_train=False, batch_size=2 * BATCH_SIZE)
+ device_1 = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # uncomment to see the shape of a given layer: #print("x: ", x.size())
+ return (
+ BATCH_SIZE,
+ CONV_KERNEL_SIZE,
+ ConvNet,
+ EPOCHS,
+ L1_SIZE,
+ L2_SIZE,
+ LEARNING_RATE,
+ NUM_BATCHES_TO_LOG,
+ NUM_CLASSES,
+ NUM_IMAGES_PER_BATCH,
+ device_1,
+ test_loader,
+ train_loader,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(
+ BATCH_SIZE,
+ CONV_KERNEL_SIZE,
+ ConvNet,
+ EPOCHS,
+ F,
+ L1_SIZE,
+ L2_SIZE,
+ LEARNING_RATE,
+ NUM_BATCHES_TO_LOG,
+ NUM_CLASSES,
+ NUM_IMAGES_PER_BATCH,
+ device_1,
+ nn,
+ test_loader,
+ torch,
+ train_loader,
+ wandb,
+):
+ # ✨ W&B: Initialize a new run to track this model's training
+ wandb.init(project='table-quickstart')
+ cfg = wandb.config
+ # ✨ W&B: Log hyperparameters using config
+ cfg.update({'epochs': EPOCHS, 'batch_size': BATCH_SIZE, 'lr': LEARNING_RATE, '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)})
+ model = ConvNet(NUM_CLASSES).to(device_1)
+ criterion = nn.CrossEntropyLoss()
+ optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+ def log_test_predictions(images, labels, outputs, predicted, test_table, log_counter):
+ # define model, loss, and optimizer
+ scores = F.softmax(outputs.data, dim=1)
+ log_scores = scores.cpu().numpy()
+ log_images = images.cpu().numpy()
+ log_labels = labels.cpu().numpy()
+ # convenience funtion to log predictions for a batch of test images
+ log_preds = predicted.cpu().numpy()
+ _id = 0 # obtain confidence scores for all classes
+ for i, l, p, s in zip(log_images, log_labels, log_preds, log_scores):
+ img_id = str(_id) + '_' + str(log_counter)
+ test_table.add_data(img_id, wandb.Image(i), p, l, *s)
+ _id = _id + 1
+ if _id == NUM_IMAGES_PER_BATCH:
+ break # adding ids based on the order of the images
+ total_step = len(train_loader)
+ for epoch in range(EPOCHS):
+ for i, (images, labels) in enumerate(train_loader): # add required info to data table:
+ images = images.to(device_1) # id, image pixels, model's guess, true label, scores for all classes
+ labels = labels.to(device_1)
+ outputs = model(images)
+ loss = criterion(outputs, labels)
+ optimizer.zero_grad()
+ loss.backward()
+ optimizer.step()
+ # train the model
+ wandb.log({'loss': loss})
+ if (i + 1) % 100 == 0:
+ print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, EPOCHS, i + 1, total_step, loss.item())) # training step
+ columns = ['id', 'image', 'guess', 'truth']
+ for digit in range(10):
+ columns.append('score_' + str(digit))
+ test_table = wandb.Table(columns=columns) # forward pass
+ model.eval()
+ log_counter = 0
+ with torch.no_grad(): # backward and optimize
+ correct = 0
+ total = 0
+ for images, labels in test_loader:
+ images = images.to(device_1)
+ labels = labels.to(device_1) # ✨ W&B: Log loss over training steps, visualized in the UI live
+ 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 = log_counter + 1
+ total = total + labels.size(0)
+ correct = correct + (predicted == labels).sum().item() # ✨ W&B: Create a Table to store predictions for each test step
+ acc = 100 * correct / total
+ wandb.log({'epoch': epoch, 'acc': acc})
+ print('Test Accuracy of the model on the 10000 test images: {} %'.format(acc))
+ wandb.log({'test_predictions': test_table})
+ # ✨ W&B: Mark the run as complete (useful for multi-cell notebook)
+ wandb.finish() # test the model # ✨ W&B: Log accuracy across training epochs, to visualize in the UI # ✨ W&B: Log predictions table to wandb
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/deepchem-w-b-x-deepchem/deepchem_w_b_x_deepchem.py b/marimo/convert/deepchem-w-b-x-deepchem/deepchem_w_b_x_deepchem.py
new file mode 100644
index 00000000..9a63ae9a
--- /dev/null
+++ b/marimo/convert/deepchem-w-b-x-deepchem/deepchem_w_b_x_deepchem.py
@@ -0,0 +1,458 @@
+# /// script
+# dependencies = ["deepchem", "dgl-cu110", "dgllife", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Introduction to Graph Convolutions
+
+ In this tutorial we will learn more about "graph convolutions." These are one of the most powerful deep learning tools for working with molecular data. The reason for this is that molecules can be naturally viewed as graphs.
+
+ 
+
+ Note how standard chemical diagrams of the sort we're used to from high school lend themselves naturally to visualizing molecules as graphs. In the remainder of this tutorial, we'll dig into this relationship in significantly more detail. This will let us get a deeper understanding of how these systems work.
+
+ ## Setup
+
+ To run DeepChem within Colab, you'll need to run the following installation commands. This will take about 5 minutes to run to completion and install your environment. You can of course run this tutorial locally if you prefer. In that case, don't run these cells since they will download and install Anaconda on your local machine.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Install Weights & Biases and log in.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: deepchem wandb !pip install -qU deepchem wandb
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ import warnings
+ warnings.filterwarnings('ignore')
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # What are Graph Convolutions?
+
+ Consider a standard convolutional neural network (CNN) of the sort commonly used to process images. The input is a grid of pixels. There is a vector of data values for each pixel, for example the red, green, and blue color channels. The data passes through a series of convolutional layers. Each layer combines the data from a pixel and its neighbors to produce a new data vector for the pixel. Early layers detect small scale local patterns, while later layers detect larger, more abstract patterns. Often the convolutional layers alternate with pooling layers that perform some operation such as max or min over local regions.
+
+ Graph convolutions are similar, but they operate on a graph. They begin with a data vector for each node of the graph (for example, the chemical properties of the atom that node represents). Convolutional and pooling layers combine information from connected nodes (for example, atoms that are bonded to each other) to produce a new data vector for each node.
+
+ # Training a GraphConvModel
+
+ Let's use the MoleculeNet suite to load the Tox21 dataset. To featurize the data in a way that graph convolutional networks can use, we set the featurizer option to `'GraphConv'`. The MoleculeNet call returns a training set, a validation set, and a test set for us to use. It also returns `tasks`, a list of the task names, and `transformers`, a list of data transformations that were applied to preprocess the dataset. (Most deep networks are quite finicky and require a set of data transformations to ensure that training proceeds stably.)
+ """)
+ return
+
+
+@app.cell
+def _():
+ import deepchem as dc
+ tasks, _datasets, transformers = dc.molnet.load_tox21(featurizer='GraphConv')
+ train_dataset, valid_dataset, test_dataset = _datasets
+ return dc, tasks, test_dataset, train_dataset, transformers, valid_dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We also need to evaluate the performance of the model while we are training. For this, we need to define a metric, a measure of model performance. `dc.metrics` holds a collection of metrics already. For this dataset, it is standard to use the ROC-AUC score, the area under the receiver operating characteristic curve (which measures the tradeoff between precision and recall). Luckily, the ROC-AUC score is already available in DeepChem.
+ """)
+ return
+
+
+@app.cell
+def _(dc):
+ metric = dc.metrics.Metric(dc.metrics.roc_auc_score)
+ return (metric,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will import and set up WandbLogger in order to log our information to Weights & Biases. WandbLogger by default will log training loss.
+
+ We also create a `ValidationCallback` to handle the validation scoring during training. At the interval specified, it will log the calculated metrics to Weights & Biases.
+ """)
+ return
+
+
+@app.cell
+def _(metric, transformers, valid_dataset):
+ from deepchem.models.wandblogger import WandbLogger
+ from deepchem.models.callbacks import ValidationCallback
+
+ wandblogger = WandbLogger(project='deepchem_graphconv', name='basic')
+ vc_valid = ValidationCallback(valid_dataset, interval=100, metrics=[metric], transformers=transformers)
+ return ValidationCallback, WandbLogger, vc_valid, wandblogger
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's now train a graph convolutional network on this dataset. DeepChem has the class `GraphConvModel` that wraps a standard graph convolutional architecture underneath the hood for user convenience. Let's instantiate an object of this class and train it on our dataset.
+ """)
+ return
+
+
+@app.cell
+def _(dc, tasks, train_dataset, vc_valid, wandblogger):
+ n_tasks = len(tasks)
+ model = dc.models.GraphConvModel(n_tasks, mode='classification', wandb_logger=wandblogger)
+ model.fit(train_dataset, nb_epoch=50, callbacks=[vc_valid])
+ wandblogger.finish()
+ return model, n_tasks
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To get the final performance of the model, we call `model.evaluate()`.
+ """)
+ return
+
+
+@app.cell
+def _(metric, model, test_dataset, train_dataset, transformers, valid_dataset):
+ train_score = model.evaluate(train_dataset, [metric], transformers)
+ valid_score = model.evaluate(valid_dataset, [metric], transformers)
+ test_score = model.evaluate(test_dataset, [metric], transformers)
+
+ print('Training set score:', train_score)
+ print('Validation set score:', valid_score)
+ print('Test set score:', test_score)
+ return test_score, train_score, valid_score
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can save our results in a wandb.Table for better visualization in the dashboard. Our table will compare the training, validation, and testing ROC-AUC score for three different models on the same dataset/task: a basic GCN, a custom GCN, and a Graph Attention Network.
+ """)
+ return
+
+
+@app.cell
+def _(test_score, train_score, valid_score, wandb):
+ columns = ["run_name", "train", "val", "test"]
+ metrics_table = wandb.Table(columns=columns)
+
+ # Add a row for the Basic GCN
+ metrics_table.add_data("Basic GCN", train_score, valid_score, test_score)
+ return (metrics_table,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The results are pretty good, and `GraphConvModel` is very easy to use. But what's going on under the hood? Could we build GraphConvModel ourselves? Of course! DeepChem provides Keras layers for all the calculations involved in a graph convolution. We are going to apply the following layers from DeepChem.
+
+ - `GraphConv` layer: This layer implements the graph convolution. The graph convolution combines per-node feature vectures in a nonlinear fashion with the feature vectors for neighboring nodes. This "blends" information in local neighborhoods of a graph.
+
+ - `GraphPool` layer: This layer does a max-pooling over the feature vectors of atoms in a neighborhood. You can think of this layer as analogous to a max-pooling layer for 2D convolutions but which operates on graphs instead.
+
+ - `GraphGather`: Many graph convolutional networks manipulate feature vectors per graph-node. For a molecule for example, each node might represent an atom, and the network would manipulate atomic feature vectors that summarize the local chemistry of the atom. However, at the end of the application, we will likely want to work with a molecule level feature representation. This layer creates a graph level feature vector by combining all the node-level feature vectors.
+
+ Apart from this we are going to apply standard neural network layers such as [Dense](https://keras.io/api/layers/core_layers/dense/), [BatchNormalization](https://keras.io/api/layers/normalization_layers/batch_normalization/) and [Softmax](https://keras.io/api/layers/activation_layers/softmax/) layer.
+ """)
+ return
+
+
+@app.cell
+def _(n_tasks):
+ from deepchem.models.layers import GraphConv, GraphPool, GraphGather
+ import tensorflow as tf
+ import tensorflow.keras.layers as layers
+
+ batch_size = 100
+
+ class MyGraphConvModel(tf.keras.Model):
+
+ def __init__(self):
+ super(MyGraphConvModel, self).__init__()
+ self.gc1 = GraphConv(128, activation_fn=tf.nn.tanh)
+ self.batch_norm1 = layers.BatchNormalization()
+ self.gp1 = GraphPool()
+
+ self.gc2 = GraphConv(128, activation_fn=tf.nn.tanh)
+ self.batch_norm2 = layers.BatchNormalization()
+ self.gp2 = GraphPool()
+
+ self.dense1 = layers.Dense(256, activation=tf.nn.tanh)
+ self.batch_norm3 = layers.BatchNormalization()
+ self.readout = GraphGather(batch_size=batch_size, activation_fn=tf.nn.tanh)
+
+ self.dense2 = layers.Dense(n_tasks*2)
+ self.logits = layers.Reshape((n_tasks, 2))
+ self.softmax = layers.Softmax()
+
+ def call(self, inputs):
+ gc1_output = self.gc1(inputs)
+ batch_norm1_output = self.batch_norm1(gc1_output)
+ gp1_output = self.gp1([batch_norm1_output] + inputs[1:])
+
+ gc2_output = self.gc2([gp1_output] + inputs[1:])
+ batch_norm2_output = self.batch_norm1(gc2_output)
+ gp2_output = self.gp2([batch_norm2_output] + inputs[1:])
+
+ dense1_output = self.dense1(gp2_output)
+ batch_norm3_output = self.batch_norm3(dense1_output)
+ readout_output = self.readout([batch_norm3_output] + inputs[1:])
+
+ logits_output = self.logits(self.dense2(readout_output))
+ return self.softmax(logits_output)
+
+ return MyGraphConvModel, batch_size
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can now see more clearly what is happening. There are two convolutional blocks, each consisting of a `GraphConv`, followed by batch normalization, followed by a `GraphPool` to do max pooling. We finish up with a dense layer, another batch normalization, a `GraphGather` to combine the data from all the different nodes, and a final dense layer to produce the global output.
+
+ Let's now create the DeepChem model which will be a wrapper around the Keras model that we just created. We will also specify the loss function so the model know the objective to minimize.
+ """)
+ return
+
+
+@app.cell
+def _(MyGraphConvModel, WandbLogger, dc):
+ wandblogger_1 = WandbLogger(project='deepchem_graphconv', name='custom')
+ model_1 = dc.models.KerasModel(MyGraphConvModel(), loss=dc.models.losses.CategoricalCrossEntropy(), wandb_logger=wandblogger_1)
+ return model_1, wandblogger_1
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ What are the inputs to this model? A graph convolution requires a complete description of each molecule, including the list of nodes (atoms) and a description of which ones are bonded to each other. In fact, if we inspect the dataset we see that the feature array contains Python objects of type `ConvMol`.
+ """)
+ return
+
+
+@app.cell
+def _(test_dataset):
+ test_dataset.X[0]
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Models expect arrays of numbers as their inputs, not Python objects. We must convert the `ConvMol` objects into the particular set of arrays expected by the `GraphConv`, `GraphPool`, and `GraphGather` layers. Fortunately, the `ConvMol` class includes the code to do this, as well as to combine all the molecules in a batch to create a single set of arrays.
+
+ The following code creates a Python generator that given a batch of data generates the lists of inputs, labels, and weights whose values are Numpy arrays. `atom_features` holds a feature vector of length 75 for each atom. The other inputs are required to support minibatching in TensorFlow. `degree_slice` is an indexing convenience that makes it easy to locate atoms from all molecules with a given degree. `membership` determines the membership of atoms in molecules (atom `i` belongs to molecule `membership[i]`). `deg_adjs` is a list that contains adjacency lists grouped by atom degree. For more details, check out the [code](https://github.com/deepchem/deepchem/blob/master/deepchem/feat/mol_graphs.py).
+ """)
+ return
+
+
+@app.cell
+def _(batch_size, n_tasks):
+ from deepchem.metrics import to_one_hot
+ from deepchem.feat.mol_graphs import ConvMol
+ import numpy as np
+
+ def data_generator(dataset, epochs=1):
+ for ind, (X_b, y_b, w_b, ids_b) in enumerate(dataset.iterbatches(batch_size, epochs,
+ deterministic=False, pad_batches=True)):
+ multiConvMol = ConvMol.agglomerate_mols(X_b)
+ inputs = [multiConvMol.get_atom_features(), multiConvMol.deg_slice, np.array(multiConvMol.membership)]
+ for i in range(1, len(multiConvMol.get_deg_adjacency_lists())):
+ inputs.append(multiConvMol.get_deg_adjacency_lists()[i])
+ labels = [to_one_hot(y_b.flatten(), 2).reshape(-1, n_tasks, 2)]
+ weights = [w_b]
+ yield (inputs, labels, weights)
+
+ return (data_generator,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we can train the model using `fit_generator(generator)` which will use the generator we've defined to train the model.
+ """)
+ return
+
+
+@app.cell
+def _(data_generator, model_1, train_dataset, wandblogger_1):
+ model_1.fit_generator(data_generator(train_dataset, epochs=50))
+ wandblogger_1.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now that we have trained our graph convolutional method, let's evaluate its performance. We again have to use our defined generator to evaluate model performance.
+ """)
+ return
+
+
+@app.cell
+def _(
+ data_generator,
+ metric,
+ metrics_table,
+ model_1,
+ test_dataset,
+ train_dataset,
+ transformers,
+ valid_dataset,
+):
+ train_score2 = model_1.evaluate_generator(data_generator(train_dataset), [metric], transformers)
+ valid_score2 = model_1.evaluate_generator(data_generator(valid_dataset), [metric], transformers)
+ test_score2 = model_1.evaluate_generator(data_generator(test_dataset), [metric], transformers)
+ metrics_table.add_data('Custom GCN', train_score2, valid_score2, test_score2)
+ # Add a row to the metrics table for our custom GCN
+ print('Training set score:', train_score2)
+ print('Validation set score:', valid_score2)
+ print('Test set score:', test_score2)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Success! The model we've constructed behaves nearly identically to `GraphConvModel`.
+
+ We can also use other graph models provided by Deepchem such as the [Graph Attention Model](https://deepchem.readthedocs.io/en/latest/api_reference/models.html#gatmodel).
+
+ In order to use it, we must first install DGL and DGL-LifeSci as specified in the GAT documentation.
+
+ Creating and training a GAT model follows the exact same process as the previous two models.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: dgl-cu110 !pip install --quiet dgl-cu110
+ # packages added via marimo's package management: dgllife !pip install --quiet dgllife
+ return
+
+
+@app.cell
+def _(ValidationCallback, WandbLogger, dc, metric):
+ from deepchem.models import GATModel
+ featurizer = dc.feat.MolGraphConvFeaturizer()
+ tasks_1, _datasets, transformers_1 = dc.molnet.load_tox21(reload=False, featurizer=featurizer, transformers=[])
+ train_dataset_1, valid_dataset_1, test_dataset_1 = _datasets
+ wandblogger_2 = WandbLogger(project='deepchem_graphconv', name='GAT')
+ model_2 = GATModel(mode='classification', n_tasks=len(tasks_1), batch_size=100, learning_rate=0.001, wandb_logger=wandblogger_2)
+ vc_valid_1 = ValidationCallback(valid_dataset_1, interval=100, metrics=[metric], transformers=transformers_1)
+ model_2.fit(train_dataset_1, nb_epoch=50, callbacks=[vc_valid_1])
+ return (
+ model_2,
+ test_dataset_1,
+ train_dataset_1,
+ transformers_1,
+ valid_dataset_1,
+ wandblogger_2,
+ )
+
+
+@app.cell
+def _(
+ metric,
+ metrics_table,
+ model_2,
+ test_dataset_1,
+ train_dataset_1,
+ transformers_1,
+ valid_dataset_1,
+):
+ train_score3 = model_2.evaluate(train_dataset_1, [metric], transformers_1)
+ valid_score3 = model_2.evaluate(valid_dataset_1, [metric], transformers_1)
+ test_score3 = model_2.evaluate(test_dataset_1, [metric], transformers_1)
+ metrics_table.add_data('Graph Attention Network', train_score3, valid_score3, test_score3)
+ # Add a row to our table for our GAT
+ print('Training set score:', train_score3)
+ print('Validation set score:', valid_score3)
+ print('Test set score:', test_score3)
+ return
+
+
+@app.cell
+def _(metrics_table, wandblogger_2):
+ # Log the final table to this run
+ wandblogger_2.wandb_run.log({'Scores': metrics_table})
+ return
+
+
+@app.cell
+def _(wandblogger_2):
+ wandblogger_2.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Congratulations! Time to join the Community!
+
+ Congratulations on completing this tutorial notebook! If you enjoyed working through the tutorial, and want to continue working with DeepChem, we encourage you to finish the rest of the tutorials in this series. You can also help the DeepChem community in the following ways:
+
+ ## Star DeepChem on [GitHub](https://github.com/deepchem/deepchem)
+ This helps build awareness of the DeepChem project and the tools for open source drug discovery that we're trying to build.
+
+ ## Join the DeepChem Gitter
+ The DeepChem [Gitter](https://gitter.im/deepchem/Lobby) hosts a number of scientists, developers, and enthusiasts interested in deep learning for the life sciences. Join the conversation!
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/diffusers-diffusers-image-generation/diffusers_diffusers_image_generation.py b/marimo/convert/diffusers-diffusers-image-generation/diffusers_diffusers_image_generation.py
new file mode 100644
index 00000000..b79ce633
--- /dev/null
+++ b/marimo/convert/diffusers-diffusers-image-generation/diffusers_diffusers_image_generation.py
@@ -0,0 +1,572 @@
+# /// script
+# dependencies = ["accelerate", "datasets", "diffusers", "ml-collections", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Unconditional Image Generation using 🤗 Diffusers + Weights & Biases 🪄🐝
+
+
+
+ **Reference:** [Official Diffusers Example for Unconditional Image Generation](https://github.com/huggingface/diffusers/tree/main/examples/unconditional_image_generation)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: accelerate diffusers datasets wandb ml_collections !pip install -qq accelerate diffusers datasets wandb ml_collections
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! accelerate config
+ subprocess.call(['accelerate', 'config'])
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! accelerate env
+ subprocess.call(['accelerate', 'env'])
+ return
+
+
+@app.cell
+def _():
+ import math
+ import os
+ from pathlib import Path
+ from typing import Optional
+
+ import torch
+ import torch.nn.functional as F
+ from torchvision.transforms import (
+ CenterCrop,
+ Compose,
+ InterpolationMode,
+ Normalize,
+ RandomHorizontalFlip,
+ Resize,
+ ToTensor,
+ )
+
+ from accelerate import Accelerator
+ from accelerate import notebook_launcher
+ from accelerate.logging import get_logger
+
+ from datasets import load_dataset
+
+ from diffusers import UNet2DModel
+ from diffusers import DDPMPipeline, DDPMScheduler
+ from diffusers import DDIMPipeline, DDIMScheduler
+ from diffusers.optimization import get_scheduler
+ from diffusers.training_utils import EMAModel
+
+ import matplotlib.pyplot as plt
+ from mpl_toolkits.axes_grid1 import ImageGrid
+
+ from ml_collections import ConfigDict
+ from tqdm.auto import tqdm
+ import wandb
+
+ return (
+ Accelerator,
+ CenterCrop,
+ Compose,
+ ConfigDict,
+ DDIMPipeline,
+ DDIMScheduler,
+ DDPMPipeline,
+ DDPMScheduler,
+ EMAModel,
+ F,
+ InterpolationMode,
+ Normalize,
+ RandomHorizontalFlip,
+ Resize,
+ ToTensor,
+ UNet2DModel,
+ get_scheduler,
+ load_dataset,
+ math,
+ notebook_launcher,
+ torch,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell
+def _(ConfigDict):
+ config = ConfigDict()
+
+ ##################### Dataset Configs #####################
+
+ # The name of the Dataset (from the HuggingFace hub or WandB Artifacts) to train on (could be your own,
+ # possibly private, dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,
+ # or to a folder containing files that HF Datasets can understand.
+ config.dataset_name = "geekyrakshit/diffusers-image-generation/anime-faces:v0" #@param {type:"string"}
+
+ # Is the config.dataset_name a Weights & Biase artifact or not
+ config.is_dataset_wandb_artifact = True #@param {type:"boolean"}
+
+ # The config of the Dataset, leave as None if there's only one config.
+ config.dataset_config_name = None #@param {type:"raw"}
+
+ # A folder containing the training data. Folder contents must follow the structure described in
+ # https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file
+ # must exist to provide the captions for the images. Ignored if `dataset_name` is specified.
+ config.train_data_dir = None #@param {type:"raw"}
+
+ # The output directory where the model predictions and checkpoints will be written.
+ config.output_dir = "ddpm-model-64" #@param {type:"string"}
+
+ # The directory where the downloaded models and datasets will be stored.
+ config.cache_dir = None #@param {type:"raw"}
+
+
+ ##################### Training Configs #####################
+
+ # Type of Diffusion pipeline
+ config.diffusion_pipeline = "ddim" #@param ["ddpm", "ddim"] {type:"string"}
+
+ # The resolution for input images, all the images in the train/validation dataset will be resized to
+ # this resolution.
+ config.resolution = 64 #@param {type:"slider", min:64, max:1024, step:4}
+
+ # Batch size (per device) for the training dataloader.
+ config.train_batch_size = 64 #@param {type:"slider", min:16, max:256, step:16}
+
+ # The number of images to generate for evaluation.
+ config.eval_batch_size = 64 #@param {type:"slider", min:16, max:256, step:16}
+
+ # The number of subprocesses to use for data loading. 0 means that the data will be loaded in the
+ # main process.
+ config.dataloader_num_workers = 0 #@param {type:"slider", min:0, max:16, step:1}
+
+ # Number of diffusion steps used to train the model.
+ config.num_train_timesteps = 1000 #@param {type:"slider", min:0, max:5000, step:100}
+
+ # Number of training epochs
+ config.num_epochs = 100 #@param {type:"slider", min:0, max:500, step:1}
+
+ # How often to save images during training.
+ config.save_images_epochs = 10 #@param {type:"slider", min:0, max:100, step:5}
+
+ # How often to save the model during training.
+ config.save_model_epochs = 10 #@param {type:"slider", min:0, max:100, step:5}
+
+ # Number of updates steps to accumulate before performing a backward/update pass.
+ config.gradient_accumulation_steps = 1 #@param {type:"slider", min:0, max:10, step:1}
+
+ # Initial learning rate (after the potential warmup period) to use.
+ config.learning_rate = 1e-4 #@param {type:"number"}
+
+ # The scheduler type to use. Choose between
+ # ["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"]
+ config.lr_scheduler = "cosine" #@param ["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"] {type:"string"}
+
+ # Number of steps for the warmup in the learning rate scheduler.
+ config.lr_warmup_steps = 500 #@param {type:"slider", min:0, max:1000, step:50}
+
+ # The exponential decay rate for the 1st moment estimates (the beta1 parameter for the Adam optimizer).
+ config.adam_beta1 = 0.95 #@param {type:"number"}
+
+ # The exponential decay rate for the 2nd moment estimates (the beta2 parameter for the Adam optimizer).
+ config.adam_beta2 = 0.999 #@param {type:"number"}
+
+ # Weight decay magnitude for the Adam optimizer.
+ config.adam_weight_decay = 1e-6 #@param {type:"number"}
+
+ # Epsilon value for the Adam optimizer.
+ config.adam_epsilon = 1e-08 #@param {type:"number"}
+
+ # Whether to use Exponential Moving Average for the final model weights.
+ config.use_ema = True #@param {type:"boolean"}
+
+ # The inverse gamma value for the EMA decay.
+ config.ema_inv_gamma = 1.0 #@param {type:"number"}
+
+ # The power value for the EMA decay.
+ config.ema_power = 3 / 4 #@param {type:"raw"}
+
+ # The maximum decay magnitude for EMA.
+ config.ema_max_decay = 0.9999 #@param {type:"number"}
+
+ # For distributed training: local_rank
+ config.local_rank = -1 #@param {type:"number"}
+
+ # Number of processes
+ config.num_processes = 1 #@param {type:"number"}
+
+ # Whether to use mixed precision.
+ # Choose between "no", "fp16" and "bf16" (bfloat16).
+ # Note that Bf16 requires PyTorch >= 1.10. and an Nvidia Ampere GPU.
+ config.mixed_precision = "no" #@param ["no", "fp16", "bf16"] {type:"raw"}
+
+
+ ##################### Weights & Biases Configs #####################
+
+ # Weights & Biases Project
+ config.wandb_project = "diffusers-image-generation" #@param {type:"string"}
+
+ # Weights & Biases Entity
+ config.wandb_entity = "geekyrakshit" #@param {type:"string"}
+
+ # Number of images to be visualized in a table
+ config.num_images_in_table = 6 #@param {type:"slider", min:1, max:50, step:1}
+ return (config,)
+
+
+@app.cell
+def _(UNet2DModel, config):
+ def build_unet_model():
+ return UNet2DModel(
+ sample_size=config.resolution,
+ in_channels=3,
+ out_channels=3,
+ layers_per_block=2,
+ block_out_channels=(128, 128, 256, 256, 512, 512),
+ down_block_types=(
+ "DownBlock2D",
+ "DownBlock2D",
+ "DownBlock2D",
+ "DownBlock2D",
+ "AttnDownBlock2D",
+ "DownBlock2D",
+ ),
+ up_block_types=(
+ "UpBlock2D",
+ "AttnUpBlock2D",
+ "UpBlock2D",
+ "UpBlock2D",
+ "UpBlock2D",
+ "UpBlock2D",
+ ),
+ )
+
+ return (build_unet_model,)
+
+
+@app.cell
+def _(
+ CenterCrop,
+ Compose,
+ InterpolationMode,
+ Normalize,
+ RandomHorizontalFlip,
+ Resize,
+ ToTensor,
+ config,
+ load_dataset,
+ torch,
+ wandb,
+):
+ def transforms(examples):
+ augmentations = Compose(
+ [
+ Resize(
+ config.resolution,
+ interpolation=InterpolationMode.BILINEAR
+ ),
+ CenterCrop(config.resolution),
+ RandomHorizontalFlip(),
+ ToTensor(),
+ Normalize([0.5], [0.5]),
+ ]
+ )
+ images = [
+ augmentations(image.convert("RGB"))
+ for image in examples["image"]
+ ]
+ return {"input": images}
+
+
+ def build_dataloader():
+ if not config.is_dataset_wandb_artifact:
+ dataset = (
+ load_dataset(
+ config.dataset_name,
+ config.dataset_config_name,
+ cache_dir=config.cache_dir,
+ split="train",
+ )
+ if config.dataset_name is not None else
+ load_dataset(
+ "imagefolder",
+ data_dir=config.train_data_dir,
+ cache_dir=config.cache_dir,
+ split="train"
+ )
+ )
+ else:
+ artifact = wandb.use_artifact(config.dataset_name, type='dataset')
+ artifact_dir = artifact.download()
+ config.train_data_dir = artifact_dir
+ dataset = load_dataset(
+ "imagefolder",
+ data_dir=config.train_data_dir,
+ cache_dir=config.cache_dir,
+ split="train"
+ )
+
+ dataset.set_transform(transforms)
+ return torch.utils.data.DataLoader(
+ dataset, batch_size=config.train_batch_size,
+ shuffle=True,
+ num_workers=config.dataloader_num_workers
+ )
+
+ return (build_dataloader,)
+
+
+@app.cell
+def _(
+ Accelerator,
+ DDIMPipeline,
+ DDIMScheduler,
+ DDPMPipeline,
+ DDPMScheduler,
+ EMAModel,
+ F,
+ build_dataloader,
+ build_unet_model,
+ config,
+ get_scheduler,
+ math,
+ torch,
+ tqdm,
+ wandb,
+):
+ def training_loop():
+ # Initialize Accelerator
+ accelerator = Accelerator(
+ gradient_accumulation_steps=config.gradient_accumulation_steps,
+ mixed_precision=config.mixed_precision,
+ log_with="wandb"
+ )
+
+ if accelerator.is_main_process:
+ accelerator.init_trackers(
+ project_name=config.wandb_project,
+ init_kwargs={
+ "wandb": {
+ 'entity': config.wandb_entity,
+ 'config': config.to_dict()
+ }
+ }
+ )
+ wandb_table = wandb.Table(
+ columns=['Epoch', 'Step', 'Generated-Images']
+ )
+
+ # Initialize Train Dataloader
+ train_dataloader = build_dataloader()
+
+ # Initialize Model
+ model = build_unet_model()
+
+ # Initialize Diffusion Pipeline
+ noise_scheduler = DDPMScheduler(
+ num_train_timesteps=config.num_train_timesteps
+ ) if config.diffusion_pipeline == "ddpm" else DDIMScheduler(
+ num_train_timesteps=config.num_train_timesteps
+ )
+
+ # Initialize AdamW optimizer
+ optimizer = torch.optim.AdamW(
+ model.parameters(),
+ lr=config.learning_rate,
+ betas=(config.adam_beta1, config.adam_beta2),
+ weight_decay=config.adam_weight_decay,
+ eps=config.adam_epsilon,
+ )
+
+ # Initialize Learning Rate Scheduler
+ lr_scheduler = get_scheduler(
+ config.lr_scheduler,
+ optimizer=optimizer,
+ num_warmup_steps=config.lr_warmup_steps,
+ num_training_steps=(
+ len(train_dataloader) * config.num_epochs
+ ) // config.gradient_accumulation_steps,
+ )
+
+ model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+ model, optimizer, train_dataloader, lr_scheduler
+ )
+
+ num_update_steps_per_epoch = math.ceil(
+ len(train_dataloader) / config.gradient_accumulation_steps
+ )
+
+ ema_model = EMAModel(
+ model,
+ inv_gamma=config.ema_inv_gamma,
+ power=config.ema_power,
+ max_value=config.ema_max_decay
+ )
+
+ global_step = 0
+ for epoch in range(config.num_epochs):
+ model.train()
+
+ progress_bar = tqdm(
+ total=num_update_steps_per_epoch,
+ disable=not accelerator.is_local_main_process
+ )
+ progress_bar.set_description(f"Epoch {epoch}")
+
+ for step, batch in enumerate(train_dataloader):
+ clean_images = batch["input"]
+ # Sample noise that we'll add to the images
+ noise = torch.randn(clean_images.shape).to(clean_images.device)
+ bsz = clean_images.shape[0]
+ # Sample a random timestep for each image
+ timesteps = torch.randint(
+ 0,
+ noise_scheduler.config.num_train_timesteps,
+ (bsz,),
+ device=clean_images.device
+ ).long()
+
+ # Add noise to the clean images according to the noise magnitude
+ # at each timestep (this is the forward diffusion process)
+ noisy_images = noise_scheduler.add_noise(
+ clean_images, noise, timesteps
+ )
+
+ with accelerator.accumulate(model):
+ # Predict the noise residual
+ noise_pred = model(noisy_images, timesteps).sample
+ loss = F.mse_loss(noise_pred, noise)
+ accelerator.backward(loss)
+
+ if accelerator.sync_gradients:
+ accelerator.clip_grad_norm_(model.parameters(), 1.0)
+ optimizer.step()
+ lr_scheduler.step()
+ if config.use_ema:
+ ema_model.step(model)
+ optimizer.zero_grad()
+
+ # Checks if the accelerator has performed an optimization step
+ # behind the scenes
+ if accelerator.sync_gradients:
+ progress_bar.update(1)
+ global_step += 1
+
+ logs = {
+ "loss": loss.detach().item(),
+ "lr": lr_scheduler.get_last_lr()[0],
+ "step": global_step
+ }
+ if config.use_ema:
+ logs["ema_decay"] = ema_model.decay
+ progress_bar.set_postfix(**logs)
+ accelerator.log(logs, step=global_step)
+
+ accelerator.log({'epoch':epoch}, step=global_step)
+ progress_bar.close()
+
+ accelerator.wait_for_everyone()
+
+ # Generate sample images for visual inspection
+ if accelerator.is_main_process:
+ if epoch % config.save_images_epochs == 0 or epoch == config.num_epochs - 1:
+ pipeline = DDPMPipeline(
+ unet=accelerator.unwrap_model(
+ ema_model.averaged_model if config.use_ema else model
+ ),
+ scheduler=noise_scheduler,
+ ) if config.diffusion_pipeline == "ddpm" else DDIMPipeline(
+ unet=accelerator.unwrap_model(
+ ema_model.averaged_model if config.use_ema else model
+ ),
+ scheduler=noise_scheduler,
+ )
+
+ generator = torch.manual_seed(0)
+ # run pipeline in inference (sample random noise and denoise)
+ images = pipeline(
+ generator=generator,
+ batch_size=config.eval_batch_size,
+ output_type="numpy"
+ ).images
+
+ # denormalize the images and save to wandb
+ images_processed = (images * 255).round().astype("uint8")
+ wandb_images = [wandb.Image(i) for i in images_processed]
+
+
+ wandb_table.add_data(
+ epoch,
+ global_step,
+ wandb_images[:config.num_images_in_table]
+ )
+
+ wandb.log({'generated_images':wandb_images,}, step=global_step)
+
+ if epoch % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
+ # save the model
+ pipeline.save_pretrained(config.output_dir)
+
+ # log wandb artifact
+ model_artifact = wandb.Artifact(
+ f'{wandb.run.id}-{config.output_dir}',
+ type='model'
+ )
+ model_artifact.add_dir(config.output_dir)
+ wandb.log_artifact(
+ model_artifact,
+ aliases=[f'step_{global_step}', f'epoch_{epoch}']
+ )
+
+ accelerator.wait_for_everyone()
+
+ wandb.log({'Generated-Images-Table': wandb_table})
+ accelerator.end_training()
+
+ return (training_loop,)
+
+
+@app.cell
+def _(config, notebook_launcher, training_loop):
+ notebook_launcher(training_loop, num_processes=config.num_processes)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/diffusers-lcm-diffusers/diffusers_lcm_diffusers.py b/marimo/convert/diffusers-lcm-diffusers/diffusers_lcm_diffusers.py
new file mode 100644
index 00000000..dc1c277c
--- /dev/null
+++ b/marimo/convert/diffusers-lcm-diffusers/diffusers_lcm_diffusers.py
@@ -0,0 +1,107 @@
+# /// script
+# dependencies = ["", "accelerate", "diffusers", "install-log", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Image Generation with Consistency Models using 🤗 Diffusers
+
+
+
+ This notebook demonstrates the following:
+ - Performing text-conditional image-generations with the [Consistency Models](https://huggingface.co/docs/diffusers/api/pipelines/consistency_models) using [🤗 Diffusers](https://huggingface.co/docs/diffusers).
+ - Manage image generation experiments using [Weights & Biases](http://wandb.ai/site).
+ - Log the prompts, generated images and experiment configs to [Weigts & Biases](http://wandb.ai/site) for visalization.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: diffusers transformers accelerate wandb > install.log !pip install diffusers transformers accelerate wandb > install.log
+ return
+
+
+@app.cell
+def _():
+ import random
+
+ import torch
+ from diffusers import DiffusionPipeline
+
+ import wandb
+ from wandb.integration.diffusers import autolog
+
+ return DiffusionPipeline, autolog, torch, wandb
+
+
+@app.cell
+def _(DiffusionPipeline, torch):
+ # Initialize the diffusion pipeline for latent consistency model
+ pipeline = DiffusionPipeline.from_pretrained("SimianLuo/LCM_Dreamshaper_v7")
+ pipeline = pipeline.to(torch_device="cuda", torch_dtype=torch.float32)
+ return (pipeline,)
+
+
+@app.cell
+def _(torch):
+ # Define the prompts, negative prompts, and seed.
+ prompt = [
+ "a photograph of an astronaut riding a horse",
+ "a photograph of a dragon"
+ ]
+
+ # Make the experiment reproducible by controlling randomness.
+ # The seed would be automatically logged to WandB.
+ generator = torch.Generator(device="cpu").manual_seed(10)
+ return generator, prompt
+
+
+@app.cell
+def _(autolog, generator, pipeline, prompt, wandb):
+ # Call WandB Autolog for Diffusers. This would automatically log
+ # the prompts, generated images, pipeline architecture and all
+ # associated experiment configs to Weights & Biases, thus making your
+ # image generation experiments easy to reproduce, share and analyze.
+ autolog(init=dict(project="diffusers_logging"))
+
+ # call the pipeline to generate the images
+ images = pipeline(
+ prompt,
+ num_images_per_prompt=2,
+ generator=generator,
+ num_inference_steps=10,
+ )
+
+ # End the experiment
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/diffusers-pixart-alpha-diffusers/diffusers_pixart_alpha_diffusers.py b/marimo/convert/diffusers-pixart-alpha-diffusers/diffusers_pixart_alpha_diffusers.py
new file mode 100644
index 00000000..4f714b33
--- /dev/null
+++ b/marimo/convert/diffusers-pixart-alpha-diffusers/diffusers_pixart_alpha_diffusers.py
@@ -0,0 +1,152 @@
+# /// script
+# dependencies = ["", "accelerate", "diffusers", "ftfy", "install-log", "sentencepiece", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Image Generation with Pixart-α using 🤗 Diffusers
+
+ This notebook demonstrates the following:
+ - Performing text-conditional image-generations with the [Pixart-α model](https://huggingface.co/docs/diffusers/v0.23.1/en/api/pipelines/pixart) using [🤗 Diffusers](https://huggingface.co/docs/diffusers).
+ - Manage image generation experiments using [Weights & Biases](http://wandb.ai/site).
+ - Log the prompts, generated images and experiment configs to [Weigts & Biases](http://wandb.ai/site) for visalization.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: diffusers transformers accelerate sentencepiece ftfy wandb > install.log !pip install diffusers transformers accelerate sentencepiece ftfy wandb > install.log
+ return
+
+
+@app.cell
+def _():
+ import random
+
+ import torch
+ from diffusers import PixArtAlphaPipeline
+
+ import wandb
+ from wandb.integration.diffusers import autolog
+
+ return PixArtAlphaPipeline, autolog, random, torch, wandb
+
+
+@app.cell
+def _(PixArtAlphaPipeline, torch):
+ # Load the pre-trained checkpoints from HuggingFace Hub to the PixArtAlphaPipeline
+ pipe = PixArtAlphaPipeline.from_pretrained(
+ "PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16
+ )
+
+ # Enable offloading the weights to the CPU and only loading them on the GPU when
+ # performing the forward pass can also save memory.
+ pipe.enable_model_cpu_offload()
+ return (pipe,)
+
+
+@app.cell
+def _(random, torch):
+ wandb_project = "pixart-alpha" # @param {type:"string"}
+
+ prompt = "a traveler navigating via a boat in countless mountains, Chinese ink painting" # @param {type:"string"}
+ negative_prompt = "" # @param {type:"string"}
+ num_inference_steps = 25 # @param {type:"slider", min:10, max:50, step:1}
+ guidance_scale = 4.5 # @param {type:"slider", min:0, max:10, step:0.1}
+ num_images_per_prompt = 1 # @param {type:"slider", min:0, max:10, step:0.1}
+ height = 1024 # @param {type:"slider", min:512, max:2560, step:32}
+ width = 1024 # @param {type:"slider", min:512, max:2560, step:32}
+ seed = None # @param {type:"raw"}
+
+
+ def autogenerate_seed():
+ max_seed = int(1024 * 1024 * 1024)
+ seed = random.randint(1, max_seed)
+ seed = -seed if seed < 0 else seed
+ seed = seed % max_seed
+ return seed
+
+
+ seed = autogenerate_seed() if seed is None else seed
+
+ # Make the experiment reproducible by controlling randomness.
+ # The seed would be automatically logged to WandB.
+ generator = torch.Generator(device="cuda").manual_seed(seed)
+ return (
+ generator,
+ guidance_scale,
+ height,
+ negative_prompt,
+ num_images_per_prompt,
+ num_inference_steps,
+ prompt,
+ wandb_project,
+ width,
+ )
+
+
+@app.cell
+def _(
+ autolog,
+ generator,
+ guidance_scale,
+ height,
+ negative_prompt,
+ num_images_per_prompt,
+ num_inference_steps,
+ pipe,
+ prompt,
+ wandb,
+ wandb_project,
+ width,
+):
+ # Call WandB Autolog for Diffusers. This would automatically log
+ # the prompts, generated images, pipeline architecture and all
+ # associated experiment configs to Weights & Biases, thus making your
+ # image generation experiments easy to reproduce, share and analyze.
+ autolog(init=dict(project=wandb_project))
+
+ # Generate the images by calling the PixArtAlphaPipeline
+ image = pipe(
+ prompt=prompt,
+ negative_prompt=negative_prompt,
+ num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale,
+ num_images_per_prompt=num_images_per_prompt,
+ height=height,
+ width=width,
+ generator=generator,
+ ).images[0]
+
+ # End the experiment
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/diffusers-sdxl-compel/diffusers_sdxl_compel.py b/marimo/convert/diffusers-sdxl-compel/diffusers_sdxl_compel.py
new file mode 100644
index 00000000..bc2d6691
--- /dev/null
+++ b/marimo/convert/diffusers-sdxl-compel/diffusers_sdxl_compel.py
@@ -0,0 +1,281 @@
+# /// script
+# dependencies = ["compel", "diffusers", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Prompt Weighing and Blending using for SDXL 1.0 using [Compel](https://github.com/damian0815/compel) and [🧨 Diffusers](https://huggingface.co/docs/diffusers)
+
+ This notebook demonstrates the following:
+ - Performing text-conditional image-generations using [🧨 Diffusers](https://huggingface.co/docs/diffusers).
+ - Using the Stable Diffusion XL Refiner pipeline to further refine the outputs of the base model.
+ - Manage image generation experiments using [Weights & Biases](http://wandb.ai/geekyrakshit).
+ - Log the prompts and generated images to [Weigts & Biases](http://wandb.ai/geekyrakshit) for visalization.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installing the Dependencies
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: diffusers[torch] transformers compel wandb !pip install -qq diffusers["torch"] transformers compel wandb
+ return
+
+
+@app.cell
+def _():
+ import torch
+ import wandb
+ from diffusers import DiffusionPipeline, EulerDiscreteScheduler
+ from compel import Compel, ReturnedEmbeddingsType
+
+ return (
+ Compel,
+ DiffusionPipeline,
+ EulerDiscreteScheduler,
+ ReturnedEmbeddingsType,
+ torch,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Experiment Management using Weights & Biases
+
+ Managing our image generation experiments is crucial for the sake of reproducibility. Hence we sync all the configs of our experiments with our Weights & Biases run. This stores all the configs of the experiments, right from the prompts to the refinement technque and the configuration of the scheduler.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project="stable-diffusion-xl", job_type="text-to-image-compel")
+
+ config = wandb.config
+ config.stable_diffusion_checkpoint = "stabilityai/stable-diffusion-xl-base-1.0"
+ config.refiner_checkpoint = "stabilityai/stable-diffusion-xl-refiner-1.0"
+ config.offload_to_cpu = False
+ config.compile_model = False
+ config.prompt_1 = "a cat playing with a ball in the (forest)---------"
+ config.prompt_2 = "Realistic, highly detailed, cold and bright color grading, 8k."
+ config.negative_prompt_1 = "low-quality"
+ config.negative_prompt_2 = "low-quality"
+ config.seed = 42
+ config.use_ensemble_of_experts = False
+ config.num_inference_steps = 100
+ config.num_refinement_steps = 150
+ config.high_noise_fraction = 0.8 # Set explicitly only if config.use_ensemble_of_experts is True
+ config.scheduler_kwargs = {
+ "beta_end": 0.012,
+ "beta_schedule": "scaled_linear", # one of ["linear", "scaled_linear"]
+ "beta_start": 0.00085,
+ "interpolation_type": "linear", # one of ["linear", "log_linear"]
+ "num_train_timesteps": 1000,
+ "prediction_type": "epsilon", # one of ["epsilon", "sample", "v_prediction"]
+ "steps_offset": 1,
+ "timestep_spacing": "leading", # one of ["linspace", "leading"]
+ "trained_betas": None,
+ "use_karras_sigmas": False,
+ }
+ config.prompt_credits = ""
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can make the experiment deterministic based on the seed specified in the experiment configs.
+ """)
+ return
+
+
+@app.cell
+def _(config, torch):
+ if config.seed is not None:
+ generator = [torch.Generator(device="cuda").manual_seed(config.seed)]
+ else:
+ generator = [torch.Generator(device="cuda")]
+ return (generator,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Creating the Diffusion Pipelines
+
+ For performing text-conditional image generation, we use the `diffusers` library to define the diffusion pipelines corresponding to the base SDXL model and the SDXL refinement model.
+ """)
+ return
+
+
+@app.cell
+def _(DiffusionPipeline, EulerDiscreteScheduler, config, torch):
+ pipe = DiffusionPipeline.from_pretrained(
+ config.stable_diffusion_checkpoint,
+ torch_dtype=torch.float16,
+ variant="fp16",
+ use_safetensors=True,
+ scheduler=EulerDiscreteScheduler(**config.scheduler_kwargs),
+ )
+
+ if config.offload_to_cpu:
+ pipe.enable_model_cpu_offload()
+ else:
+ pipe.to("cuda")
+
+ if config.compile_model:
+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+ return (pipe,)
+
+
+@app.cell
+def _(Compel, ReturnedEmbeddingsType, config, pipe, torch):
+ if config.prompt_2 == "" and config.negative_prompt_2 == "":
+ base_compel = Compel(
+ tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+ text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+ returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+ requires_pooled=[False, True]
+ )
+
+ base_positive_prompt_embeds, base_positive_prompt_pooled = base_compel(config.prompt)
+ base_negative_prompt_embeds, base_negative_prompt_pooled = base_compel(config.negative_prompt)
+ base_positive_prompt_embeds, base_negative_prompt_embeds = base_compel.pad_conditioning_tensors_to_same_length([
+ base_positive_prompt_embeds, base_negative_prompt_embeds
+ ])
+ else:
+ base_compel_1 = Compel(
+ tokenizer=pipe.tokenizer,
+ text_encoder=pipe.text_encoder,
+ returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+ requires_pooled=False,
+ )
+
+ base_positive_prompt_embeds_1 = base_compel_1(config.prompt_1)
+ base_negative_prompt_embeds_1 = base_compel_1(config.negative_prompt_1)
+
+ base_compel_2 = Compel(
+ tokenizer=pipe.tokenizer_2,
+ text_encoder=pipe.text_encoder_2,
+ returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+ requires_pooled=True,
+ )
+
+ base_positive_prompt_embeds_2, base_positive_prompt_pooled = base_compel_2(config.prompt_2)
+ base_negative_prompt_embeds_2, base_negative_prompt_pooled = base_compel_2(config.negative_prompt_2)
+
+ (
+ base_positive_prompt_embeds_2, base_negative_prompt_embeds_2
+ ) = base_compel_2.pad_conditioning_tensors_to_same_length([
+ base_positive_prompt_embeds_2, base_negative_prompt_embeds_2
+ ])
+
+ base_positive_prompt_embeds = torch.cat((base_positive_prompt_embeds_1, base_positive_prompt_embeds_2), dim=-1)
+ base_negative_prompt_embeds = torch.cat((base_negative_prompt_embeds_1, base_negative_prompt_embeds_2), dim=-1)
+ return (
+ base_negative_prompt_embeds,
+ base_negative_prompt_pooled,
+ base_positive_prompt_embeds,
+ base_positive_prompt_pooled,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Text-to-Image Generation
+
+ Now, we pass the embeddings and pooled prompts to the Stable Diffusion XL pipeline.
+ """)
+ return
+
+
+@app.cell
+def _(
+ base_negative_prompt_embeds,
+ base_negative_prompt_pooled,
+ base_positive_prompt_embeds,
+ base_positive_prompt_pooled,
+ config,
+ generator,
+ pipe,
+):
+ image = pipe(
+ prompt_embeds=base_positive_prompt_embeds,
+ pooled_prompt_embeds=base_positive_prompt_pooled,
+ negative_prompt_embeds=base_negative_prompt_embeds,
+ negative_pooled_prompt_embeds=base_negative_prompt_pooled,
+ output_type="pil",
+ num_inference_steps=config.num_inference_steps,
+ generator=generator,
+ ).images[0]
+ return (image,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Logging the Images to Weights & Biases
+
+ Now, we log the images to Weights & Biases. This enables us to:
+
+ - Visualize our generations
+ - Examine the generated images across different images
+ - Ensure reproducibility of the experiments
+ """)
+ return
+
+
+@app.cell
+def _(config, image, wandb):
+ table = wandb.Table(columns=['Prompt-1', 'Prompt-2', 'Negative-Prompt-1', 'Negative-Prompt-2', 'Generated-Image'])
+ image_1 = wandb.Image(image)
+ table.add_data(config.prompt_1, config.prompt_2, config.negative_prompt_1, config.negative_prompt_2, image_1)
+ wandb.log({'Generated-Image': image_1, 'Text-to-Image': table})
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's how you can control your prompts using Compel and manage them using Weights & Biases 👇
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/diffusers-sdxl-diffusers/diffusers_sdxl_diffusers.py b/marimo/convert/diffusers-sdxl-diffusers/diffusers_sdxl_diffusers.py
new file mode 100644
index 00000000..1f7e6c3f
--- /dev/null
+++ b/marimo/convert/diffusers-sdxl-diffusers/diffusers_sdxl_diffusers.py
@@ -0,0 +1,186 @@
+# /// script
+# dependencies = ["", "accelerate", "diffusers", "install-log", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Image Generation with Stable Diffusion XL using 🤗 Diffusers
+
+
+
+ This notebook demonstrates the following:
+ - Performing text-conditional image-generations with the [Stable Diffusion XL](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl) using [🤗 Diffusers](https://huggingface.co/docs/diffusers).
+ - Manage image generation experiments using [Weights & Biases](http://wandb.ai/site).
+ - Log the prompts, generated images and experiment configs to [Weigts & Biases](http://wandb.ai/site) for visalization.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: diffusers transformers accelerate wandb > install.log !pip install diffusers transformers accelerate wandb > install.log
+ return
+
+
+@app.cell
+def _():
+ import random
+
+ import torch
+ from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline
+
+ import wandb
+ from wandb.integration.diffusers import autolog
+
+ return (
+ StableDiffusionXLImg2ImgPipeline,
+ StableDiffusionXLPipeline,
+ autolog,
+ random,
+ torch,
+ wandb,
+ )
+
+
+@app.cell
+def _(StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline, torch):
+ base_model_id = "stabilityai/stable-diffusion-xl-base-1.0" # @param ["stabilityai/stable-diffusion-xl-base-1.0", "segmind/SSD-1B", "stabilityai/sdxl-turbo"]
+
+ base_pipeline = StableDiffusionXLPipeline.from_pretrained(
+ "stabilityai/stable-diffusion-xl-base-1.0",
+ torch_dtype=torch.float16,
+ variant="fp16",
+ use_safetensors=True,
+ )
+
+ base_pipeline.enable_model_cpu_offload()
+
+ refiner_pipeline = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+ "stabilityai/stable-diffusion-xl-refiner-1.0",
+ text_encoder_2=base_pipeline.text_encoder_2,
+ vae=base_pipeline.vae,
+ torch_dtype=torch.float16,
+ use_safetensors=True,
+ variant="fp16",
+ )
+ refiner_pipeline.enable_model_cpu_offload()
+ return base_pipeline, refiner_pipeline
+
+
+@app.cell
+def _(random, torch):
+ wandb_project = "pixart-alpha" # @param {type:"string"}
+
+ prompt_1 = "a photograph of an evil and vile looking demon in Bengali attire eating fish. The demon has large and bloody teeth. The demon is sitting on the branches of a giant Banyan tree, dimly lit, bluish and dark color palette, realistic, 8k" # @param {type:"string"}
+ prompt_2 = "" # @param {type:"string"}
+ negative_prompt_1 = "static, frame, painting, illustration, sd character, low quality, low resolution, greyscale, monochrome, nose, cropped, lowres, jpeg artifacts, deformed iris, deformed pupils, bad eyes, semi-realistic worst quality, bad lips, deformed mouth, deformed face, deformed fingers, deformed toes standing still, posing" # @param {type:"string"}
+ negative_prompt_2 = "static, frame, painting, illustration, sd character, low quality, low resolution, greyscale, monochrome, nose, cropped, lowres, jpeg artifacts, deformed iris, deformed pupils, bad eyes, semi-realistic worst quality, bad lips, deformed mouth, deformed face, deformed fingers, deformed toes standing still, posing" # @param {type:"string"}
+ num_inference_steps = 50 # @param {type:"slider", min:10, max:100, step:1}
+ guidance_scale = 5.0 # @param {type:"slider", min:0, max:10, step:0.1}
+ height = 1024 # @param {type:"slider", min:512, max:2560, step:32}
+ width = 1024 # @param {type:"slider", min:512, max:2560, step:32}
+ seed = None # @param {type:"raw"}
+
+
+ def autogenerate_seed():
+ max_seed = int(1024 * 1024 * 1024)
+ seed = random.randint(1, max_seed)
+ seed = -seed if seed < 0 else seed
+ seed = seed % max_seed
+ return seed
+
+
+ seed = autogenerate_seed() if seed is None else seed
+
+ # Make the experiment reproducible by controlling randomness.
+ # The seed would be automatically logged to WandB.
+ generator_base = torch.Generator(device="cuda").manual_seed(seed)
+ generator_refiner = torch.Generator(device="cuda").manual_seed(seed)
+ return (
+ generator_base,
+ generator_refiner,
+ guidance_scale,
+ negative_prompt_1,
+ negative_prompt_2,
+ num_inference_steps,
+ prompt_1,
+ prompt_2,
+ wandb_project,
+ )
+
+
+@app.cell
+def _(
+ autolog,
+ base_pipeline,
+ generator_base,
+ generator_refiner,
+ guidance_scale,
+ negative_prompt_1,
+ negative_prompt_2,
+ num_inference_steps,
+ prompt_1,
+ prompt_2,
+ refiner_pipeline,
+ wandb,
+ wandb_project,
+):
+ # Call WandB Autolog for Diffusers. This would automatically log
+ # the prompts, generated images, pipeline architecture and all
+ # associated experiment configs to Weights & Biases, thus making your
+ # image generation experiments easy to reproduce, share and analyze.
+ autolog(init=dict(project=wandb_project))
+
+ image = base_pipeline(
+ prompt=prompt_1,
+ prompt_2=prompt_2,
+ negative_prompt=negative_prompt_1,
+ negative_prompt_2=negative_prompt_2,
+ num_inference_steps=num_inference_steps,
+ output_type="latent",
+ generator=generator_base,
+ guidance_scale=guidance_scale,
+ ).images[0]
+
+ image = refiner_pipeline(
+ prompt=prompt_1,
+ prompt_2=prompt_2,
+ negative_prompt=negative_prompt_1,
+ negative_prompt_2=negative_prompt_2,
+ image=image[None, :],
+ num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale,
+ generator=generator_refiner,
+ ).images[0]
+
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/diffusers-sdxl-text-to-image/diffusers_sdxl_text_to_image.py b/marimo/convert/diffusers-sdxl-text-to-image/diffusers_sdxl_text_to_image.py
new file mode 100644
index 00000000..a7fd73b5
--- /dev/null
+++ b/marimo/convert/diffusers-sdxl-text-to-image/diffusers_sdxl_text_to_image.py
@@ -0,0 +1,306 @@
+# /// script
+# dependencies = ["diffusers", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Stable-Diffusion XL 1.0 using 🤗 Diffusers
+
+ This notebook demonstrates the following:
+ - Performing text-conditional image-generations using [🤗 Diffusers](https://huggingface.co/docs/diffusers).
+ - Using the Stable Diffusion XL Refiner pipeline to further refine the outputs of the base model.
+ - Manage image generation experiments using [Weights & Biases](http://wandb.ai/geekyrakshit).
+ - Log the prompts and generated images to [Weigts & Biases](http://wandb.ai/geekyrakshit) for visalization.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installing the Dependencies
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: diffusers[torch] transformers wandb !pip install -qq diffusers["torch"] transformers wandb
+ return
+
+
+@app.cell
+def _():
+ import torch
+ import wandb
+ from diffusers import (
+ StableDiffusionXLPipeline,
+ StableDiffusionXLImg2ImgPipeline,
+ EulerDiscreteScheduler
+ )
+
+ return (
+ EulerDiscreteScheduler,
+ StableDiffusionXLImg2ImgPipeline,
+ StableDiffusionXLPipeline,
+ torch,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Experiment Management using Weights & Biases
+
+ Managing our image generation experiments is crucial for the sake of reproducibility. Hence we sync all the configs of our experiments with our Weights & Biases run. This stores all the configs of the experiments, right from the prompts to the refinement technque and the configuration of the scheduler.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ project_name = "stable-diffusion-xl" # @param {type:"string"}
+
+ # initialize a wandb run
+ wandb.init(project=project_name, job_type="text-to-image")
+
+ # define experiment configs
+ config = wandb.config
+ config.stable_diffusion_checkpoint = "stabilityai/stable-diffusion-xl-base-1.0" # @param ["stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/stable-diffusion-xl-base-0.9"] {allow-input: true}
+ config.refiner_checkpoint = "stabilityai/stable-diffusion-xl-refiner-1.0" # @param ["stabilityai/stable-diffusion-xl-refiner-1.0", "stabilityai/stable-diffusion-xl-refiner-0.9"] {allow-input: true}
+ config.compile_model = False
+ config.prompt_1 = "a photograph of an evil and vile looking demon in Bengali attire eating fish. The demon has large and bloody teeth. The demon is sitting on the branches of a giant Banyan tree, dimly lit, bluish and dark color palette, realistic, 8k" # @param {type:"string"}
+ config.prompt_2 = "" # @param {type:"string"}
+ config.negative_prompt_1 = "static, frame, painting, illustration, sd character, low quality, low resolution, greyscale, monochrome, nose, cropped, lowres, jpeg artifacts, deformed iris, deformed pupils, bad eyes, semi-realistic worst quality, bad lips, deformed mouth, deformed face, deformed fingers, deformed toes standing still, posing" # @param {type:"string"}
+ config.negative_prompt_2 = "static, frame, painting, illustration, sd character, low quality, low resolution, greyscale, monochrome, nose, cropped, lowres, jpeg artifacts, deformed iris, deformed pupils, bad eyes, semi-realistic worst quality, bad lips, deformed mouth, deformed face, deformed fingers, deformed toes standing still, posing" # @param {type:"string"}
+ config.base_guidance_scale = 5.0 # @param {type:"slider", min:1, max:10, step:0.1}
+ config.seed = 0 # @param {type:"raw"}
+ config.num_inference_steps = 100 # @param {type:"slider", min:1, max:500, step:1}
+
+ config.enable_cpu_offload_base = True # @param {type:"boolean"}
+ config.enable_cpu_offload_refiner = True # @param {type:"boolean"}
+
+ config.compile_base_model = False # @param {type:"boolean"}
+
+ # Enable refinement only if high-ram instance
+ config.enable_refinement = False # @param {type:"boolean"}
+ config.compile_refinement_model = False # @param {type:"boolean"}
+ config.refiner_guidance_scale = 5.0 # @param {type:"slider", min:1, max:10, step:0.1}
+ config.num_refinement_steps = 150 # @param {type:"slider", min:1, max:500, step:1}
+
+ # Set explicitly only if config.use_ensemble_of_experts is True
+ config.high_noise_fraction = 0.8 # @param {type:"slider", min:0, max:1, step:0.1}
+
+ beta_schedule = "scaled_linear" # @param ["linear", "scaled_linear"]
+ interpolation_type = "linear" # @param ["linear", "log_linear"] {allow-input: true}
+ prediction_type = "epsilon" # @param ["epsilon", "sample", "v_prediction"]
+ timestep_spacing = "leading" # @param ["linspace", "leading"] {allow-input: true}
+
+ # configs for diffusers.EulerDiscreteScheduler
+ scheduler_kwargs = {
+ "beta_end": 0.012,
+ "beta_schedule": beta_schedule,
+ "beta_start": 0.00085,
+ "interpolation_type": interpolation_type,
+ "num_train_timesteps": 1000,
+ "prediction_type": prediction_type,
+ "steps_offset": 1,
+ "timestep_spacing": timestep_spacing,
+ "trained_betas": None,
+ "use_karras_sigmas": False,
+ }
+
+ config.scheduler_kwargs = scheduler_kwargs
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can make the experiment deterministic based on the seed specified in the experiment configs.
+ """)
+ return
+
+
+@app.cell
+def _(config, torch):
+ generator = [torch.Generator(device="cuda")]
+ if config.seed:
+ generator = [g.manual_seed(config.seed) for g in generator]
+ return (generator,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## The Base Diffusion Pipelines
+
+ For performing text-conditional image generation, we use the `diffusers` library to define the diffusion pipelines corresponding to the base SDXL model and the SDXL refinement model.
+
+ 1. We define the base diffusion pipeline using `diffusers.DiffusionPipeline` and load the pre-trained weights for SDXL 1.0 by calling the `from_pretrained` function on it. We also pass the scheduler as `diffusers.EulerDiscreteScheduler` in this step.
+
+ 2. In case we don't have a GPU with large enough GPU, it's recommended to enable CPU offloading. Otherwise, we load the model on the GPU. In case you're curious how HuggingFace manages CPU offloading in the most optimized manner, we recommend you read this port by [Sylvain Gugger](https://huggingface.co/sgugger): [How 🤗 Accelerate runs very large models thanks to PyTorch](https://huggingface.co/blog/accelerate-large-models).
+
+ 3. We can compile model using `torch.compile`, this might give a significant speedup.
+
+ 4. We generate the image from the prompts and negative prompts using the base pipeline.
+ """)
+ return
+
+
+@app.cell
+def _(
+ EulerDiscreteScheduler,
+ StableDiffusionXLPipeline,
+ config,
+ generator,
+ torch,
+):
+ # Define the Base Pipeline
+ pipe = StableDiffusionXLPipeline.from_pretrained(
+ config.stable_diffusion_checkpoint,
+ torch_dtype=torch.float16,
+ variant="fp16",
+ use_safetensors=True,
+ scheduler=EulerDiscreteScheduler(**config.scheduler_kwargs),
+ )
+
+ if config.enable_cpu_offload_base:
+ # Offload base pipeline to CPU
+ pipe.enable_model_cpu_offload()
+ else:
+ # Load base pipeline to GPU
+ pipe.to("cuda")
+
+ # Compile model using `torch.compile`, this might give a significant speedup
+ if config.compile_base_model:
+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+
+ # Generate image from the prompts and negative prompts using the base pipeline
+ generated_image = pipe(
+ prompt=config.prompt_1,
+ prompt_2=config.prompt_2,
+ negative_prompt=config.negative_prompt_1,
+ negative_prompt_2=config.negative_prompt_2,
+ guidance_scale=config.base_guidance_scale,
+ output_type="latent" if config.enable_refinement else "pil",
+ num_inference_steps=config.num_inference_steps,
+ denoising_end=config.high_noise_fraction if config.enable_refinement else None,
+ generator=generator,
+ ).images
+ return generated_image, pipe
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Refining the Generated Image
+
+ For refining the image generated by the base pipeline, we using the SDXL Refiner pipeline using the base and refiner model as an ensemble of expert of denoisers. In this case, the base model should serve as the expert for the high-noise diffusion stage and the refiner serves as the expert for the low-noise diffusion stage.
+
+ 1. We define the diffusion pipeline for the refiner using `diffusers.DiffusionPipeline` and load the pre-trained weights for SDXL 1.0 refiner by calling the `from_pretrained` function on it. We also pass the scheduler as `diffusers.EulerDiscreteScheduler` in this step.
+
+ 2. In case we don't have a GPU with large enough GPU, it's recommended to enable CPU offloading. Otherwise, we load the model on the GPU. In case you're curious how HiggingFace manages CPU offloading in the most optimized manner, we recommend you read this port by [Sylvain Gugger](https://huggingface.co/sgugger): [How 🤗 Accelerate runs very large models thanks to PyTorch](https://huggingface.co/blog/accelerate-large-models).
+
+ 3. We can compile model using `torch.compile`, this might give a significant speedup.
+
+ 4. We refine the latents generated by the base model from the same set of prompts and negative prompts using the refiner pipeline.
+ """)
+ return
+
+
+@app.cell
+def _(
+ EulerDiscreteScheduler,
+ StableDiffusionXLImg2ImgPipeline,
+ config,
+ generated_image,
+ generator,
+ pipe,
+ torch,
+):
+ if config.enable_refinement:
+ refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(config.refiner_checkpoint, text_encoder_2=pipe.text_encoder_2, vae=pipe.vae, torch_dtype=torch.float16, use_safetensors=True, variant='fp16', scheduler=EulerDiscreteScheduler(**config.scheduler_kwargs))
+ if config.enable_cpu_offload_refiner:
+ refiner.enable_model_cpu_offload()
+ else:
+ refiner.to('cuda')
+ if config.compile_refinement_model:
+ refiner.unet = torch.compile(pipe.unet, mode='reduce-overhead', fullgraph=True)
+ generated_image_1 = refiner(prompt=config.prompt_1, prompt_2=config.prompt_2, negative_prompt=config.negative_prompt_1, negative_prompt_2=config.negative_prompt_2, guidance_scale=config.refiner_guidance_scale, image=generated_image, num_inference_steps=config.num_refinement_steps, denoising_start=config.high_noise_fraction, generator=generator).images # Compile model using `torch.compile`, this might give a significant speedup
+ return (generated_image_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Logging the Images to Weights & Biases
+
+ Now, we log the images to Weights & Biases. This enables us to:
+
+ - Visualize our generations
+ - Examine the generated images across different images
+ - Ensure reproducibility of the experiments
+ """)
+ return
+
+
+@app.cell
+def _(config, generated_image_1, wandb):
+ # Create a [wandb table](https://docs.wandb.ai/guides/tables)
+ table = wandb.Table(columns=['Prompt-1', 'Prompt-2', 'Negative-Prompt-1', 'Negative-Prompt-2', 'Generated-Image'])
+ generated_image_2 = wandb.Image(generated_image_1[0])
+ table.add_data(config.prompt_1, config.prompt_2, config.negative_prompt_1, config.negative_prompt_2, generated_image_2)
+ wandb.log({'Generated-Image': generated_image_2, 'Text-to-Image': table})
+ # Add the images to the table
+ # Log the images and table to wandb
+ # finish the experiment
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's how you can examine your generations across multiple experiments 👇
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's how you can manage your prompts and your generations across experiments 👇
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/gemini-google-demo-rag/gemini_google_demo_rag.py b/marimo/convert/gemini-google-demo-rag/gemini_google_demo_rag.py
new file mode 100644
index 00000000..ca2e4296
--- /dev/null
+++ b/marimo/convert/gemini-google-demo-rag/gemini_google_demo_rag.py
@@ -0,0 +1,304 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install google-generativeai weave -qqU
+ return
+
+
+@app.cell
+def _():
+ import google.generativeai as genai
+ import weave
+
+ return genai, weave
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Set up your Google API key and log into W&B Weave
+
+ To run the following cell, your API key must be stored it in a Colab Secret named `GOOGLE_API_KEY`. If you don't already have an API key, or you're not sure how to create a Colab Secret, see the [Authentication](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) quickstart for an example.
+ """)
+ return
+
+
+@app.cell
+def _(genai):
+ from google.colab import userdata
+ GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')
+ genai.configure(api_key=GOOGLE_API_KEY)
+ return
+
+
+@app.cell
+def _():
+ import asyncio
+ from weave import Model, Evaluation, Dataset
+ import numpy as np
+
+ return Model, np
+
+
+@app.cell
+def _(weave):
+ # We call init to begin capturing data in the project, intro-example.
+ weave.init("prompt-eng/gemini-rag")
+ return
+
+
+@app.cell
+def _(genai, np):
+ from google.api_core import retry
+ emb_model = 'models/embedding-001'
+
+ def make_embed_text_fn(model):
+ @retry.Retry(timeout=300.0)
+ def embed_fn(text: str) -> list[float]:
+ embedding = genai.embed_content(model=model,
+ content=text,
+ task_type="retrieval_document")['embedding']
+ return np.array(embedding)
+ return embed_fn
+
+ return emb_model, make_embed_text_fn
+
+
+@app.cell
+def _(emb_model, make_embed_text_fn, weave):
+ @weave.op()
+ def docs_to_embeddings(docs: list) -> list:
+ # Convert documents to embeddings
+ document_embeddings = []
+ emb_fn = make_embed_text_fn(emb_model)
+ for doc in docs:
+ emb = emb_fn(doc)
+ document_embeddings.append(emb)
+ return document_embeddings
+
+ return (docs_to_embeddings,)
+
+
+@app.cell
+def _(emb_model, genai, np, weave):
+ @weave.op()
+ def get_most_relevant_document(query, docs, document_embeddings):
+ # Convert query to embedding
+ query_embedding = genai.embed_content(model=emb_model,
+ content=query,
+ task_type="retrieval_query")['embedding']
+ # Compute cosine similarity
+ similarities = np.dot(np.stack(document_embeddings), query_embedding)
+ # Get the index of the most similar document
+ most_relevant_doc_index = np.argmax(similarities)
+ return docs[most_relevant_doc_index]
+
+ return (get_most_relevant_document,)
+
+
+@app.cell
+def _(Model, genai, get_most_relevant_document, weave):
+ # define the Morgan Stanley Research RAG Model
+ class MSResearchRAGModel(Model):
+ system_message: str
+ model_name: str = 'models/gemini-pro'
+
+ @weave.op()
+ def predict(self, question: str, docs: list, add_context: bool) -> dict:
+ model = genai.GenerativeModel(self.model_name)
+
+ RAG_Context = ""
+ # Retrieve the embeddings artifact
+ embeddings = weave.ref("MSRAG_Embeddings").get()
+
+ if add_context:
+ # Using Google Embeddings, get the relevant document for context
+ RAG_Context = get_most_relevant_document(question, docs, embeddings)
+
+ query = f"""Use the following information to answer the subsequent question. If the answer cannot be found, write "I don't know."
+
+ Context from Morgan Stanley Research:
+ \"\"\"
+ {RAG_Context}
+ \"\"\"
+
+ Question: {question}"""
+ prompt = f'{self.system_message}\n\n{query}'
+ response = model.generate_content(prompt)
+ model_output = response.text
+
+ return model_output
+
+ return (MSResearchRAGModel,)
+
+
+@app.cell
+def _(MSResearchRAGModel, docs_to_embeddings, weave):
+ model = MSResearchRAGModel(
+ system_message="You are an expert in finance and answer questions related to finance, financial services, and financial markets. When responding based on provided information, be sure to cite the source."
+ )
+
+ contexts = [
+ f"""Morgan Stanley has moved in new market managers on the East and West Coasts as part of changes that sent some of other management veterans into new roles, according to two sources.
+
+ On the West Coast, Ken Sullivan, a 37-year-industry veteran who came to Morgan Stanley five years ago from RBC Wealth Management, has assumed an expanded role as market executive for a consolidated Beverly Hills and Los Angeles market, according to a source.
+
+ Meanwhile, Greg Laetsch, a 44-year industry veteran who had been the complex manager in Los Angeles for the last 15 years, has moved to a non-producing senior advisor role to the LA market for Morgan Stanley, according to the same source.
+
+ On the East Coast, Morgan Stanley hired Nikolas Totaro, a 19-year industry veteran, from Merrill Lynch, where he had worked for 14 years and had been most recently a market executive in Greenwich, Connecticut. Totaro will be a market manager reporting to John Palazzetti in the Midtown Wealth Management Center in Manhattan, according to the same source.
+
+ Totaro is replacing Bill DeMatteo, a 21-year industry veteran who spent the last 14 years at Morgan Stanley, and who has returned to full-time production. DeMatteo has joined the Continuum Group at Morgan Stanley, which Barron’s ranked 20th among on its 2022 Top 100 Private Wealth Management Teams and listed as managing $7.2 billion in client assets.
+
+ “His extensive 17 years of management experience at Morgan Stanley will be instrumental in shaping our approach to wealth management, fostering client relationships, and steering the team towards sustained growth and success,” Scott Siegel, leader of the Continuum Group, wrote on LinkedIn.
+
+ Totaro and Laestch did not respond immediately to requests for comments sent through LinkedIn. Sullivan did not respond immediately to an emailed request. Both Morgan Stanley and Merrill spokespersons declined to comment about the changes.
+
+ Totaro’s former Southern Connecticut market at Merrill included over 325 advisors and support staff across six offices in Greenwich, Stamford, Darien, Westport, Fairfield, and New Canaan, according to his LinkedIn profile.
+
+ Separately, a former Raymond James Financial divisional director has joined Janney Montgomery Scott in Ponte Vedra Beach, Florida. Tom M. Galvin, who has spent the last 25 years with Raymond James, joins Janney as a complex director, according to an announcement.
+
+ Galvin had most recently worked as a divisional director for Raymond James & Associates’ Southern Division. The firm consolidated the territory as part of a reorganization that took effect December 1. Galvin’s registration with Raymond James ended November 8, according to BrokerCheck.
+
+ During his career, Galvin has held a range of branch and complex management roles in the North Atlantic and along the East Coast, according to his LinkedIn profile.
+
+ “We’re looking forward to his experience and strong industry relationships as we continue to expand our team and geographic footprint,” stated Janney’s Florida Regional Director Frank Amigo, who joined from Raymond James in 2017.
+
+ Galvin started his career in 1995 with RBC predecessor firm J. B. Hanauer & Co. and joined Raymond James two years later, according to BrokerCheck. He did not immediately respond to a request for comment sent through social media.""",
+ f"""Don’t Count on a March Rate Cut - Raise Rates
+ Inflation will be stickier than expected, delaying the start of long-awaited interest rate cuts.
+ Investors expecting a rate cut in March may be disappointed.
+ Six-month core consumer price inflation is likely to increase in the first quarter, prompting the Fed to watch and wait.
+ Unless there is an unexpectedly sharp economic downturn or weakening in the labor market, rate cuts are more likely to begin in June.
+
+ Investors betting that the U.S. Federal Reserve will begin trimming interest rates in the first quarter of 2024 may be in for a disappointment.
+
+ After the Fed’s December meeting, market expectations for a March rate cut jumped to surprising heights. Markets are currently putting a 75% chance, approximately, on rate cuts beginning in March. However, Morgan Stanley Research forecasts indicate that cuts are unlikely to come before June.
+
+ Central bank policymakers have likewise pushed back on investors’ expectations. As Federal Reserve Chairman Jerome Powell said in December 2023, when it comes to inflation, “No one is declaring victory. That would be premature.”1
+
+ Here’s why we still expect that rates are likely to hold steady until the middle of 2024.
+
+ Inflation Outlook
+ A renewed uptick in core consumer prices is likely in the first quarter, as prices for services remain elevated, led by healthcare, housing and car insurance. Additionally, in monitoring inflation, the Fed will be watching the six-month average—which means that weaker inflation numbers from summer 2023 will drop out of the comparison window. Although annual inflation rates should continue to decline, the six-month gauge could nudge higher, to 2.4% in January and 2.69% in February.
+
+ Labor markets have also proven resilient, giving Fed policymakers room to watch and wait.
+
+ Data-Driven Expectations
+ Data is critical to the Fed’s decisions and Morgan Stanley’s forecasts, and both could change as new information emerges. At the March policy meeting, the Fed will have only data from January and February in hand, which likely won’t provide enough information for the central bank to be ready to announce a rate cut. The Fed is likely to hold rates steady in March unless nonfarm payrolls add fewer than 50,000 jobs in February and core prices gain less than 0.2% month-over-month. However, unexpected swings in employment and consumer prices, or a marked change in financial conditions or labor force participation, could trigger a cut earlier than we anticipate.
+
+ There are scenarios in which the Fed could cut rates before June, including: a pronounced deterioration in credit conditions, signs of a sharp economic downturn, or slower-than-expected job growth coupled with weak inflation. Weaker inflation and payrolls could bolster the chances of a May rate cut especially.
+
+ When trying to assess timing, statements from Fed insiders are good indicators because they tend to communicate premeditated changes in policy well in advance. If the Fed plans to hold rates steady in March, they might emphasize patience, or talk about inflation remaining elevated. If they’re considering a cut, their language will shift, and they may begin to say that a change in policy may be appropriate “in coming meetings,” “in coming months” or even “soon.” But a long heads up is not guaranteed.
+ https://www.morganstanley.com/ideas/fed-rate-cuts-2024
+ """,
+ f"""What Global Turmoil Could Mean for Investors
+ Weighing the investment impacts of global conflict and geopolitical tensions to international trade, oil prices and China equities.
+ Morgan Stanley Research expects cargo shipping to remain robust despite Red Sea disruption.
+ Crude oil shipments and oil prices should see limited negative impact from regional conflict.
+ Long-term trends could bring growth in Japan and India.
+ In a multipolar world, competition for global power is increasingly leading countries to protect their military and economic interests by erecting new barriers to cross-border commerce in key industries such as technology and renewable energy. As geopolitics and national security are to a growing degree driving how goods flow and where big capital investments are made, it’s that much more crucial for investors to know how to pick through a dizzying amount of information and focus on what’s relevant. But it’s hard to do with a seemingly endless series of alerts lighting up your phone.
+
+ In particular, potential ripples from U.S.-China relations as well as U.S. military involvement in the Middle East could be important for investors. Morgan Stanley Research pared back the headlines and market noise to home in on three key takeaways.
+
+ Gauging Red Sea Disruption
+ Commercial cargo ships in the Red Sea handle about 12% of global trade. Attacks on these ships by Houthi militants, and ongoing U.S. military strikes to quell the disruption, have raised concerns that supply chains could see pandemic-type disruption—and a corresponding spike in inflation.
+
+ However, my colleagues and I expect the flow of container ships to remain robust, even if that flow is redirected to avoid the Red Sea, which serves as an outlet for vessels coming out of the Suez Canal. Although there has been a recent 200% surge in freight rates, there have not been fundamental cost increases for shipping. Additionally, there’s currently a surplus of container ships. Lengthy reroutes around the Southern tip of Africa by carriers to avoid the conflict zone may cause delays, but they should have minimal impact to inflation in Europe. The risks to the U.S. retail sector should be similarly manageable.
+
+ Resilience in Oil and the Countries That Produce it
+ The Middle East is responsible for supplying and producing the majority of the world’s oil, so escalating conflict in the region naturally puts pressure on energy supply, as well the economic growth of relevant countries. However, the threat of weaker growth, higher inflation and erosion of support from allies offer these countries an incentive to contain the conflict. As a result, there’s unlikely to be negative impact to the debt of oil-producing countries in the region. Crude oil shipments should also see limited impacts, though oil prices could spike and European oil refiners, in particular, could face pressure if disruption in the Strait of Hormuz, which traffics about a fifth of oil supplies daily, accelerates.
+
+ Opportunities in Asia Emerging in Japan and India
+ China has significant work to do to retool its economic engine away from property, infrastructure and debt, leading Morgan Stanley economists to predict gross-domestic product growth of 4.2% for 2024 (below the government’s 5% target), slowing to 1.7% from 2025 to 2027. As a result, China’s relatively low equity market valuation still faces challenges, including risks such as U.S. policy restricting future investment. But elsewhere in Asia—particularly in standouts Japan and India—positive long-term trends should drive markets higher. These include fiscal rebalancing, increased digitalization and increasing shifts of manufacturing and supply hubs in a multipolar world.
+
+ For a deeper insights and analysis, ask your Morgan Stanley Representative or Financial Advisor for the full report, “Paying Attention to Global Tension.”
+ https://www.morganstanley.com/ideas/geopolitical-risk-2024
+ """,
+ f"""What 'Edge AI' Means for Smartphones
+ As generative artificial intelligence gets embedded in devices, consumers should see brand new features while smartphone manufacturers could see a sales lift.
+ Advances in artificial intelligence are pushing computing from the cloud directly onto consumer devices, such as smartphones, notebooks, wearables, automobiles and drones.
+ This trend is expected to drive smartphone sales during the next two years, reversing a slowdown that began in 2021.
+ Consumers can expect new features, such as touch-free control of their phones, desktop-quality gaming and real-time photo retouching.
+ As the adoption of generative artificial intelligence accelerates, more computing will be done in the hands of end users—literally. Increasingly, AI will be embedded in consumer devices such as smartphones, notebooks, wearables, automobiles and drones, creating new opportunities and challenges for the manufacturers of these devices.
+
+ Generative AI’s phenomenal capabilities are power-intensive. So far, the processing needed to run sophisticated, mainstream generative AI models can only take place in the cloud. While the cloud will remain the foundation of AI infrastructure, more AI applications, functions and services require faster or more secure computing closer to the consumer. “That’s driving the need for AI algorithms that run locally on the devices rather than on a centralized cloud—or what’s known as the AI at the Edge,” says Ed Stanley, Morgan Stanley’s Head of Thematic Research in London.
+
+ By 2025, Edge AI will be responsible for half of all enterprise data created, according to an estimate by technology market researcher Gartner Inc. While there are many hurdles to reaching commercial viability, the opportunity to tap into 30 billion devices could reduce cost, increase personalization, and improve security and privacy. In addition, faster algorithms on the Edge can reduce latency (i.e., the lag in an app’s response time as it communicates with the cloud).
+
+ “If 2023 was the year of generative AI, 2024 could be the year the technology moves to the Edge,” says Stanley. “We think this trend will pick up steam in 2024, and along with it, opportunities for hardware makers and component suppliers that can help put AI directly into consumers' hands.”
+
+ New Smartphones Lead the Charge
+ Smartphones currently on the market rely on traditional processors and cloud-based computing, and the only AI-enabled programs are features like face recognition, voice assist and low-light photography. Device sales have slowed in recent years, and many investors expect that smartphones will follow the trajectory of personal computers, with multi-year downturns as consumers hold onto their devices for longer due to lack of new features, sensitivity to pricing and other factors.
+
+ But thanks in part to Edge AI, Morgan Stanley analysts think the smartphone market is poised for an upswing and predict that shipments, which have slowed since 2021, will rise by 3.9% this year and 4.4% next year.
+
+ “Given the size of the smartphone market and consumers’ familiarity with them, it makes sense that they will lead the way in bringing AI to the Edge,” says Morgan Stanley’s U.S. Hardware analyst Erik Woodring. “This year should bring a rollout of generative AI-enabled operating systems, as well as next-generation devices and voice assistants that could spur a cycle of smartphone upgrades.”
+
+ However, the move to the Edge will require new smartphone capabilities, especially to improve battery life, power consumption, processing speed and memory. Manufacturers with the strongest brands and balance sheets are best positioned to take the lead in the hardware arms race.
+
+ Killer Apps
+ In addition to hardware, AI itself continues to evolve. New generations of AI models are designed be more flexible and adaptable for a wide range of uses, including Edge devices. Other beneficiaries include smartphone memory players, integrated circuit makers and camera parts suppliers that support new AI applications.
+
+ What can you expect from your phone in the next year?
+
+ “Always-sensing cameras” that automatically activate or lock the screen by detecting if the user is looking at it without the need to touch the screen. This feature could also automatically launch applications such as online payment and food ordering by detecting bar codes.
+
+ Gesture controls for when the user is unable to hold their devices, such as while cooking or exercising.
+
+ Desktop-quality gaming experiences that offer ultra-realistic graphics with cinematic detail, all with smoother interactions and blazing-fast response times.
+
+ Professional-level photography in which image processors enhance photos and video in real time by recognizing each element in a frame—faces, hair, glasses, objects—and fine tune each, eliminating the need for retouching later.
+
+ Smarter voice assistance that is more responsive and tuned the user’s voice and speech patterns, and can launch or suggest apps based on auditory clues.
+
+ “With Edge AI becoming part of everyday life, we see significant opportunities ahead as new hardware provides a platform for developers to create ground-breaking generative AI apps, which could trigger a new hardware product cycle that liftsservices sales,” says Woodring.
+
+ For deeper insights and analysis, ask your Morgan Stanley Representative or Financial Advisor for the full reports, “Tech Diffusion: Edge AI—Growing Impetus” (Nov. 7, 2023), “Edging Into a Smartphone Upcycle” (Nov. 9, 2023) and “Edge AI: Product Releases on Track, But Where Are Killer Apps?”
+ https://www.morganstanley.com/ideas/edge-ai-devices-diffusion""",
+ ]
+
+ questions = [
+ "Can you summarize the latest changes to Morgan Stanley market managers?",
+ "When will the fed lower rates?",
+ "What are the top market risks?",
+ "How will AI impact the smartphone market?",
+ ]
+
+ # Calculate the document embeddings and store in weave
+ document_embeddings = docs_to_embeddings(contexts)
+ embeddings_ref = weave.publish(document_embeddings, "MSRAG_Embeddings")
+
+ for i in range(0, len(questions)):
+ # Using Google Embeddings
+ model.predict(questions[i], contexts, True)
+
+ # Not using Google Embeddings
+ model.predict(questions[1], contexts, False)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/gemini-how-to-use-gemini-pro-api-with-wb-weave/gemini_how_to_use_gemini_pro_api_with_wb_weave.py b/marimo/convert/gemini-how-to-use-gemini-pro-api-with-wb-weave/gemini_how_to_use_gemini_pro_api_with_wb_weave.py
new file mode 100644
index 00000000..0c3da0b1
--- /dev/null
+++ b/marimo/convert/gemini-how-to-use-gemini-pro-api-with-wb-weave/gemini_how_to_use_gemini_pro_api_with_wb_weave.py
@@ -0,0 +1,320 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # How to use Gemini Pro API with W&B Weave
+
+ Read [our article](https://wandb.ai/prompt-eng/gemini-weave/reports/How-to-use-Gemini-Pro-API-with-W-B-Weave--Vmlldzo3NzEwNTA1) and follow along in this colab.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installation
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install google-generativeai weave -qqU
+ return
+
+
+@app.cell
+def _():
+ import google.generativeai as genai
+ import weave
+
+ return genai, weave
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Set up your Google API key and log into W&B Weave
+
+ To run the following cell, your API key must be stored it in a Colab Secret named `GOOGLE_API_KEY`. If you don't already have an API key, or you're not sure how to create a Colab Secret, see the [Authentication](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) quickstart for an example.
+ """)
+ return
+
+
+@app.cell
+def _(genai):
+ from google.colab import userdata
+ GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')
+ genai.configure(api_key=GOOGLE_API_KEY)
+ return
+
+
+@app.cell
+def _(weave):
+ weave.init('prompt-eng/gemini-weave')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Generate a summary and track it in Weave
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !wget https://raw.githubusercontent.com/wandb/llm-workshop-fc2024/main/part_2_structured_outputs/longpaper.txt
+ # with open('longpaper.txt', 'r') as file:
+ # long_paper_text = file.read()
+ return
+
+
+@app.cell
+def _(genai):
+ model_info = genai.get_model('models/gemini-1.5-pro-latest')
+ print(model_info.input_token_limit)
+ return
+
+
+@app.cell
+def _(genai, long_paper_text):
+ model = genai.GenerativeModel('models/gemini-1.5-pro-latest')
+ model.count_tokens(long_paper_text)
+ return (model,)
+
+
+@app.cell
+def _(long_paper_text, model, weave):
+ @weave.op()
+ def generate_summary(text):
+ prompt = "Generate a concise summary of below text:\n"
+ response = model.generate_content(prompt + long_paper_text)
+ return {
+ 'summary': response.text
+ }
+
+ return (generate_summary,)
+
+
+@app.cell
+def _(generate_summary, long_paper_text):
+ summary = generate_summary(long_paper_text)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Gemini API JSON Mode
+ """)
+ return
+
+
+@app.cell
+def _(genai):
+ model_1 = genai.GenerativeModel('gemini-1.5-pro-latest', generation_config={'response_mime_type': 'application/json'})
+ return (model_1,)
+
+
+@app.cell
+def _():
+ from pydantic import BaseModel, Field
+
+ class Summary(BaseModel):
+ title: str
+ summary: str = Field(description="plain short text summary without markdown")
+
+ schema = Summary.model_json_schema()
+ schema
+ return (schema,)
+
+
+@app.cell
+def _():
+ import json
+
+ return (json,)
+
+
+@app.cell
+def _(json, model_1, weave):
+ @weave.op()
+ def create_prompt(text, schema):
+ prompt = f'Generate a concise summary of below text using below JSON schema.\nPlease output plain text without markdown and limit it to 200 words.\nText:\n{text}\nJSON schema:\n{schema}\n'
+ return prompt
+
+ @weave.op()
+ def generate_summary_1(text, schema):
+ prompt = create_prompt(text, schema)
+ response = model_1.generate_content(prompt)
+ try:
+ output = json.loads(response.text)
+ except:
+ output = response.text
+ return {'summary': output}
+
+ return (generate_summary_1,)
+
+
+@app.cell
+def _(generate_summary_1, long_paper_text, schema):
+ new_summary = generate_summary_1(long_paper_text, schema)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Evaluation with Weave
+ """)
+ return
+
+
+@app.cell
+def _(genai, json, weave):
+ from pydantic import model_validator
+ import os
+ import time
+
+ os.environ['WEAVE_PARALLELISM'] = '1' # remove parallelism due to our Gemini quota, remove it if not needed
+
+
+ class SummaryModel(weave.Model):
+ model_name: str
+ prompt_template: str
+ json_schema: dict
+ model: genai.GenerativeModel
+
+ @model_validator(mode="before")
+ def create_model(cls, v):
+ model_name = v["model_name"]
+ model = genai.GenerativeModel(model_name,
+ generation_config={"response_mime_type": "application/json"})
+ v["model"] = model
+ return v
+
+ @weave.op()
+ async def predict(self, text: str) -> dict:
+ time.sleep(15) # remove if your Gemini quota allows for it :)
+ prompt = self.prompt_template.format(text=text, schema=self.schema)
+ response = self.model.generate_content(prompt)
+ try:
+ output = json.loads(response.text)
+ return output[0]
+ except:
+ return {'summary': response.text}
+
+ return (SummaryModel,)
+
+
+@app.cell
+def _():
+ prompt_template = """Generate a concise summary of below text using below JSON schema.
+ Please output plain text without markdown and limit it to 200 words.
+ Text:
+ {text}
+ JSON schema:
+ {schema}
+ """
+ return (prompt_template,)
+
+
+@app.cell
+def _(SummaryModel, prompt_template, schema):
+ model_2 = SummaryModel(model_name='gemini-1.5-pro-latest', prompt_template=prompt_template, json_schema=schema)
+ return (model_2,)
+
+
+@app.cell
+async def _(long_paper_text, model_2):
+ await model_2.predict(long_paper_text)
+ return
+
+
+@app.cell
+def _(weave):
+ dataset_uri = "weave:///prompt-eng/gemini-weave/object/long_papers:9N9vkE4XY1SYoXLbvbCtP0YKqyqXErilG4XW8jYmQgE"
+ dataset = weave.ref(dataset_uri).get()
+ return (dataset,)
+
+
+@app.cell
+def _(weave):
+ # Scoring function checking format adherence
+ @weave.op()
+ def check_formatting(model_output: dict) -> dict:
+ # Check if length is smaller than threshold
+ result = False
+ if type(model_output) == list:
+ model_output = model_output[0]
+ if type(model_output) == dict:
+ if 'summary' in model_output.keys():
+ if type(model_output['summary']) == str:
+ result = True
+ return {'formatting': result}
+
+ return (check_formatting,)
+
+
+@app.cell
+def _(weave):
+ # Scoring function checking length of summary
+ @weave.op()
+ def check_conciseness(model_output: dict) -> dict:
+ # Check if length is smaller than threshold
+ result = False
+ if type(model_output) == list:
+ model_output = model_output[0]
+ if type(model_output) == dict:
+ if 'summary' in model_output.keys():
+ summary = model_output['summary']
+ if type(summary) == str:
+ result = len(summary.split()) < 300
+ return {'conciseness': result}
+
+ return (check_conciseness,)
+
+
+@app.cell
+def _(check_conciseness, check_formatting, dataset, weave):
+ evaluation = weave.Evaluation(
+ dataset=dataset, scorers=[check_formatting, check_conciseness],
+ )
+ return (evaluation,)
+
+
+@app.cell
+async def _(evaluation, model_2):
+ await evaluation.evaluate(model_2)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/grouping-feature/grouping_feature.py b/marimo/convert/grouping-feature/grouping_feature.py
new file mode 100644
index 00000000..66c03a2d
--- /dev/null
+++ b/marimo/convert/grouping-feature/grouping_feature.py
@@ -0,0 +1,93 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Weights & Biases Grouping
+ From your script, use grouping to organize individual runs into larger experiments. This is useful for distributed training and cross validation.
+
+ In `wandb.init()`:
+ - **group**: the first level of organization, usually this is your unique experiment name
+ - **job_type**: the second level of grouping, this is often `train`, `eval`, `optimizer`, `rollout` etc.
+
+ **Links**
+ - [Documentation](https://docs.wandb.ai/library/grouping)
+ - [Example project](https://wandb.ai/carey/group-demo?workspace=user-carey)
+ - [Example dedicated group page](https://wandb.ai/carey/group-demo/groups/exp_5?workspace=user-carey)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this example, I'm setting the experiment index up front, and then incrementing it every time I re-run the cell below.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Set experiment index (for demo purposes)
+ experiment_index = 1
+ return (experiment_index,)
+
+
+@app.cell
+def _(experiment_index):
+ # Simulate launching multiple different jobs that log to the same experiment
+ import wandb
+ import math
+ import random
+ for i in range(5):
+ job_type = 'rollout'
+ if i == 2:
+ job_type = 'eval'
+ if i == 3:
+ job_type = 'eval2'
+ if i == 4:
+ job_type = 'optimizer'
+ wandb.init(project='group-demo', group='exp_' + str(experiment_index), job_type=job_type)
+ for j in range(100):
+ acc = 0.1 * (math.log(1 + j + 0.1) + random.random())
+ val_acc = 0.1 * (math.log(1 + j + 2) + random.random() + random.random()) # Set group and job_type to see auto-grouping in the UI
+ if j % 10 == 0:
+ wandb.log({'acc': acc, 'val_acc': val_acc})
+ wandb.finish()
+ # I'm incrementing this so you can re-run this cell and get another experiment
+ # grouped in the W&B UI
+ experiment_index_1 = experiment_index + 1 # Using this to mark a run complete in a notebook context
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/huggingface-custom-progress-callback/huggingface_custom_progress_callback.py b/marimo/convert/huggingface-custom-progress-callback/huggingface_custom_progress_callback.py
new file mode 100644
index 00000000..91bf9065
--- /dev/null
+++ b/marimo/convert/huggingface-custom-progress-callback/huggingface_custom_progress_callback.py
@@ -0,0 +1,414 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Visualize LLM training progress with Wandb Tables
+
+ In this example we will see how to instrument a custom [callback](https://huggingface.co/docs/transformers/main_classes/callback) for the huggingface [Trainer](https://huggingface.co/docs/transformers/main_classes/trainer) to periodically visualize model predictions using Weight & Biases [Tables](https://wandb.ai/site/tables)
+
+ 
+
+
+
+ ## 🤔 Why should I use W&B?
+
+
+
+ - **Unified dashboard**: Central repository for all your model metrics and predictions
+ - **Lightweight**: No code changes required to integrate with Hugging Face
+ - **Accessible**: Free for individuals and academic teams
+ - **Secure**: All projects are private by default
+ - **Trusted**: Used by machine learning teams at OpenAI, Toyota, Lyft and more
+
+ Think of W&B like GitHub for machine learning models— save machine learning experiments to your private, hosted dashboard. Experiment quickly with the confidence that all the versions of your models are saved for you, no matter where you're running your scripts.
+
+ W&B lightweight integrations works with any Python script, and all you need to do is sign up for a free W&B account to start tracking and visualizing your models.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🚀 Install, Import, and Log in
+
+ To get started with this example you will need to Install the Transformers, Weights & Biases, datasets libraries.
+ - [Hugging Face Transformers](https://github.com/huggingface/transformers)
+ - [Weights & Biases](https://docs.wandb.com/)
+ - [Huggingface Datasets](https://github.com/huggingface/datasets)
+
+ **Uncomment the following cell install the libraries.**
+ """)
+ return
+
+
+@app.cell
+def _():
+ # ! pip install -qqq datasets "transformers[torch]" wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🖊️ [Sign up for a free account →](https://app.wandb.ai/login?signup=true)
+
+ ## 🔑 Put in your API key
+
+ Once you've signed up, run the next cell. You'll be prompted to create a new API key at [wandb.ai/settings](https://wandb.ai/settings) if you haven't already. Store your API key securely. It can only be viewed once when created.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train a language model
+
+ In this notebook, we'll see how to train a [🤗 Transformers](https://github.com/huggingface/transformers) model on the Causal language modeling task i.e. the model has to predict the next token in the sentence (so the labels are the same as the inputs shifted to the right). To make sure the model does not cheat, it gets an attention mask that will prevent it to access the tokens after token i when trying to predict the token i+1 in the sentence.
+
+ We will see how to load and preprocess the dataset for the task and train a model on it using the `Trainer` API.
+ We will be building a custom trainer callback to visualize the model predictions using Weights & Biases Tables by priodically logging the predictions to the table.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Prepare the dataset
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ we will use the [Wikitext 2](https://paperswithcode.com/dataset/wikitext-2) dataset as an example for this task.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from datasets import load_dataset
+ datasets = load_dataset('wikitext', 'wikitext-2-raw-v1')
+
+ # look at a sample from the train dataset
+ datasets["train"][10]
+ return (datasets,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For causal language modeling (CLM) task we are going to take all the texts in our dataset and concatenate them after they are tokenized.
+ Then we split them in examples of a certain sequence length. This way the model will receive chunks of contiguous text that may look like:
+ ```
+ part of text 1
+ ```
+ or
+ ```
+ end of text 1 [BOS_TOKEN] beginning of text 2
+ ```
+ depending on whether they span over several of the original texts in the dataset or not. The labels will be the same as the inputs, shifted to the left.
+
+ We will use the [`gpt2`](https://huggingface.co/gpt2) architecture for this example. You can pick any of the checkpoints listed [here](https://huggingface.co/models?filter=causal-lm) instead.
+ """)
+ return
+
+
+@app.cell
+def _():
+ MODEL_NAME ="gpt2"
+
+ from transformers import AutoTokenizer, AutoModelForCausalLM
+
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+ return AutoModelForCausalLM, MODEL_NAME, tokenizer
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can now call the tokenizer on all our texts using the [`map`](https://huggingface.co/docs/datasets/process#map) method from the Datasets library.
+ """)
+ return
+
+
+@app.cell
+def _(datasets, tokenizer):
+ def tokenize_function(examples):
+ return tokenizer(examples["text"])
+
+ tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=["text"])
+ return (tokenized_datasets,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we need to concatenate all our texts together then split the result in small chunks of a certain `block_size`. Here we will use a block_size of `128`.
+ """)
+ return
+
+
+@app.cell
+def _(tokenized_datasets, tokenizer):
+ BLOCK_SIZE = 128
+
+ def group_texts(examples):
+ # Concatenate all texts.
+ concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+ total_length = len(concatenated_examples[list(examples.keys())[0]])
+ # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+ # customize this part to your needs.
+ total_length = (total_length // BLOCK_SIZE) * BLOCK_SIZE
+ # Split by chunks of max_len.
+ result = {
+ k: [t[i : i + BLOCK_SIZE] for i in range(0, total_length, BLOCK_SIZE)]
+ for k, t in concatenated_examples.items()
+ }
+ result["labels"] = result["input_ids"].copy()
+ return result
+
+
+ lm_datasets = tokenized_datasets.map(
+ group_texts,
+ batched=True,
+ batch_size=1000,
+ num_proc=4,
+ )
+
+ # look at a sample from the preprocessed dataset
+ tokenizer.decode(lm_datasets["train"][1]["input_ids"])
+ return (lm_datasets,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Model & TrainingArguments
+
+ Now that we have our dataset prepared we are readt to instantiate the model and the Training Arguments.
+ For simpilicity we will train the model for 3 epochs and log the model predictions and metrics after each epoch.
+ """)
+ return
+
+
+@app.cell
+def _(AutoModelForCausalLM, MODEL_NAME):
+ from transformers import AutoConfig
+ from transformers import Trainer, TrainingArguments
+ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+ training_args = TrainingArguments(f'{MODEL_NAME}-wikitext2', evaluation_strategy='epoch', num_train_epochs=3, learning_rate=2e-05, do_train=True, do_eval=True, weight_decay=0.01, logging_strategy='epoch', fp16=True, dataloader_num_workers=4)
+ return Trainer, model, training_args
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## WandbPredictionProgressCallback
+
+ To periodically visualize the results we will subclass the [`WandbCallback`](https://huggingface.co/docs/transformers/main_classes/callback#transformers.integrations.WandbCallback) from the transformers library. This callback already instrumented to log model metrics, checkpoints and system metrics to Weights & Biases.
+
+ Here we will customize the callback to periodically log model predictions and labels to a `wandb.Table` so that we can visualize the model predictions as the training progresses. To do this, we will also need to pass the trainier and tokenizer to our callback in order to predict over the validation dataset.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from transformers.integrations import WandbCallback
+ import pandas as pd
+ import os
+
+ os.environ["WANDB_LOG_MODEL"] = "checkpoint"
+
+ def decode_predictions(tokenizer, predictions):
+ labels = tokenizer.batch_decode(predictions.label_ids)
+ prediction_text = tokenizer.batch_decode(predictions.predictions.argmax(axis=-1))
+ return {"labels": labels, "predictions": prediction_text}
+
+
+ class WandbPredictionProgressCallback(WandbCallback):
+ """Custom WandbCallback to log model predictions during training.
+
+ This callback logs model predictions and labels to a wandb.Table at each logging step during training.
+ It allows to visualize the model predictions as the training progresses.
+
+ Attributes:
+ trainer (Trainer): The Hugging Face Trainer instance.
+ tokenizer (AutoTokenizer): The tokenizer associated with the model.
+ sample_dataset (Dataset): A subset of the validation dataset for generating predictions.
+ num_samples (int, optional): Number of samples to select from the validation dataset for generating predictions. Defaults to 100.
+ """
+
+ def __init__(self, trainer, tokenizer, val_dataset, num_samples=100, freq=2):
+ """Initializes the WandbPredictionProgressCallback instance.
+
+ Args:
+ trainer (Trainer): The Hugging Face Trainer instance.
+ tokenizer (AutoTokenizer): The tokenizer associated with the model.
+ val_dataset (Dataset): The validation dataset.
+ num_samples (int, optional): Number of samples to select from the validation dataset for generating predictions. Defaults to 100.
+ freq (int, optional): Control the frequency of logging. Defaults to 2.
+ """
+ super().__init__()
+ self.trainer = trainer
+ self.tokenizer = tokenizer
+ self.sample_dataset = val_dataset.select(range(num_samples))
+ self.freq = freq
+
+
+ def on_evaluate(self, args, state, control, **kwargs):
+ super().on_evaluate(args, state, control, **kwargs)
+ # control the frequency of logging by logging the predictions every `freq` epochs
+ if state.epoch % self.freq == 0:
+ # generate predictions
+ predictions = self.trainer.predict(self.sample_dataset)
+ # decode predictions and labels
+ predictions = decode_predictions(self.tokenizer, predictions)
+ # add predictions to a wandb.Table
+ predictions_df = pd.DataFrame(predictions)
+ predictions_df["epoch"] = state.epoch
+ records_table = self._wandb.Table(dataframe=predictions_df)
+ # log the table to wandb
+ self._wandb.log({"sample_predictions": records_table})
+
+ return (WandbPredictionProgressCallback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Trainer
+
+ First we inistantiate the [`Trainer`] class with the model, training arguments, and the train, eval datasets. Since our callback needs to store predictions after each evaluation we will be passing the trainer to callback and then adding the callback to the trainer.
+
+ **Note**: Here we donot add the callback while inistatiating the `Trainer` but instead we will use the `add_callback` method to include the callback in the trainer after instantiation.
+ """)
+ return
+
+
+@app.cell
+def _(
+ Trainer,
+ WandbPredictionProgressCallback,
+ lm_datasets,
+ model,
+ tokenizer,
+ training_args,
+):
+ trainer = Trainer(
+ model=model,
+ args=training_args,
+ train_dataset=lm_datasets["train"],
+ eval_dataset=lm_datasets["validation"],
+
+ )
+
+
+ progress_callback = WandbPredictionProgressCallback(trainer, tokenizer, lm_datasets["validation"], 10)
+ trainer.add_callback(progress_callback)
+ return (trainer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training
+
+ And that's it, we are ready to train the model and visualize the predictions.
+ """)
+ return
+
+
+@app.cell
+def _(trainer):
+ trainer.train()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👀 Visualize results in dashboard
+
+ Click the link printed out 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:
+
+ ```
+ Tracking run with wandb version
+ Run data is saved locally in
+ Syncing run to Weights & Biases (docs)
+ View project at
+ View run at
+ ```
+
+ Click on the to visualize the sample model predictions epoch. You should see a table similar to the one shown in the screenshot below.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🤓 Learn more!
+ - [Documentation](https://docs.wandb.com/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
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/huggingface-huggingface-wandb/huggingface_huggingface_wandb.py b/marimo/convert/huggingface-huggingface-wandb/huggingface_huggingface_wandb.py
new file mode 100644
index 00000000..d0633922
--- /dev/null
+++ b/marimo/convert/huggingface-huggingface-wandb/huggingface_huggingface_wandb.py
@@ -0,0 +1,227 @@
+# /// script
+# dependencies = ["accelerate", "datasets", "evaluate", "transformers @ git+https://github.com/huggingface/transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Hugging Face + W&B
+ Visualize your [Hugging Face](https://github.com/huggingface/transformers) model's performance quickly with a seamless [W&B](https://wandb.ai/site) integration.
+
+ Compare hyperparameters, output metrics, and system stats like GPU utilization across your models.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🤔 Why should I use W&B?
+
+
+
+ - **Unified dashboard**: Central repository for all your model metrics and predictions
+ - **Lightweight**: No code changes required to integrate with Hugging Face
+ - **Accessible**: Free for individuals and academic teams
+ - **Secure**: All projects are private by default
+ - **Trusted**: Used by machine learning teams at OpenAI, Toyota, Lyft and more
+
+ Think of W&B like GitHub for machine learning models— save machine learning experiments to your private, hosted dashboard. Experiment quickly with the confidence that all the versions of your models are saved for you, no matter where you're running your scripts.
+
+ W&B lightweight integrations works with any Python script, and all you need to do is sign up for a free W&B account to start tracking and visualizing your models.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+
+ 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
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🚀 Install, Import, and Log in
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Install the Hugging Face and Weights & Biases libraries, and the GLUE dataset and training script for this tutorial.
+ - [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
+
+
+@app.cell
+def _(subprocess):
+ # packages added via marimo's package management: datasets wandb evaluate accelerate !pip install datasets wandb evaluate accelerate -qU
+ #! wget https://raw.githubusercontent.com/huggingface/transformers/master/examples/pytorch/text-classification/run_glue.py
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/huggingface/transformers/master/examples/pytorch/text-classification/run_glue.py'])
+ return
+
+
+@app.cell
+def _():
+ # the run_glue.py script requires transformers dev
+ # packages added via marimo's package management: git+https://github.com/huggingface/transformers !pip install -q git+https://github.com/huggingface/transformers
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🖊️ [Sign up for a free account →](https://app.wandb.ai/login?signup=true)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🔑 Put in your API key
+ Once you've signed up, run the next cell. You'll be prompted to create a new API key at [wandb.ai/settings](https://wandb.ai/settings) if you haven't already. Store your API key securely. It can only be viewed once when created.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Optionally, we can set environment variables to customize W&B logging. See [documentation](https://docs.wandb.com/library/integrations/huggingface).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Optional: log both gradients and parameters
+ import os
+ os.environ['WANDB_WATCH'] = 'all'
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(os, subprocess):
+ os.environ['WANDB_PROJECT'] = 'huggingface-demo'
+ os.environ['TASK_NAME'] = 'MRPC'
+ #! python run_glue.py --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 /tmp/$TASK_NAME/ --overwrite_output_dir --logging_steps 50
+ subprocess.call(['python', 'run_glue.py', '--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', '/tmp/$TASK_NAME/', '--overwrite_output_dir', '--logging_steps', '50'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👀 Visualize results in dashboard
+ Click the link printed out 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.
+
+ 
+
+ **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.
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/huggingface-llm-finetuning-notebook/huggingface_llm_finetuning_notebook.py b/marimo/convert/huggingface-llm-finetuning-notebook/huggingface_llm_finetuning_notebook.py
new file mode 100644
index 00000000..19341e31
--- /dev/null
+++ b/marimo/convert/huggingface-llm-finetuning-notebook/huggingface_llm_finetuning_notebook.py
@@ -0,0 +1,320 @@
+# /// script
+# dependencies = ["accelerate", "bitsandbytes", "ctranslate2", "datasets", "loralib", "peft @ git+https://github.com/huggingface/peft.git", "transformers.git@main @ git+https://github.com/huggingface/transformers.git@main", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # LLM Finetuning with HuggingFace and Weights and Biases
+
+ - Fine-tune a lightweight LLM (OPT-125M) with LoRA and 8-bit quantization using Launch
+ - Checkpoint the LoRA adapter weights as artifacts
+ - Link the best checkpoint in Model Registry
+ - Run inference on a quantized model
+
+ The same workflow and principles from this notebook can be applied to fine-tuning some of the stronger OSS LLMs (e.g. Llama2)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Fine-tune large models using 🤗 `peft` adapters, `transformers` & `bitsandbytes`
+
+ In this tutorial we will cover how we can fine-tune large language models using the very recent `peft` library and `bitsandbytes` for loading large models in 8-bit.
+ The fine-tuning method will rely on a recent method called "Low Rank Adapters" (LoRA), instead of fine-tuning the entire model you just have to fine-tune these adapters and load them properly inside the model.
+ After fine-tuning the model you can also share your adapters on the 🤗 Hub and load them very easily. Let's get started!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Install requirements
+
+ First, run the cells below to install the requirements:
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: bitsandbytes datasets accelerate loralib !pip install -q bitsandbytes datasets accelerate loralib
+ # packages added via marimo's package management: git+https://github.com/huggingface/transformers.git@main git+https://github.com/huggingface/peft.git !pip install -q git+https://github.com/huggingface/transformers.git@main git+https://github.com/huggingface/peft.git
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ # packages added via marimo's package management: ctranslate2 !pip install -q ctranslate2
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Model Loading
+
+ - Here we leverage 8-bit quantization to reduce the memory footprint of the model during training
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.environ["CUDA_VISIBLE_DEVICES"]="0"
+ import torch
+ import torch.nn as nn
+ import bitsandbytes as bnb
+ from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
+
+ model = AutoModelForCausalLM.from_pretrained(
+ "facebook/opt-125m",
+ load_in_8bit=True,
+ device_map='auto',
+ )
+
+ tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+ return AutoModelForCausalLM, AutoTokenizer, model, nn, os, tokenizer, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Post-processing on the model
+
+ Finally, we need to apply some post-processing on the 8-bit model to enable training, let's freeze all our layers, and cast the layer-norm in `float32` for stability. We also cast the output of the last layer in `float32` for the same reasons.
+ """)
+ return
+
+
+@app.cell
+def _(model, nn, torch):
+ for param in model.parameters():
+ param.requires_grad = False # freeze the model - train adapters later
+ if param.ndim == 1:
+ # cast the small parameters (e.g. layernorm) to fp32 for stability
+ param.data = param.data.to(torch.float32)
+
+ model.gradient_checkpointing_enable() # reduce number of stored activations
+ model.enable_input_require_grads()
+
+ class CastOutputToFloat(nn.Sequential):
+ def forward(self, x): return super().forward(x).to(torch.float32)
+ model.lm_head = CastOutputToFloat(model.lm_head)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Apply LoRA
+
+ Here comes the magic with `peft`! Let's load a `PeftModel` and specify that we are going to use low-rank adapters (LoRA) using `get_peft_model` utility function from `peft`.
+ """)
+ return
+
+
+@app.function
+def print_trainable_parameters(model):
+ """
+ Prints the number of trainable parameters in the model.
+ """
+ trainable_params = 0
+ all_param = 0
+ for _, param in model.named_parameters():
+ all_param = all_param + param.numel()
+ if param.requires_grad:
+ trainable_params = trainable_params + param.numel()
+ print(f'trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}')
+
+
+@app.cell
+def _(model):
+ from peft import LoraConfig, get_peft_model
+ config = LoraConfig(r=16, lora_alpha=32, target_modules=['q_proj', 'v_proj'], lora_dropout=0.05, bias='none', task_type='CAUSAL_LM')
+ model_1 = get_peft_model(model, config)
+ print_trainable_parameters(model_1)
+ return (model_1,)
+
+
+@app.cell
+def _(model_1):
+ model_1
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Training
+ - [W&B HuggingFace integration](https://docs.wandb.ai/guides/integrations/huggingface) automatically tracks important metrics during the course of training
+ - Also track the HF checkpoints as artifacts and register them in the model registry!
+ - Change the number of steps to 200+ for real results!
+ """)
+ return
+
+
+@app.cell
+def _(model_1, os, tokenizer):
+ import transformers
+ from datasets import load_dataset
+ import wandb
+ project_name = 'llm-finetuning'
+ entity = 'wandb' #@param
+ os.environ['WANDB_LOG_MODEL'] = 'checkpoint' #@param
+ wandb.init(project=project_name, entity=entity, job_type='training')
+ data = load_dataset('Abirate/english_quotes')
+ data = data.map(lambda samples: tokenizer(samples['quote']), batched=True)
+ trainer = transformers.Trainer(model=model_1, train_dataset=data['train'], args=transformers.TrainingArguments(per_device_train_batch_size=4, gradient_accumulation_steps=4, report_to='wandb', warmup_steps=5, max_steps=25, learning_rate=0.0002, fp16=True, logging_steps=1, save_steps=5, output_dir='outputs'), data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False))
+ model_1.config.use_cache = False
+ trainer.train()
+ wandb.finish() # silence the warnings. Please re-enable for inference!
+ return entity, project_name, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Adding Model Weights to W&B Model Registry
+ - Here we get our best checkpoint from the finetuning run and register it as our best model
+ """)
+ return
+
+
+@app.cell
+def _(entity, project_name, wandb):
+ last_run_id = 'zz0lxkc8' #@param
+ wandb.init(project=project_name, entity=entity, job_type='registering_best_model')
+ _best_model = wandb.use_artifact(f'{entity}/{project_name}/checkpoint-{last_run_id}:latest')
+ registered_model_name = 'OPT-125M-english' #@param {type: "string"}
+ wandb.run.link_artifact(_best_model, f'{entity}/model-registry/{registered_model_name}', aliases=['staging'])
+ wandb.finish()
+ return (registered_model_name,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Consuming Model From Registry and Quantizing using ctranslate2
+ - LLMs are typically too large to run in full-precision on even decent hardware.
+ - You can quantize the model to run it more efficiently with minimal loss in accuracy.
+ - CTranslate2 is a great first pass at quantization but doesn't do "smart" quantization. It just converts all weights to half precision.
+ - Checkout out GPTQ and AutoGPTQ for SOTA quantization at scale
+ """)
+ return
+
+
+@app.cell
+def _(entity, project_name, registered_model_name, wandb):
+ # Pull model from the registry
+ wandb.init(project=project_name, entity=entity, job_type='ctranslate2')
+ _best_model = wandb.use_artifact(f'{entity}/model-registry/{registered_model_name}:latest')
+ _best_model.download(root=f'model-registry/{registered_model_name}:latest')
+ wandb.finish()
+ return
+
+
+@app.cell
+def _(AutoModelForCausalLM, AutoTokenizer, os, registered_model_name):
+ from peft import PeftModel, PeftConfig
+
+ def convert_qlora2ct2(adapter_path=f'model-registry/{registered_model_name}:latest',
+ full_model_path="opt125m-finetuned",
+ offload_path="opt125m-offload",
+ ct2_path="opt125m-finetuned-ct2",
+ quantization="int8"):
+
+
+ peft_model_id = adapter_path
+ peftconfig = PeftConfig.from_pretrained(peft_model_id)
+
+ model = AutoModelForCausalLM.from_pretrained(
+ "facebook/opt-125m",
+ offload_folder = offload_path,
+ device_map='auto',
+ )
+
+ tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+
+ model = PeftModel.from_pretrained(model, peft_model_id)
+
+ print("Peft model loaded")
+
+ merged_model = model.merge_and_unload()
+
+ merged_model.save_pretrained(full_model_path)
+ tokenizer.save_pretrained(full_model_path)
+
+ if quantization == False:
+ os.system(f"ct2-transformers-converter --model {full_model_path} --output_dir {ct2_path} --force")
+ else:
+ os.system(f"ct2-transformers-converter --model {full_model_path} --output_dir {ct2_path} --quantization {quantization} --force")
+ print("Convert successfully")
+
+ return (convert_qlora2ct2,)
+
+
+@app.cell
+def _(convert_qlora2ct2, registered_model_name):
+ convert_qlora2ct2(adapter_path=f'model-registry/{registered_model_name}:latest')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Run Inference Using Quantized CTranslate2 Model
+ - Record the results in a W&B Table!
+ """)
+ return
+
+
+@app.cell
+def _(entity, project_name, tokenizer, wandb):
+ import ctranslate2
+
+
+ run = wandb.init(project=project_name, entity=entity, job_type="inference")
+ generator = ctranslate2.Generator("opt125m-finetuned-ct2")
+
+ prompts = ["Hey, are you conscious? Can you talk to me?",
+ "What is machine learning?",
+ "What is W&B?"]
+
+
+ wandb_table = wandb.Table(columns=['prompt', 'completion'])
+ for prompt in prompts:
+ start_tokens = tokenizer.convert_ids_to_tokens(tokenizer.encode(prompt))
+ results = generator.generate_batch([start_tokens], max_length=30)
+ output = tokenizer.decode(results[0].sequences_ids[0])
+ wandb_table.add_data(prompt, output)
+
+ wandb.log({"inference_table": wandb_table})
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/huggingface-optimize-hugging-face-models-with-weights-biases/huggingface_optimize_hugging_face_models_with_weights_biases.py b/marimo/convert/huggingface-optimize-hugging-face-models-with-weights-biases/huggingface_optimize_hugging_face_models_with_weights_biases.py
new file mode 100644
index 00000000..2e28a33c
--- /dev/null
+++ b/marimo/convert/huggingface-optimize-hugging-face-models-with-weights-biases/huggingface_optimize_hugging_face_models_with_weights_biases.py
@@ -0,0 +1,737 @@
+# /// script
+# dependencies = ["accelerate", "datasets", "evaluate", "transformers @ git+https://github.com/huggingface/transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Optimize 🤗 Hugging Face models with Weights & Biases
+ """)
+ return
+
+
+@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.
+
+ Coupled with [Weights & Biases integration](https://docs.wandb.ai/integrations/huggingface), you can quickly train and monitor models for full traceability and reproducibility without any extra line of code! You just need to install the library, sign in, and your experiments will automatically be logged:
+
+ ```bash
+ pip install wandb
+ wandb login
+ ```
+
+ **Note**: To enable logging to W&B, set `report_to` to `wandb` in your `TrainingArguments` or script.
+ """)
+ return
+
+
+@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)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Here's what the W&B interactive dashboard will look like:
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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/python/latest/) 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
+
+ We also install `wandb` to automatically instrument our training.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # packages added via marimo's package management: datasets wandb evaluate accelerate !pip install datasets wandb evaluate accelerate -qU
+ #! wget https://raw.githubusercontent.com/huggingface/transformers/master/examples/pytorch/text-classification/run_glue.py
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/huggingface/transformers/master/examples/pytorch/text-classification/run_glue.py'])
+ return
+
+
+@app.cell
+def _():
+ # the run_glue.py script requires transformers dev
+ # packages added via marimo's package management: git+https://github.com/huggingface/transformers !pip install -q git+https://github.com/huggingface/transformers
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We finally make sure we're logged into W&B so that our experiments can be associated to our account.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@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 models as artifacts: just set environment varilable `WANDB_LOG_MODEL` to `true`
+ * 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/integrations/huggingface).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's log every trained model.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.environ['WANDB_LOG_MODEL'] = 'true'
+ return
+
+
+@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 summary of task](https://huggingface.co/transformers/task_summary.html) 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.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! wget https://raw.githubusercontent.com/huggingface/transformers/master/examples/pytorch/text-classification/run_glue.py
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/huggingface/transformers/master/examples/pytorch/text-classification/run_glue.py'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ These scripts are automatically instrumented with logging when `wandb` is installed and logged in.
+
+ Just set `report_to` to `wandb` to enable logging through W&B.
+
+ **Note**: This cell can take up to 5 minutes to run.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python run_glue.py --report_to wandb --model_name_or_path bert-base-uncased --task_name MRPC --learning_rate 1e-4 --do_train --do_eval --max_steps 300 --logging_steps 30 --evaluation_strategy steps --output_dir /tmp/MRPC --overwrite_output_dir --run_name demo
+ subprocess.call(['python', 'run_glue.py', '--report_to', 'wandb', '--model_name_or_path', 'bert-base-uncased', '--task_name', 'MRPC', '--learning_rate', '1e-4', '--do_train', '--do_eval', '--max_steps', '300', '--logging_steps', '30', '--evaluation_strategy', 'steps', '--output_dir', '/tmp/MRPC', '--overwrite_output_dir', '--run_name', 'demo'])
+ 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!
+
+ Just click the [run page](https://docs.wandb.com/ref/app/pages/run-page)
+ link that appears in the `wandb` section output of the cell above,
+ just before training launches
+ and just after ir finishes.
+
+ It should look something like this:
+ """)
+ 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.com/ref/app/pages/run-page#artifacts-tab) 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):
+ mo.md(r"""
+ # 🔬 Advanced usage & custom training
+
+ Let's create our own logic for a more customized training.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### ✏️ Preparing a dataset
+
+ The dataset will vary based on the task you work on. Let's work on sequence classification!
+
+ 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:
+
+ 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
+
+ The objective of our trained model will be to correctly identify the class associated to new sentences.
+ """)
+ 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/yahoo_answers_topics) and visualize it with the [Datasets viewer](https://huggingface.co/datasets/viewer/?dataset=yahoo_answers_topics).
+
+ 
+
+ Each topic number reprersent 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(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🏷️ Loading a dataset
+
+ Any dataset from the Datasets Hub can easily be loaded and is automatically downloaded if not present locally.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from datasets import load_dataset
+
+ dataset = load_dataset("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 _(dataset):
+ dataset
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can easily access any element.
+ """)
+ return
+
+
+@app.cell
+def _(dataset):
+ 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 _(dataset):
+ 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
+
+
+@app.cell
+def _(dataset):
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(label_list):
+ num_labels = len(label_list)
+ num_labels
+ return (num_labels,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The "topic" class needs to be renamed to "labels" for the `Trainer` to find it.
+ """)
+ return
+
+
+@app.cell
+def _(dataset):
+ dataset_1 = dataset.rename_column('topic', 'labels')
+ return (dataset_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### ⚙️ Tokenizing the dataset
+
+ 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
+
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from transformers import AutoTokenizer
+ tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
+ return (tokenizer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The tokenizer let us quickly preprocess our data.
+ """)
+ return
+
+
+@app.cell
+def _(dataset_1):
+ sample_input = dataset_1['train'][0]['question_title']
+ sample_input
+ return (sample_input,)
+
+
+@app.cell
+def _(sample_input, tokenizer):
+ 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 _(dataset_1, tokenizer):
+ 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 _(dataset_2):
+ dataset_2['train'][0]
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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 _(num_labels):
+ from transformers import AutoModelForSequenceClassification
+ model = AutoModelForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=num_labels)
+ return (model,)
+
+
+@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.
+ """)
+ return
+
+
+@app.cell
+def _(dataset_2, model, tokenizer):
+ import torch
+
+ def get_topic(sentence, tokenize=tokenizer, model=model):
+ inputs = tokenizer(sentence, return_tensors='pt') # tokenize the input
+ inputs = {name: tensor.cuda() for name, tensor in inputs.items()}
+ model = model.cuda() # ensure model and inputs are on the same device (GPU)
+ with torch.no_grad():
+ predictions = model(**inputs)[0].cpu().numpy()
+ top_prediction = predictions.argmax().item() # get prediction - 10 classes "probabilities" (not really true because they still need to be normalized)
+ return dataset_2['train'].features['labels'].int2str(top_prediction) # get the top prediction class and convert it to its associated label
+
+ return (get_topic,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's test a prediction on a sample sentence.
+ """)
+ return
+
+
+@app.cell
+def _(get_topic):
+ 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
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+
+ **Note:** set `report_to` to `wandb` in `TrainingArguments` to enable logging through W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from transformers import Trainer, TrainingArguments
+
+ args = TrainingArguments(
+ report_to = 'wandb', # enable logging to W&B
+ output_dir = 'topic_classification', # output directory
+ overwrite_output_dir = True,
+ evaluation_strategy = 'steps', # check evaluation metrics at each epoch
+ learning_rate = 5e-5, # we can customize learning rate
+ max_steps = 30000,
+ logging_steps = 100, # we will log every 100 steps
+ eval_steps = 5000, # we will perform evaluation every 500 steps
+ save_steps = 10000,
+ load_best_model_at_end = True,
+ metric_for_best_model = 'accuracy',
+ run_name = 'custom_training' # name of the W&B run
+ )
+ return Trainer, args
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For more customization, refer to [`TrainingArguments` documentation](https://huggingface.co/transformers/main_classes/trainer.html#transformers.TrainingArguments).
+ """)
+ return
+
+
+@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.
+
+ Several [metrics](https://huggingface.co/metrics) are readily available from the datasets library to monitor model performance.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from datasets import load_metric
+ import numpy as np
+
+ accuracy_metric = load_metric("accuracy")
+
+ def compute_metrics(eval_pred):
+ predictions, labels = eval_pred
+ predictions = np.argmax(predictions, axis=1)
+ # metrics from the datasets library have a `compute` method
+ return accuracy_metric.compute(predictions=predictions, references=labels)
+
+ return (compute_metrics,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The `Trainer` handles all the training & evaluation logic.
+ """)
+ return
+
+
+@app.cell
+def _(Trainer, args, compute_metrics, dataset_2, model, tokenizer):
+ trainer = Trainer(model=model, args=args, train_dataset=dataset_2['train'], eval_dataset=dataset_2['test'], tokenizer=tokenizer, compute_metrics=compute_metrics) # model to be trained # training args # for padding batched data # for custom metrics
+ return (trainer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can verify that we initially have an accuracy of about 10% (random predictions over 10 classes).
+ """)
+ return
+
+
+@app.cell
+def _(trainer):
+ trainer.evaluate()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We start training by simply calling `train()`.
+ """)
+ return
+
+
+@app.cell
+def _(trainer):
+ trainer.train()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can monitor losses, metrics, gradients and parameters as the model trains.
+
+ 
+
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(get_topic):
+ get_topic('Why is cheese so much better with wine?')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ When we want to close our W&B run, we can call `wandb.finish()` (mainly useful in notebooks, called automatically in scripts).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@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 word](https://huggingface.co/transformers/model_sharing.html) on the Model Hub!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📚 Resources
+
+ * [Hugging Face and W&B integration documentation](https://docs.wandb.ai/integrations/huggingface) contains a few tips for taking most advantage of W&B
+ * [🤗 Transformers documentation](https://huggingface.co/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, please reach out to the [slack community](http://bit.ly/wandb-forum).
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/huggingface-simple-accelerate-integration-wandb/huggingface_simple_accelerate_integration_wandb.py b/marimo/convert/huggingface-simple-accelerate-integration-wandb/huggingface_simple_accelerate_integration_wandb.py
new file mode 100644
index 00000000..860f7698
--- /dev/null
+++ b/marimo/convert/huggingface-simple-accelerate-integration-wandb/huggingface_simple_accelerate_integration_wandb.py
@@ -0,0 +1,284 @@
+# /// script
+# dependencies = ["accelerate", "fastprogress", "timm", "torcheval", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using Huggingface Accelerate with Weights and Biases
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [Accelerate](https://github.com/huggingface/accelerate) is this amazing little framework that simplifies your PyTorch training scripts enabling you to train with all the tricks out there!
+ - Quickly convert your code to support multiple hardward (GPUS, TPUs, Metal,...)
+ - One code to support mixed precision, bfloat16 and even 8 bit Adam.
+
+ Minimal code and no boilerplate. Weights and Biases integration out of the box!
+
+ ```diff
+ import torch
+ import torch.nn.functional as F
+ from datasets import load_dataset
+ + from accelerate import Accelerator
+
+ + accelerator = Accelerator(log_with="wandb")
+ + accelerator.init_trackers("my_wandb_project", config=cfg)
+ - device = 'cpu'
+ + device = accelerator.device
+
+ model = torch.nn.Transformer().to(device)
+ optimizer = torch.optim.Adam(model.parameters())
+
+ dataset = load_dataset('my_dataset')
+ data = torch.utils.data.DataLoader(dataset, shuffle=True)
+
+ + model, optimizer, data = accelerator.prepare(model, optimizer, data)
+
+ model.train()
+ for epoch in range(10):
+ for source, targets in data:
+ source = source.to(device)
+ targets = targets.to(device)
+
+ optimizer.zero_grad()
+
+ output = model(source)
+ loss = F.cross_entropy(output, targets)
+
+ - loss.backward()
+ + accelerator.backward(loss)
+
+ optimizer.step()
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training and Image Classifier
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: accelerate wandb torcheval timm fastprogress !pip install accelerate wandb torcheval timm fastprogress
+ return
+
+
+@app.cell
+def _():
+ import os
+ from types import SimpleNamespace
+
+ import wandb
+
+ import torch
+ import torch.nn as nn
+ import torch.nn.functional as F
+ from torch.utils.data import DataLoader
+ from torch.optim import AdamW
+ from torchvision.datasets import FashionMNIST
+ import torchvision.transforms as T
+ from torcheval.metrics.toolkit import sync_and_compute
+ from fastprogress import progress_bar
+
+ from accelerate import Accelerator
+
+ return (
+ Accelerator,
+ AdamW,
+ DataLoader,
+ F,
+ FashionMNIST,
+ SimpleNamespace,
+ T,
+ nn,
+ progress_bar,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Store your configuration parameters
+ """)
+ return
+
+
+@app.cell
+def _(SimpleNamespace):
+ cfg = SimpleNamespace(
+ path=".",
+ bs=256,
+ epochs=5,
+ size=28,
+ num_workers=8,
+ )
+
+ WANDB_PROJECT = "accelerate_fmnist"
+ return WANDB_PROJECT, cfg
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ setup transforms
+ """)
+ return
+
+
+@app.cell
+def _(T, cfg):
+ tfms = T.Compose([
+ T.RandomCrop(cfg.size, padding=1),
+ T.RandomHorizontalFlip(),
+ T.ToTensor()
+ ])
+ return (tfms,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Create a simple CNN
+ """)
+ return
+
+
+@app.cell
+def _(nn):
+ def conv_block(in_ch, out_ch, ks=3): return nn.Sequential(nn.BatchNorm2d(in_ch),
+ nn.Conv2d(in_ch, out_ch, ks, stride=2, padding=0),
+ nn.ReLU())
+
+ def create_cnn():
+ return nn.Sequential(nn.Conv2d(1, 16, 5, stride=1, padding="same"),
+ conv_block(16, 32),
+ conv_block(32, 64),
+ conv_block(64, 128),
+ conv_block(128, 256, 1),
+ nn.Sequential(nn.Flatten(), nn.Linear(256,10), nn.BatchNorm1d(10)),
+ )
+
+ return (create_cnn,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Wrap everything into a training functions (this is necessary to run on multiple GPUS, if it is only one, you can skip the wrapping)
+ """)
+ return
+
+
+@app.cell
+def _(
+ Accelerator,
+ AdamW,
+ DataLoader,
+ F,
+ FashionMNIST,
+ WANDB_PROJECT,
+ create_cnn,
+ progress_bar,
+ tfms,
+):
+ def train(cfg):
+
+ # data
+ ds = FashionMNIST(cfg.path, transform=tfms, download=True)
+ dl = DataLoader(ds, batch_size=cfg.bs, num_workers=cfg.num_workers)
+
+ # model
+ model = create_cnn()
+
+ # training setup
+ optimizer = AdamW(model.parameters(), lr=1e-3)
+
+
+ # accelerate
+ accelerator = Accelerator(log_with="wandb")
+
+ # this will call wandb.init(...)
+ accelerator.init_trackers(WANDB_PROJECT, config=cfg)
+
+ # prepare
+ model, optimizer, dl = accelerator.prepare(model, optimizer, dl)
+
+ # train
+ model.train()
+ for epoch in progress_bar(range(cfg.epochs)):
+ accurate, num_elems = 0., 0
+ for source, targets in dl:
+ optimizer.zero_grad()
+ output = model(source)
+ loss = F.cross_entropy(output, targets)
+ accelerator.backward(loss)
+
+ # under the hood this calls wandb.log(...) on the main process
+ accelerator.log({"train_loss": loss})
+
+ accurate_preds = output.argmax(dim=1) == targets
+ num_elems += accurate_preds.shape[0]
+ accurate += accurate_preds.long().sum()
+ optimizer.step()
+ accuracy = accurate.item() / num_elems
+ accelerator.log({"epoch":epoch, "accuracy":accuracy}, log_kwargs={"wandb": {"commit": False}})
+ print(f"epoch: {epoch:3} || loss: {loss:5.3f} || accuracy: {accuracy:5.3f}")
+
+ # this will call wandb.finish()
+ accelerator.end_training()
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's train on 2 GPUs! This is really nice, as accelerate will take care of only calling `log` on the main process, so only one run get's created, so no need to manually check the rank of the process when using multiple GPUs.
+ """)
+ return
+
+
+@app.cell
+def _(cfg, train):
+ num_GPUSs = 2
+
+ from accelerate import notebook_launcher
+
+ notebook_launcher(train, (cfg,), num_processes=num_GPUSs)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/huggingface-wandb-hf-example/huggingface_wandb_hf_example.py b/marimo/convert/huggingface-wandb-hf-example/huggingface_wandb_hf_example.py
new file mode 100644
index 00000000..83bd1900
--- /dev/null
+++ b/marimo/convert/huggingface-wandb-hf-example/huggingface_wandb_hf_example.py
@@ -0,0 +1,573 @@
+# /// script
+# dependencies = ["//", "accelerate", "datasets", "evaluate", "github-com/huggingface/transformers", "package @ git+https:", "qqq", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏃♀️ Introduction
+ [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.
+
+ ## 🤔 Why should I use W&B?
+
+
+
+ - **Unified dashboard**: Central repository for all your model metrics and predictions
+ - **Lightweight**: No code changes required to integrate with Hugging Face
+ - **Accessible**: Free for individuals and academic teams
+ - **Secure**: All projects are private by default
+ - **Trusted**: Used by machine learning teams at OpenAI, Toyota, Lyft and more
+
+ Think of Weights & Biases like GitHub for machine learning models — save machine learning experiments to your private, hosted dashboard. Experiment quickly with the confidence that all the versions of your models are saved for you, no matter where you're running your scripts.
+
+ W&B lightweight integrations works with any Python script, and all you need to do is sign up for a free W&B account to start tracking and visualizing your models.
+
+ In the HuggingFace Transformers repo, we've instrumented the Trainer to automatically log training and evaluation metrics to W&B at each logging step.
+
+ 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
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌴 Installation and Setup
+
+ First, let us install the latest version of Weights and Biases. We will then setup a few environment variables to enable Weights & Biases logging and finally authenticate this colab instance to use W&B.
+
+ **Note**: To enable logging to W&B, you will also need to set the `report_to` argument in your `TrainingArguments` or script to `wandb`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install required transformer libraries along with wandb
+ # packages added via marimo's package management: qqq evaluate datasets wandb accelerate git+https: // github.com/huggingface/transformers !pip install - qqq evaluate datasets wandb accelerate git+https: // github.com/huggingface/transformers
+ return
+
+
+@app.cell
+def _():
+ # Setup enviroment variables to enable logging to Weights & Biases
+
+ import os
+ # can be "end", "checkpoint" or "false"
+ os.environ['WANDB_LOG_MODEL'] = "checkpoint"
+ # the name of the wandb project defaults to `huggingface`
+ os.environ['WANDB_PROJECT'] = "hf_transformers"
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🖊️ Sign-up/login
+ If this is your first time using Weights & Biases or you are not logged in, the link that appears after running `wandb.login()` in the following code cell will take you to sign-up/login page. Signing up for a [free account](https://wandb.ai/signup) is as easy as a few clicks.
+
+ ## 🔑 Authentication
+ Once you've signed up, run the next cell. You'll be prompted to create a new API key at [wandb.ai/settings](https://wandb.ai/settings) if you haven't already. Store your API key securely. It can only be viewed once when created.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Login and authenticate Weights & Biases
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Task
+
+ Text classification is a common NLP task that assigns a label or class to text. Some of the largest companies run text classification in production for a wide range of practical applications. In this example we will use the [TweetEval](https://arxiv.org/abs/2010.12421) dataset to classify tweets into identify the emotions evoked by a tweet. The dataset is used as a benchmark to train models for tweet classification tasks. We will use then use a distilled verison of RoBERTa model - [distilroberta-base](https://huggingface.co/distilroberta-base) to recoganize the emotions evoked by the tweets.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Data
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Loading the data
+ Start by loading the tweet_eval dataset from the 🤗 Datasets library:
+ """)
+ return
+
+
+@app.cell
+def _():
+ from datasets import load_dataset
+
+ dataset = load_dataset("tweet_eval", "emotion")
+ return (dataset,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Understanding the dataset
+ """)
+ return
+
+
+@app.cell
+def _(dataset):
+ # What does the dataset look like ?
+ print(dataset)
+
+ # look at an example record
+ print("\nSample Record:", end="\t")
+ print(dataset["validation"][0])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ There are two fields in this dataset:
+
+ - `text`: The text of the tweet.
+ - `label`: The integer label of the emotion corresponding to the tweet
+ """)
+ return
+
+
+@app.cell
+def _(dataset):
+ # What do the labels mean ?
+ idx2label = dict(enumerate(dataset["train"].features["label"].names))
+ label2idx = {v: k for k, v in idx2label.items()}
+
+ print(idx2label)
+ return idx2label, label2idx
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Preprocessing
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We need to convert the `text` to integer tokens so that they can be passed into the model as inputs. To do this we will use the `distilroberta` tokenizer to preprocess the `text` field in the dataset.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from transformers import AutoTokenizer
+ MODEL_NAME = "distilroberta-base"
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+ return MODEL_NAME, tokenizer
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Create a preprocessing function to tokenize `text` and truncate sequences to be no longer than distilroberta's maximum input length:
+ """)
+ return
+
+
+@app.cell
+def _(tokenizer):
+ def preprocess_function(examples):
+ return tokenizer(examples["text"], truncation=True)
+
+ return (preprocess_function,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To apply the preprocessing function over the entire dataset, use 🤗 Datasets [map](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.map) function. You can speed up `map` by setting `batched=True` to process multiple elements of the dataset at once:
+ """)
+ return
+
+
+@app.cell
+def _(dataset, preprocess_function):
+ tokenized_ds = dataset.map(preprocess_function, batched=True,)
+ tokenized_ds
+ return (tokenized_ds,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The above step added two new columns to our dataset. `input_ids` and `attention_mask`. These are the inputs we will be passing to our model.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Since all our examples are of different lengths and the model expects a batch of tokens with the same length we will need to pad our inputs. We can use the `DataCollatorWithPadding` utility to do this. To further speed up training we will pre-compute the length of texts in the tokenized dataset and sort the dataset by this column. This ensures that the batches of data have as minimal padding as possible.
+ """)
+ return
+
+
+@app.cell
+def _(tokenized_ds):
+ def length_function(examples):
+ return {'length': [len(example) for example in examples['input_ids']]}
+ tokenized_ds_1 = tokenized_ds.map(length_function, batched=True)
+ tokenized_ds_1 = tokenized_ds_1.sort('length')
+ return (tokenized_ds_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now create a batch of examples using [DataCollatorWithPadding](https://huggingface.co/docs/transformers/main/en/main_classes/data_collator#transformers.DataCollatorWithPadding). It's more efficient to *dynamically pad* the sentences to the longest length in a batch during collation, instead of padding the whole dataset to the maximium length.
+ """)
+ return
+
+
+@app.cell
+def _(tokenizer):
+ from transformers import DataCollatorWithPadding
+
+ data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+ return (data_collator,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Evaluation
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Including a metric during training is often helpful for evaluating your model's performance. You can quickly load a evaluation method with the 🤗 [Evaluate](https://huggingface.co/docs/evaluate/index) library. For this task, load the [f1-score](https://huggingface.co/spaces/evaluate-metric/f1) metric. This is the metric used in the TweetEval benchmark.
+ You will notice that this metric get logged automatically to your weights & biases run while training.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import evaluate
+
+ f1_score = evaluate.load("f1")
+ return (f1_score,)
+
+
+@app.cell
+def _(f1_score):
+ import numpy as np
+
+
+ def compute_metrics(eval_pred):
+ predictions, labels = eval_pred
+ predictions = np.argmax(predictions, axis=1)
+ return f1_score.compute(predictions=predictions,
+ references=labels,
+ average="weighted")
+
+ return (compute_metrics,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Your `compute_metrics` function is ready to go now, and you'll return to it when you setup your training.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train
+ """)
+ return
+
+
+@app.cell
+def _(MODEL_NAME, idx2label, label2idx):
+ from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer
+
+ model = AutoModelForSequenceClassification.from_pretrained(
+ MODEL_NAME,
+ num_labels=len(idx2label),
+ id2label=idx2label,
+ label2id=label2idx,
+ attention_probs_dropout_prob=0.2,
+ hidden_dropout_prob=0.3)
+ return Trainer, TrainingArguments, model
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We are almost ready to train our model. The steps that remain include:
+
+ 1. Define your training hyperparameters in [TrainingArguments](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments). The only required parameter is `output_dir` which specifies where to save your model. You'll also add the `report_to="wandb"` argument here. At the end of each epoch, the [Trainer](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.Trainer) will evaluate the accuracy and save the training checkpoint. These metrics and checkpoints are automatically pushed to your wandb project.
+ 2. Pass the training arguments to [Trainer](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.Trainer) along with the model, dataset, tokenizer, data collator, and `compute_metrics` function.
+ 3. Call [train()](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.Trainer.train) to finetune your model.
+ """)
+ return
+
+
+@app.cell
+def _(TrainingArguments):
+ training_args = TrainingArguments(
+ output_dir="my_emotion_model",
+ learning_rate=2e-5,
+ per_device_train_batch_size=128,
+ per_device_eval_batch_size=128,
+ num_train_epochs=5,
+ weight_decay=0.01,
+ evaluation_strategy="epoch",
+ save_strategy="epoch",
+ logging_strategy="steps",
+ logging_steps=25,
+ load_best_model_at_end=True,
+ warmup_steps=50,
+ save_total_limit=2,
+ report_to="wandb", # enable logging metrics and model checkpoints to Weights & Biases
+ )
+ return (training_args,)
+
+
+@app.cell
+def _(
+ Trainer,
+ compute_metrics,
+ data_collator,
+ model,
+ tokenized_ds_1,
+ tokenizer,
+ training_args,
+):
+ trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_ds_1['train'], eval_dataset=tokenized_ds_1['validation'], tokenizer=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics)
+ trainer.train()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can visuzalize the training logs by looking at the wandb.run object or by clicking the link printed out above, or go to wandb.ai to see your results stream in live. The link to see your run in the browser will appear just before the training begins — look for the following output: "wandb: 🚀 View run at [URL to your unique run]"
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.run
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Finally, we can optionally call the `wandb.finish()` method to indicate that the experiment is complete.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Resuming Training
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ But wait!! Looks like the model did not converge. Perhaps we should train for a few more epochs. Additionally, since we are training the model on colab it is possible that the preemptible instance was shutdown midway and that the model was not fully trained. Don't worry the wandb integration got us fully covered. We can easily resume training from the last checkpoint by doing the following.
+
+ 1. Initialize the last wandb run by passing the `run id` from your Weights & Biases workspace to `wandb.init`
+ 2. Download the lastest checkpoint using `wandb.artifact`.
+ 3. Reinitialize the trainer and pass the `artifact_dir` to the `resume_from_checkpoint` argument in the `trainer.train` method.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Note: Change the `last_run_id` in the below cell to the id from your wandb run`**
+ """)
+ return
+
+
+@app.cell
+def _(os, wandb):
+ last_run_id = "25d6hznl" # fetch the run_id from your wandb workspace
+
+ # resume the wandb run from the run_id
+ run = wandb.init(
+ project=os.environ["WANDB_PROJECT"],
+ id=last_run_id,
+ resume="must",
+ )
+ return last_run_id, run
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Note: Change the `latest_checpoint`in the below cell to the checkpoint artifact from your run**
+ """)
+ return
+
+
+@app.cell
+def _(last_run_id, run):
+ # fetch the checkpoint artifact from the run
+ # this is usually in the format "/checkpoint-:
+ latest_checkpoint = f'parambharat/hf_transformers/checkpoint-{last_run_id}:v5'
+ _artifact = run.use_artifact(latest_checkpoint, type='model')
+ artifact_dir = _artifact.download()
+ return (artifact_dir,)
+
+
+@app.cell
+def _(
+ Trainer,
+ TrainingArguments,
+ compute_metrics,
+ data_collator,
+ model,
+ tokenized_ds_1,
+ tokenizer,
+):
+ # recreate the training arguments with more epochs
+ training_args_1 = TrainingArguments(output_dir='my_emotion_model', learning_rate=2e-05, per_device_train_batch_size=128, per_device_eval_batch_size=128, num_train_epochs=12, weight_decay=0.01, evaluation_strategy='epoch', save_strategy='epoch', logging_strategy='steps', logging_steps=25, load_best_model_at_end=True, warmup_steps=50, save_total_limit=2, report_to='wandb')
+ # reinitialize the trainer object
+ trainer_1 = Trainer(model=model, args=training_args_1, train_dataset=tokenized_ds_1['train'], eval_dataset=tokenized_ds_1['validation'], tokenizer=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics) # change the number of epochs to train
+ return (trainer_1,)
+
+
+@app.cell
+def _(artifact_dir, trainer_1):
+ trainer_1.train(resume_from_checkpoint=artifact_dir)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.run
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Inference
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Great, now that you've finetuned a model, you can use it for inference!
+
+ Grab some text you'd like to run inference on:
+ """)
+ return
+
+
+@app.cell
+def _():
+ text = "This was a masterpiece. Not completely faithful to the books, but enthralling from beginning to end. Might be my favorite of the three."
+ return (text,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The simplest way to try out your finetuned model for inference is to use it in a [pipeline()](https://huggingface.co/docs/transformers/main/en/main_classes/pipelines#transformers.pipeline). Instantiate a `pipeline` for sentiment analysis with your model, and pass your text to it. Here we will create a new wandb.run to download the model artifact. Then we simply pass the `artifact_dir` as the pretrained model to the `model` argument in the pipeline.
+ """)
+ return
+
+
+@app.cell
+def _(last_run_id, os, wandb):
+ run_1 = wandb.init(project=os.environ['WANDB_PROJECT'], job_type='inference')
+ latest_model = f'parambharat/hf_transformers/model-{last_run_id}:latest'
+ _artifact = run_1.use_artifact(latest_model, type='model')
+ artifact_dir_1 = _artifact.download()
+ return (artifact_dir_1,)
+
+
+@app.cell
+def _(artifact_dir_1, text):
+ from transformers import pipeline
+ classifier = pipeline('sentiment-analysis', model=artifact_dir_1)
+ predictions = classifier(text)
+ print(predictions)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/intro-3-in-1-intro-to-weights-biases-cv-nlp-and-rl/intro_3_in_1_intro_to_weights_biases_cv_nlp_and_rl.py b/marimo/convert/intro-3-in-1-intro-to-weights-biases-cv-nlp-and-rl/intro_3_in_1_intro_to_weights_biases_cv_nlp_and_rl.py
new file mode 100644
index 00000000..4036922c
--- /dev/null
+++ b/marimo/convert/intro-3-in-1-intro-to-weights-biases-cv-nlp-and-rl/intro_3_in_1_intro_to_weights_biases_cv_nlp_and_rl.py
@@ -0,0 +1,1095 @@
+# /// script
+# dependencies = ["datasets", "fastprogress", "gym", "keras-rl2", "timm", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # A 3-in-1 Intro to Weights & Biases: Computer Vision, Natural Language Processing and Reinforcement Learning
+
+ Weights & Biases is a developer toolkit for machine learning experiment tracking, dataset and model versioning, and collaboration
+
+
+
+ In this mega 3-in-1 notebook you'll see how Weights and Biases seamlessly integrates into ML code across modalities -- one example each from Computer Vision, NLP and Reinforcement Learning -- as well as across frameworks like PyTorch, Keras, and more.
+
+ For the full range of supported integrations, plus more examples, see
+ [our docs](https://docs.wandb.ai/integrations).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 0. ✍️ The `wandb` Library
+
+ When working directly with the `wandb` library,
+ the functions you'll use most often are:
+
+ - `wandb.login` - Login to W&B at the start of your session
+ - [`wandb.init`](https://docs.wandb.ai/guides/track/launch) - Initialise a new W&B, returns a "run" object
+ - [`wandb.log`](https://docs.wandb.ai/guides/track/log) - Add information to the logs for your run
+
+ Most of the code in this notebook is for setting up and executing
+ our ML experiments and is not specific to experiment tracking with W&B.
+ To see where W&B is added in the code below you can search for these functions in the notebook or look for the ✍️ emoji.
+
+ When a W&B run begins, a link labeled **Run Page**
+ will be printed to the standard out.
+ This link will take you to a W&B
+ [dashboard](https://docs.wandb.ai/ref/app/pages/run-page)
+ where you can view live, interactive charts
+ and information about your experiment.
+
+ ## Data & Privacy
+ We take security seriously, and our cloud-hosted dashboard uses industry best practices for encryption. If you're working with datasets that cannot leave your enterprise cluster, we have [on-prem](https://docs.wandb.com/self-hosted) installations available. It's also easy to download all your data and export it to other tools, for example, for custom analysis in a Jupyter notebook. Here's more on our [API](https://docs.wandb.com/library/api).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Creating a W&B Account, Importing Libraries, and Logging in
+ [Create an account](http://wandb.ai/login?signup=true), then run the following code cell to install `wandb` and log in.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb timm fastprogress transformers datasets !pip install wandb timm fastprogress transformers datasets -Uqqq
+ return
+
+
+@app.cell
+def _():
+ import os
+ from os import listdir
+ from os.path import isfile, join
+ import time
+
+ from fastprogress.fastprogress import master_bar, progress_bar
+ import numpy as np
+ import pandas as pd
+ import PIL
+ from torchvision import transforms
+ import torch
+ from torch.utils.data import Dataset, DataLoader
+
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+ from IPython.display import clear_output
+ clear_output()
+ return (
+ DataLoader,
+ Dataset,
+ PIL,
+ clear_output,
+ device,
+ isfile,
+ join,
+ listdir,
+ master_bar,
+ np,
+ os,
+ pd,
+ progress_bar,
+ time,
+ torch,
+ transforms,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ✍️ Login to wandb
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ # Use wandb-core
+ wandb.require("core")
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 1. 👓 Computer Vision - Fine-tune MobileNet with W&B and PyTorch
+
+ First, we'll see an example of image classification
+ on a super-mini version of the
+ [iNaturalist 2021 dataset](https://www.kaggle.com/c/inaturalist-2021)
+ using a [`mobilenet-V3` model](https://pytorch.org/vision/stable/_modules/torchvision/models/mobilenetv3.html)
+ from the
+ [PyTorch IMage Models, or "`timm`" library](https://github.com/rwightman/pytorch-image-models).
+
+ **Credit:** This code based on the official PyTorch tutorial
+ [here](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Download Data from W&B Artifacts
+
+ We use W&B Artifacts to download the super-mini iNaturalist 2021 dataset, a subset of the [iNaturalist 2021 dataset](https://www.kaggle.com/c/inaturalist-2021) with only 50 classes with 50 examples each.
+
+ W&B Artifacts is a versioning system for
+ models, datasets, and other large files.
+ Read more in our [docs here](https://docs.wandb.ai/guides/artifacts).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # ✍️ Start a wandb run and add any additional info you'd like to the config
+ WANDB_CV_PROJECT = 'iNat2021'
+ _run = wandb.init(project=WANDB_CV_PROJECT) # config is optional here
+ artifact = _run.use_artifact('wandb/iNat2021/data_supermini_iNat2021:v1', type='dataset')
+ # ✍️ Identify the W&B Artifact where the dataset is stored and download it
+ data_dir = artifact.download()
+ # ✍️ Close your W&B Run
+ _run.finish()
+ return WANDB_CV_PROJECT, data_dir
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup Data and DataLoaders
+
+ The following cells are boilerplate that
+ defines our dataset and how it should be loaded off disk.
+ No Weights & Biases-specific code here
+ -- you can just execute and move to the next section.
+ """)
+ return
+
+
+@app.cell
+def _(data_dir, pd):
+ labels_csv_pth = "supertiny_data/supertiny_data.csv"
+ data = pd.read_csv(f"{data_dir}/{labels_csv_pth}")
+
+ N_CLASSES = data.category_id.nunique()
+ BS, N_EPOCHS, LR = 128, 10, 0.003
+ return BS, LR, N_CLASSES, N_EPOCHS, labels_csv_pth
+
+
+@app.cell
+def _(
+ Dataset,
+ PIL,
+ data_dir,
+ isfile,
+ join,
+ labels_csv_pth,
+ listdir,
+ pd,
+ torch,
+ transforms,
+):
+ # Define the Dataset
+ class INatDataset(Dataset):
+
+ def __init__(self, data_dir, labels_csv_pth, is_train=True, transform=None):
+ self.data_dir = data_dir
+ self.all_data = pd.read_csv(f'{data_dir}/{labels_csv_pth}')
+ self.cat_ids = self.all_data['category_id'].unique()
+ if is_train:
+ self.data = self.all_data.loc[self.all_data.is_train == 1]
+ else:
+ self.data = self.all_data.loc[self.all_data.is_train == 0]
+ self._setup_files()
+ self._setup_data()
+ self.file_path_col = self.data.columns.get_loc('file_path')
+ self.file_name_col = self.data.columns.get_loc('file_name')
+ self._do_cat2label()
+ self.transform = transform
+
+ def _setup_files(self):
+ self.file_list = [join(self.data_dir, f) for f in listdir(self.data_dir) if isfile(join(self.data_dir, f)) and f.endswith('jpg')]
+ self.file_names = [f_path.split('/')[-1] for f_path in self.file_list]
+ self.files_df = pd.DataFrame({'file_path': self.file_list, 'f_name': self.file_names})
+
+ def _setup_data(self):
+ self.data = pd.merge(self.files_df, self.data, left_on='f_name', right_on='file_name')
+ self.data.drop(columns=['f_name'])
+
+ def _do_cat2label(self):
+ self.cat2label = {}
+ idx = 0
+ for c in self.cat_ids:
+ if c not in self.cat2label:
+ self.cat2label[c] = idx
+ idx = idx + 1
+
+ def __len__(self):
+ return len(self.data)
+
+ def __getitem__(self, idx):
+ f_path = self.data.iloc[idx, self.file_path_col]
+ fn = self.data.iloc[idx, self.file_name_col]
+ cat = int(self.data.loc[self.data.file_name == fn, 'category_id'].values)
+ label = torch.tensor([self.cat2label[cat]])
+ X = PIL.Image.open(f_path)
+ if self.transform is not None:
+ X = self.transform(X)
+ return {'image': X, 'label': label}
+ data_transforms = {'train': transforms.Compose([transforms.RandomResizedCrop(224, scale=(0.5, 1.0)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]), 'val': transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
+ train_dataset = INatDataset(data_dir, labels_csv_pth, is_train=True, transform=data_transforms['train'])
+ val_dataset = INatDataset(data_dir, labels_csv_pth, is_train=False, transform=data_transforms['val'])
+ # Define transforms to apply to our data during training and validation
+ # Construct the datasets
+ dataset_lens = {'train': len(train_dataset), 'val': len(val_dataset)}
+ return dataset_lens, train_dataset, val_dataset
+
+
+@app.cell
+def _(BS, DataLoader, clear_output, train_dataset, val_dataset):
+ # Create DataLoaders
+
+ train_dataloader = DataLoader(
+ train_dataset, batch_size=BS, shuffle=True, num_workers=2)
+ val_dataloader = DataLoader(
+ val_dataset, batch_size=2 * BS, shuffle=False, num_workers=2)
+
+ dataloaders = {"train": train_dataloader, "val": val_dataloader}
+ clear_output()
+ return (dataloaders,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train the Model
+
+ The cell below defines our entire training setup.
+
+ First, we define and download the model,
+ then set up its optimizer.
+
+ Finally, we define the model training logic --
+ and that's where we integrate the `wandb` code
+ for logging metrics and more during training.
+ """)
+ return
+
+
+@app.cell
+def _(
+ LR,
+ N_CLASSES,
+ N_EPOCHS,
+ clear_output,
+ dataloaders,
+ dataset_lens,
+ device,
+ master_bar,
+ progress_bar,
+ time,
+ torch,
+ wandb,
+):
+ import timm
+ from torch.optim import lr_scheduler, AdamW
+ model = timm.create_model('mobilenetv3_large_100', pretrained=True, num_classes=N_CLASSES)
+ # Fetch the model and pretrained weights from timm
+ clear_output()
+ criterion = torch.nn.CrossEntropyLoss()
+ optimizer = AdamW(model.parameters(), lr=LR, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01)
+ scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=LR, steps_per_epoch=len(dataloaders['train']), epochs=N_EPOCHS)
+
+ # Set up the optimizer
+ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
+ since = time.time()
+ model.to(device)
+ mb = master_bar(range(num_epochs))
+ for epoch in mb:
+ print(f'Epoch {epoch}/{num_epochs - 1}')
+ print('-' * 10)
+ for phase in ['train', 'val']:
+ # Define the training loop
+ if phase == 'train':
+ model.train()
+ else:
+ model.eval()
+ running_loss, running_corrects = (0.0, 0)
+ for b_idx, batch in enumerate(progress_bar(dataloaders[phase], parent=mb)):
+ inputs = batch['image'].to(device)
+ labels = batch['label'].to(device)
+ labels = labels.squeeze(1)
+ optimizer.zero_grad()
+ with torch.set_grad_enabled(phase == 'train'): # Set model to training mode
+ outputs = model(inputs) # Set model to evaluate mode
+ _, preds = torch.max(outputs, 1)
+ loss = criterion(outputs, labels)
+ if phase == 'train':
+ log_dict = {f'{phase}/loss': loss} # Iterate over data
+ if b_idx % 50 == 0:
+ log_dict[f'{phase}/train_examples'] = [wandb.Image(i) for i in inputs[:10]]
+ wandb.log(log_dict)
+ if phase == 'train':
+ loss.backward()
+ optimizer.step()
+ running_loss = running_loss + loss.item() * inputs.size(0)
+ running_corrects = running_corrects + torch.sum(preds == labels.data)
+ if phase == 'train': # forward pass
+ scheduler.step()
+ epoch_loss = running_loss / dataset_lens[phase]
+ epoch_acc = running_corrects.double() / dataset_lens[phase]
+ if phase == 'val':
+ wandb.log({f'{phase}/accuracy': epoch_acc, f'{phase}/loss': epoch_loss})
+ print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') # ✍️ Log your loss for this step to wandb
+ print()
+ wandb.run.summary['final_accuracy'] = epoch_acc
+ wandb.run.finish() # ✍️ Occasionally log 10 images from this batch for inspection
+ time_elapsed = time.time() - since
+ print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
+ return model # ✍️ Log to W&B # backward pass + optimize # ✍️ Log validation metrics to W&B # ✍️ Log your final accuracy as a summary metric
+
+ return criterion, model, optimizer, scheduler, train_model
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **✍️ Start a `wandb` Run**
+
+ Here, we inititalize a `wandb` run,
+ log the hyperparameters.
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_CV_PROJECT, wandb):
+ BS_1 = 128
+ N_EPOCHS_1 = 10
+ LR_1 = 0.003
+ _run = wandb.init(project=WANDB_CV_PROJECT, config={'lr': LR_1, 'batch_size': BS_1, 'n_epochs': N_EPOCHS_1})
+ return (N_EPOCHS_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Start Training**
+
+ Training may take several minutes,
+ but you can watch your metrics stream in **live**
+ at the Run Page link produced by the cell above.
+ """)
+ return
+
+
+@app.cell
+def _(N_EPOCHS_1, criterion, model, optimizer, scheduler, train_model):
+ model_1 = train_model(model, criterion, optimizer, scheduler, num_epochs=N_EPOCHS_1)
+ return (model_1,)
+
+
+@app.cell
+def _(model_1, optimizer, torch):
+ # Clear GPU memory in preparation for next model training
+ import gc
+ del model_1
+ del optimizer
+ gc.collect()
+ torch.cuda.empty_cache()
+ return (gc,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Bonus: Object Detection with YOLOv5 and W&B
+ - Interested in object detection? You're in luck! Check out our [YOLOv5 integration here](https://wandb.ai/cayush/yolov5-dsviz-demo/reports/Object-Detection-with-YOLO-and-Weights-Biases--Vmlldzo0NTgzMjk?galleryTag=computer-vision).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 2. 📜 Natural Language Processing - Fine-tune a BERT Model with W&B and Hugging Face
+
+ In this example,
+ we will use [Hugging Face](https://huggingface.co/)
+ to fine-tune a [distilled version of the BERT model](https://huggingface.co/distilbert-base-uncased)
+ for topic classification using the
+ [Yahoo! Answers dataset](https://huggingface.co/datasets/viewer/?dataset=yahoo_answers_topics).
+
+ To learn more about the HuggingFace W&B integration see
+ [our docs](https://docs.wandb.ai/integrations/huggingface).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup `wandb`
+
+ **✍️ `wandb` Environment Variables**
+
+ The easiest way to configure the W&B integration with Hugging Face
+ is by setting the values of certain environment variables.
+
+ Here we will define the project to which
+ we'll log our results and upload our model.
+ See the [docs](https://docs.wandb.ai/integrations/huggingface)
+ for more on how to configure the Hugging Face integration.
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ os.environ["WANDB_LOG_MODEL"] = "true"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup Data and DataLoaders
+
+ The following cell is boilerplate that
+ defines our dataset and how it should be loaded off disk.
+ No Weights & Biases-specific code here
+ -- you can just execute and move to the next section.
+ """)
+ return
+
+
+@app.cell
+def _(clear_output):
+ from datasets import load_dataset
+ from transformers import AutoTokenizer
+
+ # Load the ~1GB Yahoo! Answers dataset
+ dataset = load_dataset("yahoo_answers_topics")
+
+ # Select just a subset of the dataset for this demo
+ dataset["train"] = dataset["train"].select(list(range(40000)))
+ dataset["test"] = dataset["test"].select(list(range(5000)))
+
+ # Extract topics
+ label_list = dataset["train"].flatten_indices().unique("topic")
+ num_labels = len(label_list)
+
+ # Rename topic->label column (hf expects a labels column)
+ dataset = dataset.rename_column("topic", "labels")
+
+ clear_output()
+
+ # Tokenize the text
+
+ tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+ dataset = dataset.map(lambda x: tokenizer(x["question_title"], truncation=True), batched=True)
+ clear_output()
+ return dataset, num_labels, tokenizer
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train the Model
+
+ In the following cell,
+ we use the Hugging Face
+ [`transformers`](https://huggingface.co/transformers/)
+ and [`datasets`](https://huggingface.co/docs/datasets/)
+ libraries to set up our training and logging
+ and to download the pretrained model.
+ """)
+ return
+
+
+@app.cell
+def _(clear_output, np, num_labels):
+ from datasets import load_metric
+ from transformers import Trainer, TrainingArguments
+ from transformers import AutoModelForSequenceClassification
+ args = TrainingArguments(report_to='wandb', output_dir='topic_classification', overwrite_output_dir=True, per_device_train_batch_size=64, per_device_eval_batch_size=128, num_train_epochs=1, learning_rate=0.0001, logging_steps=25, evaluation_strategy='steps', eval_steps=100, load_best_model_at_end=True, save_total_limit=3, metric_for_best_model='accuracy')
+ # Define the TrainingArguments to configure training
+ accuracy_metric = load_metric('accuracy')
+
+ def compute_metrics(eval_pred): # ✍️ enable logging to W&B
+ predictions, labels = eval_pred # output directory
+ predictions = np.argmax(predictions, axis=1)
+ return accuracy_metric.compute(predictions=predictions, references=labels)
+ model_2 = AutoModelForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=num_labels)
+ # Define our evaluation metric and a helper function
+ # Download a pretrained model from Hugging Face
+ clear_output() # do logging every 25 steps # calculate eval metrics based on steps, not epochs # calculate eval metrics every 100 steps # metrics from the datasets library have a `compute` method
+ return Trainer, args, compute_metrics, model_2
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **✍️ Start a `wandb.Run`**
+ """)
+ return
+
+
+@app.cell
+def _(Trainer, args, compute_metrics, dataset, model_2, tokenizer, wandb):
+ WANDB_HF_PROJECT = 'yahoo_answers_topics'
+ _run = wandb.init(project=WANDB_HF_PROJECT, name='yahoo_training')
+ trainer = Trainer(model=model_2, args=args, train_dataset=dataset['train'], eval_dataset=dataset['test'], tokenizer=tokenizer, compute_metrics=compute_metrics)
+ return (trainer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Start Training**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Training may take several minutes,
+ but you can watch your metrics stream in **live**
+ at the 🚀 link produced by the cell above
+ once training starts.
+ """)
+ return
+
+
+@app.cell
+def _(trainer, wandb):
+ trainer.train()
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Test the Predictions**
+ """)
+ return
+
+
+@app.cell
+def _(dataset, model_2, tokenizer, torch):
+ def get_topic(sentence, tokenize=tokenizer, model=model_2):
+ inputs = tokenizer(sentence, return_tensors='pt') # tokenize the input
+ inputs = {name: tensor.cuda() for name, tensor in inputs.items()}
+ model = model.cuda() # ensure model and inputs are on the same device (GPU)
+ with torch.no_grad():
+ predictions = model(**inputs)[0].cpu().numpy()
+ top_prediction = predictions.argmax().item() # get prediction - 10 classes un-normalized probabilities
+ return dataset['train'].features['labels'].int2str(top_prediction) # get the top prediction class and convert it to its associated label
+
+ return (get_topic,)
+
+
+@app.cell
+def _(get_topic):
+ get_topic("Why is cheese so much better with wine?")
+ return
+
+
+@app.cell
+def _(gc, model_2, torch, trainer):
+ # Clear GPU memory in preparation for next model training
+ del model_2
+ del trainer
+ gc.collect()
+ torch.cuda.empty_cache()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 3. 🤖 Reinforcement Learning - Train a Simple Agent with W&B, Keras, and OpenAI Gym
+
+ In this section,
+ we train a [Keras-RL2](https://github.com/wau/keras-rl2)
+ model to solve
+ [Cartpole](https://gsurma.medium.com/cartpole-introduction-to-reinforcement-learning-ed0eb5b58288),
+ a classic introductory RL problem, using
+ [OpenAI Gym](https://gym.openai.com/).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell
+def _():
+ # install dependencies
+ # packages added via marimo's package management: gym[all] !pip install "gym[all]" -Uqqq
+ # install OpenAI RL Gym 🏋♀️
+ # packages added via marimo's package management: keras-rl2 !pip install keras-rl2 -Uqqq
+ # install rl framework on top of Keras 🥕
+ return
+
+
+@app.cell
+def _(clear_output):
+ # import from stdlib -- mostly for videos
+ import base64
+ import glob
+ import io
+ import timeit
+ import warnings
+ from IPython.display import HTML
+ from IPython.display import display
+ # Yes, this is also for videos!
+ import gym
+ import tensorflow as tf
+ from tensorflow import keras
+ # OpenAI Gym 🏋♀️
+ from keras.layers import Dense, Activation, Flatten
+ from tensorflow.keras.optimizers import Adam
+ # Keras, TensorFlow, Numpy
+ import rl
+ import rl.agents
+ import rl.memory
+ import rl.policy
+ # keras-rl2 framework
+ clear_output()
+ return (
+ Activation,
+ Adam,
+ Dense,
+ Flatten,
+ glob,
+ gym,
+ keras,
+ rl,
+ timeit,
+ warnings,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup Data and DataLoaders
+
+ Tools for Rendering OpenAI Gym Videos in Colab
+ """)
+ return
+
+
+@app.cell
+def _(glob, os):
+ # starting a fake screen in the background
+ # in order to render videos
+ os.system("Xvfb :1 -screen 0 1024x768x24 &")
+ os.environ["DISPLAY"] = ":1"
+
+ # utility to get video file from directory
+ def get_video_filename(dir="video"):
+ glob_mp4 = os.path.join(dir, "*.mp4")
+ mp4list = glob.glob(glob_mp4)
+ assert len(mp4list) > 0, "couldnt find video files"
+ return mp4list[-1]
+
+ return (get_video_filename,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ✍️ Integrating W&B with Keras-RL - WandB Test Logger
+ """)
+ return
+
+
+@app.cell
+def _(env, np, rl, timeit, wandb, warnings):
+ class WandbTrainLogger(rl.callbacks.TrainEpisodeLogger):
+ def __init__(self, env, **kwargs):
+ kwargs = {
+ "project": "cartpole",
+ **kwargs
+ }
+ self.wandb_kwargs = kwargs
+ super().__init__()
+
+ def init_logging(self):
+ # ✍️ Initialize your wandb run
+ return wandb.init(**self.wandb_kwargs)
+
+ # at the start of training, we start up wandb
+ def on_train_begin(self, logs):
+ if wandb.run is None:
+ # ✍️ Initialize your wandb run and log configs
+ self.init_logging()
+ wandb.config.update({"env.spec": env.spec.__dict__,
+ })
+ wandb.config.update({
+ "params": self.params,
+ "agent": self.model.__dict__
+ })
+ super().on_train_begin(logs)
+
+ # when an episode finishes, we log its stats
+ def on_episode_end(self, episode, logs):
+ duration = timeit.default_timer() - self.episode_start[episode]
+ episode_steps = len(self.observations[episode])
+ metrics_dict = self.build_metrics_dict(np.array(self.metrics[episode]))
+ # ✍️ Log your metrics to wandb
+ wandb.log({
+ # duration and timing metadata
+ "step": self.step,
+ "episode": episode + 1,
+ "duration": duration,
+ "episode_steps": episode_steps,
+ "sps": float(episode_steps) / duration,
+
+ # reward stats
+ "episode_reward": np.sum(self.rewards[episode]),
+ "reward_mean": np.mean(self.rewards[episode]),
+ "reward_min": np.min(self.rewards[episode]),
+ "reward_max": np.max(self.rewards[episode]),
+
+ # action and observation stats
+ "action_mean": np.mean(self.actions[episode]),
+ "action_min": np.min(self.actions[episode]),
+ "action_max": np.max(self.actions[episode]),
+ "obs_mean": np.mean(self.observations[episode]),
+ "obs_min": np.min(self.observations[episode]),
+ "obs_max": np.max(self.observations[episode]),
+ **metrics_dict
+ })
+
+ # clean up after ourselves
+ del self.episode_start[episode]
+ del self.observations[episode]
+ del self.rewards[episode]
+ del self.actions[episode]
+ del self.metrics[episode]
+
+ # building dict of metric means in a super robust way
+ def build_metrics_dict(self, metrics):
+ metrics_dict = {}
+ with warnings.catch_warnings():
+ warnings.filterwarnings("error")
+ for idx, name in enumerate(self.metrics_names):
+ try:
+ metrics_dict[name] = np.nanmean(metrics[:, idx])
+ except Warning:
+ metrics_dict[name] = float("nan")
+ return metrics_dict
+
+ return (WandbTrainLogger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ WandB Test Logger
+ """)
+ return
+
+
+@app.cell
+def _(env_dict, get_video_filename, np, rl, wandb):
+ class WandbTestLogger(rl.callbacks.TestLogger):
+ def __init__(self, env, **kwargs):
+ kwargs = {
+ "project": "cartpole",
+ **kwargs
+ }
+ self.wandb_kwargs = kwargs
+
+ self.env = env
+ self.rewards = {}
+ self.env.env.theta_threshold_radians = np.pi / 2
+
+ def init_logging(self):
+ # ✍️ Initialize wandb logging
+ return wandb.init(**self.wandb_kwargs)
+
+ # at the start of testing, we start up wandb
+ def on_train_begin(self, logs):
+ # ✍️ Initialize your wandb run and log configs
+ if wandb.run is None:
+ self.init_logging()
+ wandb.config.update(env_dict)
+ wandb.config.update({
+ "params": self.params,
+ "agent": self.model.__dict__,
+ })
+
+ super().on_train_begin(logs)
+
+ def on_episode_end(self, episode, logs):
+ """ Compute and log training statistics of the episode when done """
+ super().on_episode_end(episode, logs)
+
+ wandb.summary.update({
+ "test_reward": logs["episode_reward"]
+ })
+
+ # log gameplay video in wandb
+ self.env.close()
+ mp4 = get_video_filename()
+ wandb.log({"test_gameplay": wandb.Video(mp4, fps=4, format="mp4")})
+
+ return (WandbTestLogger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Setup Gym Environment
+ """)
+ return
+
+
+@app.cell
+def _(gym, np):
+ # First, we build the environment
+ ENV_NAME = "CartPole-v1"
+
+ # Get the environment
+ env = gym.make(ENV_NAME)
+ env.env.theta_threshold_radians = np.pi / 4 # extend the angle allowed before failure
+ return (env,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Setting Hyperparameters
+ """)
+ return
+
+
+@app.cell
+def _():
+ config = {}
+ config["activation"] = "linear" # what is the activation function of our hidden layer?
+ config["n_hidden"] = 1 # how many hidden units does that layer have?
+
+ # how often do we update the target model,
+ # as a fraction of how often we update the online model
+ config["target_model_update"] = 1.0
+ config["lr"] = 1e-10 # learning rate for Adam optimizer
+
+ config["nb_steps"] = 500
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup Model
+ """)
+ return
+
+
+@app.cell
+def _(Activation, Dense, Flatten, keras):
+ def build_model(config, env):
+ model = keras.Sequential()
+ model.add(Flatten(input_shape=(1,) + env.observation_space.shape))
+ model.add(Dense(config["n_hidden"]))
+ model.add(Activation(config["activation"]))
+ model.add(Dense(env.action_space.n, activation="linear"))
+
+ return model
+
+ return (build_model,)
+
+
+@app.cell
+def _(build_model, config, env):
+ model_3 = build_model(config, env)
+ model_3.summary()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Construct an Agent
+ """)
+ return
+
+
+@app.cell
+def _(Adam, rl):
+ # an agent combines a policy and a Q value-predicting network
+ # and optionally a memory for replay
+
+ def build_dqn_agent(model, config):
+ training_policy = rl.policy.EpsGreedyQPolicy()
+ test_policy = rl.policy.GreedyQPolicy()
+
+ memory = rl.memory.SequentialMemory(limit=50000, window_length=1)
+ dqn = rl.agents.dqn.DQNAgent(
+ model=model, nb_actions=model.output_shape[-1], memory=memory,
+ nb_steps_warmup=10,
+ target_model_update=config["target_model_update"],
+ policy=training_policy, test_policy=test_policy)
+
+ dqn.compile(Adam(lr=config["lr"]), metrics=["mae"])
+
+ return dqn
+
+ return (build_dqn_agent,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train
+
+ Define your train function
+ """)
+ return
+
+
+@app.cell
+def _(WandbTestLogger, WandbTrainLogger, build_dqn_agent, build_model, wandb):
+ def train(config, env, verbosity=2):
+ train_logger = WandbTrainLogger(env, project='cartpole', job_type='train', config=config) # set up our training and testing loggers
+ test_logger = WandbTestLogger(env)
+ with train_logger.init_logging() as _run:
+ if verbosity > 1:
+ print(env.spec.__dict__)
+ model = build_model(wandb.config, env) # tell wandb we're ready to go!
+ if verbosity:
+ model.summary() # describe the environment if at high verbosity
+ agent = build_dqn_agent(model, wandb.config) # this gets logged to W&B too!
+ agent.fit(env, nb_steps=wandb.config['nb_steps'], visualize=False, callbacks=[train_logger])
+ return agent # summarize the model's contents if non-zero verbosity # this gets logged to W&B too! # call .fit on the agent # # test the agent once and generate a video # test_env = gym.wrappers.Monitor(env, "./video", force=True) # agent.test(test_env, visualize=False, verbose=0, callbacks=[test_logger])
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Train your agent
+ """)
+ return
+
+
+@app.cell
+def _(config, env, train):
+ dqn = train(config, env)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ✍️ W&B Best Practices
+
+ 1. **Projects**: Log multiple runs to a project to compare them. `wandb.init(project="project-name")`
+ 2. **Groups**: For multiple processes or cross validation folds, log each process as a runs and group them together. `wandb.init(group="experiment-1")`
+ 3. **Tags**: Add tags to track your current baseline or production model.
+ 4. **Notes**: Type notes in the table to track the changes between runs.
+ 5. **Reports**: Take quick notes on progress to share with colleagues and make dashboards and snapshots of your ML projects.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ✍️ W&B Advanced Setup
+ 1. [Environment variables](https://docs.wandb.com/library/environment-variables): Set API keys in environment variables so you can run training on a managed cluster.
+ 2. [Offline mode](https://docs.wandb.com/library/technical-faq#can-i-run-wandb-offline): Use `dryrun` mode to train offline and sync results later.
+ 3. [On-prem](https://docs.wandb.com/self-hosted): 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.com/sweeps): Set up hyperparameter search quickly with our lightweight tool for tuning.
+ 5. [Artifacts](https://docs.wandb.com/artifacts): Track and version models and datasets in a streamlined way that automatically picks up your pipeline steps as you train models.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/intro-intro-to-weights-biases-keras/intro_intro_to_weights_biases_keras.py b/marimo/convert/intro-intro-to-weights-biases-keras/intro_intro_to_weights_biases_keras.py
new file mode 100644
index 00000000..023ff867
--- /dev/null
+++ b/marimo/convert/intro-intro-to-weights-biases-keras/intro_intro_to_weights_biases_keras.py
@@ -0,0 +1,244 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏃♀️ Quickstart
+ Use **[Weights & Biases](https://wandb.ai/site?utm_source=keras_intro_colab&utm_medium=code&utm_campaign=keras_intro)** for machine learning experiment tracking, model checkpointing, and collaboration with your team. See the full Weights & Biases Documentation **[here](https://docs.wandb.ai/guides/integrations/keras)**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🤩 A shared dashboard for your experiments
+
+ With just a few lines of code,
+ you'll get rich, interactive, shareable dashboards [which you can see yourself here](https://wandb.ai/wandb/wandb_example).
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🔒 Data & Privacy
+
+ We take security very seriously, and our cloud-hosted dashboard uses industry standard best practices for encryption. If you're working with datasets that cannot leave your enterprise cluster, we have [on-prem](https://docs.wandb.com/self-hosted) installations available.
+
+ It's also easy to download all your data and export it to other tools — like custom analysis in a Jupyter notebook. Here's [more on our API](https://docs.wandb.com/library/api).
+
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Start by installing the library and logging in to your free account.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ # Log in to your W&B account
+ import wandb
+
+ # Use wandb-core
+ wandb.require("core")
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 👟 Run an experiment
+ 1️⃣. **Start a new run** and pass in hyperparameters to track
+
+ 2️⃣. **Log metrics** from training or evaluation
+
+ 3️⃣. **Visualize results** in the dashboard
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ import random
+ for _run in range(5):
+ # Launch 5 simulated experiments
+ wandb.init(project='basic-intro', config={'learning_rate': 0.02, 'architecture': 'CNN', 'dataset': 'CIFAR-100'})
+ offset = random.random() / 5 # 1️⃣ Start a new run to track this script
+ for ii in range(2, 10):
+ acc = 1 - 2 ** (-ii) - random.random() / ii - offset # Set entity to specify your username or team name
+ loss = 2 ** (-ii) + random.random() / ii + offset # ex: entity="carey",
+ wandb.log({'acc': acc, 'loss': loss}) # Set the project where this run will be logged
+ wandb.finish() # Track hyperparameters and run metadata # This simple block simulates a training loop logging metrics # 2️⃣ Log metrics from your script to W&B # Mark the run as finished
+ return (random,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You have now trained your first model using wandb! 👆 Click on the wandb link above to see your metrics
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🥕 Simple Keras Classifier
+ Run this model to train a simple MNIST classifier, and click on the project page link to see your results stream in live to a W&B project. For a full guide on how to use Weights & Biases with Keras, **[see here](https://docs.wandb.ai/guides/integrations/keras)**
+ """)
+ return
+
+
+@app.cell
+def _(random, wandb):
+ import numpy as np
+ import tensorflow as tf
+ from wandb.integration.keras import WandbMetricsLogger, WandbModelCheckpoint
+ for _run in range(5):
+ wandb.init(project='keras-intro', config={'layer_1': 512, 'activation_1': 'relu', 'dropout': random.uniform(0.01, 0.8), 'layer_2': 10, 'activation_2': 'softmax', 'optimizer': 'sgd', 'loss': 'sparse_categorical_crossentropy', 'metric': 'accuracy', 'epoch': 6, 'batch_size': 256})
+ config = wandb.config
+ mnist = tf.keras.datasets.mnist
+ (x_train, y_train), (x_test, y_test) = mnist.load_data()
+ x_train, x_test = (x_train / 255.0, x_test / 255.0)
+ x_train, y_train = (x_train[::5], y_train[::5])
+ x_test, y_test = (x_test[::20], y_test[::20])
+ labels = [str(digit) for digit in range(np.max(y_train) + 1)]
+ model = tf.keras.models.Sequential([tf.keras.layers.Flatten(input_shape=(28, 28)), tf.keras.layers.Dense(config.layer_1, activation=config.activation_1), tf.keras.layers.Dropout(config.dropout), tf.keras.layers.Dense(config.layer_2, activation=config.activation_2)])
+ model.compile(optimizer=config.optimizer, loss=config.loss, metrics=[config.metric])
+ wandb_callbacks = [WandbMetricsLogger(), WandbModelCheckpoint(filepath='my_model_{epoch:02d}.keras')]
+ model.fit(x=x_train, y=y_train, epochs=config.epoch, batch_size=config.batch_size, validation_data=(x_test, y_test), callbacks=wandb_callbacks)
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You have now trained your first model using wandb! 👆 Click on the wandb link above to see your metrics.
+
+ For a full guide on how to use Weights & Biases with Keras, **[see here](https://docs.wandb.ai/guides/integrations/keras)**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔔 Try W&B Alerts
+
+ **[W&B Alerts](https://docs.wandb.ai/guides/track/alert)** allows you to send alerts, triggered from your Python code, to your Slack or email. There are 2 steps to follow the first time you'd like to send a Slack or email alert, triggered from your code:
+
+ 1) Turn on Alerts in your W&B [User Settings](https://wandb.ai/settings)
+
+ 2) Add `wandb.alert()` to your code:
+
+ ```python
+ wandb.alert(
+ title="Low accuracy",
+ text=f"Accuracy is below the acceptable threshold"
+ )
+ ```
+
+ See the minimal example below to see how to use `wandb.alert`. You can find the full docs for **[W&B Alerts here](https://docs.wandb.ai/guides/track/alert)**
+ """)
+ return
+
+
+@app.cell
+def _(random, wandb):
+ # Start a wandb run
+ wandb.init(project="keras-intro")
+
+ # Simulating a model training loop
+ acc_threshold = 0.3
+ for training_step in range(1000):
+
+ # Generate a random number for accuracy
+ accuracy = round(random.random() + random.random(), 3)
+ print(f"Accuracy is: {accuracy}, {acc_threshold}")
+
+ # 🐝 Log accuracy to wandb
+ wandb.log({"Accuracy": accuracy})
+
+ # 🔔 If the accuracy is below the threshold, fire a W&B Alert and stop the run
+ if accuracy <= acc_threshold:
+ # 🐝 Send the wandb Alert
+ wandb.alert(
+ title="Low Accuracy",
+ text=f"Accuracy {accuracy} at step {training_step} is below the acceptable theshold, {acc_threshold}",
+ )
+ print("Alert triggered")
+ break
+
+ # Mark the run as finished (useful in Jupyter notebooks)
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # What's next 🚀 ?
+ The next tutorial you will learn how to do hyperparameter optimization using W&B Sweeps:
+ ## 👉 [Hyperparameters sweeps using PyTorch](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/pytorch/Organizing_Hyperparameter_Sweeps_in_PyTorch_with_W%26B.ipynb)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/intro-intro-to-weights-biases/intro_intro_to_weights_biases.py b/marimo/convert/intro-intro-to-weights-biases/intro_intro_to_weights_biases.py
new file mode 100644
index 00000000..330d87ab
--- /dev/null
+++ b/marimo/convert/intro-intro-to-weights-biases/intro_intro_to_weights_biases.py
@@ -0,0 +1,364 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use [W&B](https://wandb.ai/site?utm_source=intro_colab&utm_medium=code&utm_campaign=intro) for machine learning experiment tracking, model checkpointing, collaboration with your team and more. See the full W&B Documentation [here](https://docs.wandb.ai/).
+
+ In this notebook, you will create and track a machine learning experiment using a simple PyTorch model. By the end of the notebook, you will have an interactive project dashboard that you can share and customize with other members of your team. [View an example dashboard here](https://wandb.ai/wandb/wandb_example).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Prerequisites
+
+ Install the W&B Python SDK and log in:
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ # Log in to your W&B account
+ import wandb
+ import random
+ import math
+
+ return math, random, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Simulate and track a machine learning experiment with W&B
+
+ Create, track, and visualize a machine learning experiment. To do this:
+
+ 1. Initialize a [W&B run](https://docs.wandb.ai/guides/runs) and pass in the hyperparameters you want to track.
+ 2. Within your training loop, log metrics such as the accuracy and loss.
+ """)
+ return
+
+
+@app.cell
+def _(random, wandb):
+ total_runs = 5
+ for run in range(total_runs):
+ wandb.init(project='basic-intro', name=f'experiment_{run}', config={'learning_rate': 0.02, 'architecture': 'CNN', 'dataset': 'CIFAR-100', 'epochs': 10})
+ epochs = 10
+ offset = random.random() / 5
+ for _epoch in range(2, epochs):
+ acc = 1 - 2 ** (-_epoch) - random.random() / _epoch - offset
+ loss = 2 ** (-_epoch) + random.random() / _epoch + offset
+ wandb.log({'acc': acc, 'loss': loss})
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ View how your machine learning peformed in your W&B project. Copy and paste the URL link that is printed from the previous cell. The URL will redirect you to a W&B project that contains a dashboard showing graphs the show how
+
+ The following image shows what a dashboard can look like:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now that we know how to integrate W&B into a pseudo machine learning training loop, let's track a machine learning experiment using a basic PyTorch neural network. The following code will also upload model checkpoints to W&B that you can then share with other teams in in your organization.
+
+ ## Track a machine learning experiment using PyTorch
+
+ The following code cell defines and trains a simple MNIST classifier. During training, you will see W&B prints out URLs. Click on the project page link to see your results stream in live to a W&B project.
+
+ W&B runs automatically log [metrics](https://docs.wandb.ai/ref/app/pages/run-page#charts-tab),
+ [system information](https://docs.wandb.ai/ref/app/pages/run-page#system-tab),
+ [hyperparameters](https://docs.wandb.ai/ref/app/pages/run-page#overview-tab),
+ [terminal output](https://docs.wandb.ai/ref/app/pages/run-page#logs-tab) and
+ you'll see an [interactive table](https://docs.wandb.ai/guides/data-vis)
+ with model inputs and outputs.
+
+ ### Set up PyTorch Dataloader
+ The following cell defines some useful functions that we will need to train our machine learning model. The functions themselves are not unique to W&B so we'll not cover them in detail here. See the PyTorch documentation for more information on how to define [forward and backward training loop](https://pytorch.org/tutorials/beginner/nn_tutorial.html), how to use [PyTorch DataLoaders](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html) to load data in for training, and how define PyTorch models using the [`torch.nn.Sequential` Class](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html).
+ """)
+ return
+
+
+@app.cell
+def _(log_image_table):
+ #@title
+ import torch, torchvision
+ import torch.nn as nn
+ from torchvision.datasets import MNIST
+ import torchvision.transforms as T
+
+ MNIST.mirrors = [mirror for mirror in MNIST.mirrors if "http://yann.lecun.com/" not in mirror]
+
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+ def get_dataloader(is_train, batch_size, slice=5):
+ "Get a training dataloader"
+ full_dataset = MNIST(root=".", train=is_train, transform=T.ToTensor(), download=True)
+ sub_dataset = torch.utils.data.Subset(full_dataset, indices=range(0, len(full_dataset), slice))
+ loader = torch.utils.data.DataLoader(dataset=sub_dataset,
+ batch_size=batch_size,
+ shuffle=True if is_train else False,
+ pin_memory=True, num_workers=2)
+ return loader
+
+ def get_model(dropout):
+ "A simple model"
+ model = nn.Sequential(nn.Flatten(),
+ nn.Linear(28*28, 256),
+ nn.BatchNorm1d(256),
+ nn.ReLU(),
+ nn.Dropout(dropout),
+ nn.Linear(256,10)).to(device)
+ return model
+
+ def validate_model(model, valid_dl, loss_func, log_images=False, batch_idx=0):
+ "Compute performance of the model on the validation dataset and log a wandb.Table"
+ model.eval()
+ val_loss = 0.
+ with torch.inference_mode():
+ correct = 0
+ for i, (images, labels) in enumerate(valid_dl):
+ images, labels = images.to(device), labels.to(device)
+
+ # Forward pass ➡
+ outputs = model(images)
+ val_loss += loss_func(outputs, labels)*labels.size(0)
+
+ # Compute accuracy and accumulate
+ _, predicted = torch.max(outputs.data, 1)
+ correct += (predicted == labels).sum().item()
+
+ # Log one batch of images to the dashboard, always same batch_idx.
+ if i==batch_idx and log_images:
+ log_image_table(images, predicted, labels, outputs.softmax(dim=1))
+ return val_loss / len(valid_dl.dataset), correct / len(valid_dl.dataset)
+
+ return device, get_dataloader, get_model, nn, torch, validate_model
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Create a table to compare the predicted values versus the true value
+
+ The following cell is unique to W&B, so let's go over it.
+
+ In the cell we define a function called `log_image_table`. Though technically, optional, this function creates a W&B Table object. We will use the table object to create a table that shows what the model predicted for each image.
+
+ More specifically, each row will conists of the image fed to the model, along with predicted value and the actual value (label).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ def log_image_table(images, predicted, labels, probs):
+ "Log a wandb.Table with (img, pred, target, scores)"
+ # Create a wandb Table to log images, labels and predictions to
+ table = wandb.Table(columns=["image", "pred", "target"]+[f"score_{i}" for i in range(10)])
+ for img, pred, targ, prob in zip(images.to("cpu"), predicted.to("cpu"), labels.to("cpu"), probs.to("cpu")):
+ table.add_data(wandb.Image(img[0].numpy()*255), pred, targ, *prob.numpy())
+ wandb.log({"predictions_table":table}, commit=False)
+
+ return (log_image_table,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Train your model and upload checkpoints
+
+ The following code trains and saves model checkpoints to your project. Use model checkpoints like you normally would to assess how the model performed during training.
+
+ W&B also makes it easy to share your saved models and model checkpoints with other members of your team or organization. To learn how to share your model and model checkpoints with members outside of your team, see [W&B Registry](https://docs.wandb.ai/guides/registry).
+ """)
+ return
+
+
+@app.cell
+def _(
+ device,
+ get_dataloader,
+ get_model,
+ math,
+ nn,
+ random,
+ torch,
+ validate_model,
+ wandb,
+):
+ # Launch 3 experiments, trying different dropout rates
+ for _ in range(3):
+ wandb.init(project='pytorch-intro', config={'epochs': 5, 'batch_size': 128, 'lr': 0.001, 'dropout': random.uniform(0.01, 0.8)}) # initialise a wandb run
+ config = wandb.config
+ train_dl = get_dataloader(is_train=True, batch_size=config.batch_size)
+ valid_dl = get_dataloader(is_train=False, batch_size=2 * config.batch_size)
+ n_steps_per_epoch = math.ceil(len(train_dl.dataset) / config.batch_size)
+ model = get_model(config.dropout)
+ loss_func = nn.CrossEntropyLoss()
+ optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
+ example_ct = 0
+ step_ct = 0
+ for _epoch in range(config.epochs): # Copy your config
+ model.train()
+ for step, (images, labels) in enumerate(train_dl):
+ images, labels = (images.to(device), labels.to(device)) # Get the data
+ outputs = model(images)
+ train_loss = loss_func(outputs, labels)
+ optimizer.zero_grad()
+ train_loss.backward()
+ optimizer.step() # A simple MLP model
+ example_ct += len(images)
+ metrics = {'train/train_loss': train_loss, 'train/epoch': (step + 1 + n_steps_per_epoch * _epoch) / n_steps_per_epoch, 'train/example_ct': example_ct}
+ if step + 1 < n_steps_per_epoch: # Make the loss and optimizer
+ wandb.log(metrics)
+ step_ct += 1
+ val_loss, _accuracy = validate_model(model, valid_dl, loss_func, log_images=_epoch == config.epochs - 1)
+ val_metrics = {'val/val_loss': val_loss, 'val/val_accuracy': _accuracy} # Training
+ wandb.log({**metrics, **val_metrics})
+ torch.save(model, 'my_model.pt')
+ wandb.log_model('./my_model.pt', 'my_mnist_model', aliases=[f'epoch-{_epoch + 1}_dropout-{round(wandb.config.dropout, 4)}'])
+ print(f'Epoch: {_epoch + 1}, Train Loss: {train_loss:.3f}, Valid Loss: {val_loss:3f}, Accuracy: {_accuracy:.2f}')
+ wandb.summary['test_accuracy'] = 0.8
+ wandb.finish() # Log train metrics to wandb # Log train and validation metrics to wandb # Save the model checkpoint to wandb # If you had a test set, this is how you could log it as a Summary metric # Close your wandb run
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You have now trained your first model using W&B. Click on one of the links above to see your metrics and see your saved model checkpoints in the Artifacts tab in the W&B App UI
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## (Optional) Set up a W&B Alert
+
+ Create a [W&B Alerts](https://docs.wandb.ai/guides/track/alert) to send alerts to your Slack or email from your Python code.
+
+ There are 2 steps to follow the first time you'd like to send a Slack or email alert, triggered from your code:
+
+ 1) Turn on Alerts in your W&B [User Settings](https://wandb.ai/settings)
+ 2) Add `wandb.alert()` to your code. For example:
+
+ ```python
+ wandb.alert(
+ title="Low accuracy",
+ text=f"Accuracy is below the acceptable threshold"
+ )
+ ```
+
+ The following cell shows a minimal example below to see how to use `wandb.alert`
+ """)
+ return
+
+
+@app.cell
+def _(random, wandb):
+ # Start a wandb run
+ wandb.init(project='pytorch-intro')
+ acc_threshold = 0.3
+ # Simulating a model training loop
+ for training_step in range(1000):
+ _accuracy = round(random.random() + random.random(), 3)
+ print(f'Accuracy is: {_accuracy}, {acc_threshold}')
+ wandb.log({'Accuracy': _accuracy}) # Generate a random number for accuracy
+ if _accuracy <= acc_threshold:
+ wandb.alert(title='Low Accuracy', text=f'Accuracy {_accuracy} at step {training_step} is below the acceptable theshold, {acc_threshold}')
+ print('Alert triggered')
+ break # Log accuracy to wandb
+ # Mark the run as finished (useful in Jupyter notebooks)
+ wandb.finish() # If the accuracy is below the threshold, fire a W&B Alert and stop the run # Send the wandb Alert
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can find the full docs for [W&B Alerts here](https://docs.wandb.ai/guides/track/alert).
+
+ ## Next steps
+ The next tutorial you will learn how to do hyperparameter optimization using W&B Sweeps:
+ [Hyperparameters sweeps using PyTorch](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/pytorch/Organizing_Hyperparameter_Sweeps_in_PyTorch_with_W%26B.ipynb)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/intro-report-api-quickstart/intro_report_api_quickstart.py b/marimo/convert/intro-report-api-quickstart/intro_report_api_quickstart.py
new file mode 100644
index 00000000..3166d397
--- /dev/null
+++ b/marimo/convert/intro-report-api-quickstart/intro_report_api_quickstart.py
@@ -0,0 +1,466 @@
+# /// script
+# dependencies = ["wandb", "wandb-workspaces"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+
+ [wandb-workspaces](https://github.com/wandb/wandb-workspaces) is a Python library for programmatically creating and customizing W&B Workspaces and Reports.
+
+ In this tutorial you will see how to use wandb-workspaces to create and customize W&B Reports.
+ """)
+ return
+
+
+@app.cell
+def _():
+ #install dependencies
+ # packages added via marimo's package management: wandb wandb-workspaces !pip install wandb wandb-workspaces -qqq
+ return
+
+
+@app.cell
+def _():
+ #@title ## Log Runs { run: "auto", display-mode: "form" }
+ #@markdown If this is your first time here, consider running the setup code for a better docs experience!
+ #@markdown If you have run the setup code before, you can uncheck the box below to avoid unnecessary logging.
+ LOG_DUMMY_RUNS = True
+ import requests #@param {type: "boolean"}
+ from PIL import Image
+ from io import BytesIO
+ import wandb
+ import pandas as pd
+ from itertools import product
+ import random
+ import math
+ import string
+ ENTITY = wandb.apis.PublicApi().default_entity
+ PROJECT = 'report-api-quickstart'
+ LINEAGE_PROJECT = 'lineage-example'
+
+ def get_image(url):
+ r = requests.get(url)
+ return Image.open(BytesIO(r.content))
+
+ def log_dummy_data(): #@param {type: "string"}
+ 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'] #@param {type: "string"}
+ 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) as run:
+ for step in range(1000):
+ wandb.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]) 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})
+
+ 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 = []
+
+ def create_artifact(name: str, type: str, content: str):
+ art = wandb.Artifact(name, type)
+ with open('boom.txt', 'w') as f:
+ f.write(content)
+ art.add_file('boom.txt', 'test-name')
+ 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
+
+ def log_dummy_lineage():
+ pipeline = [Step('dataset-generator', 1, 0, 3, 'dataset'), Step('trainer', 4, (1, 2), 3), Step('evaluator', 2, 1, 3), Step('ensemble', 1, 1, 1)]
+ for i, step in enumerate(pipeline):
+ for _ in range(step.runs):
+ with wandb.init(project=LINEAGE_PROJECT, job_type=step.job_type) as run:
+ 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)
+ for j in range(step.outputs_per_run):
+ 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()
+ step.artifacts.append(art)
+ if LOG_DUMMY_RUNS:
+ log_dummy_data()
+ log_dummy_lineage() # use # log output artifacts # name = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6)) # save in pipeline
+ return ENTITY, LINEAGE_PROJECT, PROJECT
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🚀 Quickstart!
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb_workspaces.reports.v2 as wr
+
+ return (wr,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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
+
+
+@app.cell
+def _(PROJECT, wr):
+ report = wr.Report(
+ project=PROJECT,
+ title='Quickstart Report',
+ description="That was easy!"
+ ) # Create
+ report.save() # Save
+ wr.Report.from_url(report.url) # Load
+ return (report,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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
+ """)
+ return
+
+
+@app.cell
+def _(report, wr):
+ report.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.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?"),
+ ]
+ report.save()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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`
+ """)
+ return
+
+
+@app.cell
+def _(ENTITY, PROJECT, report, wr):
+ pg = wr.PanelGrid(
+ runsets=[
+ wr.Runset(ENTITY, PROJECT, "First Run Set"),
+ 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}),
+ ]
+ )
+
+ report.blocks = report.blocks[:1] + [wr.H1("Panel Grid Example"), pg] + report.blocks[1:]
+ report.save()
+ return (pg,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Add data lineage with Artifact blocks
+ - There are equivalent weave panels as well
+ """)
+ return
+
+
+@app.cell
+def _(ENTITY, LINEAGE_PROJECT, report, wr):
+ artifact_lineage = wr.WeaveBlockArtifact(entity=ENTITY, project=LINEAGE_PROJECT, artifact='model-1', tab='lineage')
+
+ report.blocks = report.blocks[:1] + [wr.H1("Artifact lineage example"), artifact_lineage] + report.blocks[1:]
+ report.save()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Customize run colors
+ - Pass in a `dict[run_name, color]`
+ """)
+ return
+
+
+@app.cell
+def _(pg, report):
+ 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.save()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ❓ FAQ
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## My report is too wide/narrow
+ - Change the report's width to the right size for you.
+ """)
+ return
+
+
+@app.cell
+def _(report):
+ report2 = report.save(clone=True)
+ report2.width = 'fluid'
+ report2.save()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, wr):
+ report_1 = wr.Report(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))])])
+ report_1.save() # Adjusting the layout for the plot size
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.` to get autocomplete.
+ """)
+ return
+
+
+@app.cell
+def _(ENTITY, LINEAGE_PROJECT, PROJECT, wr):
+ report_2 = wr.Report(project=PROJECT, title='W&B Block Gallery', description='Check out all of the blocks available in W&B', blocks=[wr.H1(text='Heading 1'), wr.P(text='Normal paragraph'), wr.H2(text='Heading 2'), wr.P(text=['here is some text, followed by', wr.InlineCode(text='select * from code in line'), 'and then latex', wr.InlineLatex(text='e=mc^2')]), wr.H3(text='Heading 3'), wr.CodeBlock(code='this:\n- is\n- a\ncool:\n- yaml\n- file', language='yaml'), wr.WeaveBlockSummaryTable(entity=ENTITY, project=PROJECT, table_name='my-table'), wr.WeaveBlockArtifact(entity=ENTITY, project=LINEAGE_PROJECT, artifact='model-1', tab='lineage'), wr.WeaveBlockArtifactVersionedFile(entity=ENTITY, project=LINEAGE_PROJECT, artifact='model-1', version='v0', file='dataframe.table.json'), wr.MarkdownBlock(text='Markdown cell with *italics* and **bold** and $e=mc^2$'), wr.LatexBlock(text='\\gamma^2+\\theta^2=\\omega^2\n\\\\ a^2 + b^2 = c^2'), wr.Image(url='https://api.wandb.ai/files/megatruong/images/projects/918598/350382db.gif', caption="It's a me, Pikachu"), wr.UnorderedList(items=['Bullet 1', 'Bullet 2']), wr.OrderedList(items=['Ordered 1', 'Ordered 2']), wr.CheckedList(items=[wr.CheckedListItem(text='Unchecked', checked=False), wr.CheckedListItem(text='Checked', checked=True)]), wr.BlockQuote(text='Block Quote 1\nBlock Quote 2\nBlock Quote 3'), wr.CalloutBlock(text='Callout 1\nCallout 2\nCallout 3'), wr.HorizontalRule(), wr.Video(url='https://www.youtube.com/embed/6riDJMI-Y8U')])
+ report_2.save()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## What panels are available?
+ - See `wr.panels` for a list of panels
+ - In an IDE or notebook, you can also do `wr.panels.` to get autocomplete.
+ - Panels have a lot of settings. Inspect the panel to see what you can do!
+ """)
+ return
+
+
+@app.cell
+def _(LINEAGE_PROJECT, PROJECT, wr):
+ report_3 = wr.Report(project=PROJECT, title='W&B Panel Gallery', description='Check out all of the panels available in W&B', width='fluid', blocks=[wr.PanelGrid(runsets=[wr.Runset(project=LINEAGE_PROJECT), wr.Runset()], panels=[wr.MediaBrowser(media_keys=['img']), wr.MarkdownPanel(markdown='Hello *italic* **bold** $e=mc^2$ `something`'), wr.LinePlot(title='Validation Accuracy over Time', x='Step', y=['val_acc'], range_x=(0, 1000), range_y=(1, 4), log_x=True, log_y=False, title_x='Training steps', title_y='Validation Accuracy', ignore_outliers=True, groupby='encoder', groupby_aggfunc='mean', groupby_rangefunc='minmax', smoothing_factor=0.5, smoothing_type='gaussian', smoothing_show_original=True, max_runs_to_show=10, font_size='large', legend_position='west'), wr.ScatterPlot(title='Validation Accuracy vs. Validation Loss', x='val_acc', y='val_loss', log_x=False, log_y=False, running_ymin=True, running_ymean=True, running_ymax=True, font_size='small', regression=True), wr.BarPlot(title='Validation Loss by Encoder', metrics=['val_loss'], orientation='h', range_x=(0, 0.11), title_x='Validation Loss', groupby='encoder', groupby_aggfunc='median', groupby_rangefunc='stddev', max_runs_to_show=20, max_bars_to_show=3, font_size='auto'), wr.ScalarChart(title='Maximum Number of Steps', metric='Step', groupby_aggfunc='max', groupby_rangefunc='stderr', font_size='large'), wr.CodeComparer(diff='split'), wr.ParallelCoordinatesPlot(columns=[wr.ParallelCoordinatesPlotColumn('Step'), wr.ParallelCoordinatesPlotColumn('c::model'), wr.ParallelCoordinatesPlotColumn('c::optimizer'), wr.ParallelCoordinatesPlotColumn('val_acc'), wr.ParallelCoordinatesPlotColumn('val_loss')]), wr.ParameterImportancePlot(with_respect_to='val_loss'), wr.RunComparer(diff_only=True), wr.CustomChart(query={'summary': ['val_loss', 'val_acc']}, chart_name='wandb/scatter/v0', chart_fields={'x': 'val_loss', 'y': 'val_acc'})]), wr.WeaveBlockSummaryTable(entity='your_entity', project='your_project', table_name='my-table'), wr.WeaveBlockArtifact(entity='your_entity', project='your_project', artifact='model-1', tab='lineage'), wr.WeaveBlockArtifactVersionedFile(entity='your_entity', project='your_project', artifact='model-1', version='v0', file='dataframe.table.json')])
+ report_3.save() # LinePlot with various settings enabled # Add WeaveBlock types directly to the blocks list
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## How can I link related reports together?
+ - Suppose have have two reports like below:
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, wr):
+ report1 = wr.Report(project=PROJECT, title='Report 1', description='Great content coming from Report 1', blocks=[wr.H1(text='Heading from Report 1'), wr.P(text='Lorem ipsum dolor sit amet. Aut fuga minus nam vero saepeA aperiam eum omnis dolorum et ducimus tempore aut illum quis aut alias vero. Sed explicabo illum est eius quianon vitae sed voluptatem incidunt. Vel architecto assumenda Ad voluptatem quo dicta provident et velit officia. Aut galisum inventoreSed dolore a illum adipisci a aliquam quidem sit corporis quia cum magnam similique.'), wr.PanelGrid(panels=[wr.LinePlot(title='Episodic Return', x='global_step', y=['charts/episodic_return'], smoothing_factor=0.85, groupby_aggfunc='mean', groupby_rangefunc='minmax', layout=wr.Layout(x=0, y=0, w=12, h=8)), wr.MediaBrowser(media_keys=['videos'], num_columns=4, layout=wr.Layout(w=12, h=8))], runsets=[wr.Runset(entity='openrlbenchmark', project='cleanrl', query='bigfish', groupby=['env_id', 'exp_name'])], custom_run_colors={wr.RunsetGroup(runset_name='Run set', keys=(wr.RunsetGroupKey(key='bigfish', value='ppg_procgen'),)): '#2980b9', wr.RunsetGroup(runset_name='Run set', keys=(wr.RunsetGroupKey(key='bigfish', value='ppo_procgen'),)): '#e74c3c'})])
+ report1.save()
+ report2_1 = wr.Report(project=PROJECT, title='Report 2', description='Great content coming from Report 2', blocks=[wr.H1(text='Heading from Report 2'), wr.P(text='Est quod ducimus ut distinctio corruptiid optio qui cupiditate quibusdam ea corporis modi. Eum architecto vero sed error dignissimosEa repudiandae a recusandae sint ut sint molestiae ea pariatur quae. In pariatur voluptas ad facere neque 33 suscipit et odit nostrum ut internos molestiae est modi enim. Et rerum inventoreAut internos et dolores delectus aut Quis sunt sed nostrum magnam ab dolores dicta.'), wr.PanelGrid(panels=[wr.LinePlot(title='SPS', x='global_step', y=['charts/SPS']), wr.LinePlot(title='Episodic Length', x='global_step', y=['charts/episodic_length']), wr.LinePlot(title='Episodic Return', x='global_step', y=['charts/episodic_return'])], runsets=[wr.Runset(entity='openrlbenchmark', project='cleanrl', name='DQN', groupby=['exp_name']), wr.Runset(entity='openrlbenchmark', project='cleanrl', name='SAC-discrete 0.8', groupby=['exp_name']), wr.Runset(entity='openrlbenchmark', project='cleanrl', name='SAC-discrete 0.88', groupby=['exp_name'])], custom_run_colors={wr.RunsetGroup(runset_name='DQN', keys=(wr.RunsetGroupKey(key='dqn_atari', value='exp_name'),)): '#e84118', wr.RunsetGroup(runset_name='SAC-discrete 0.8', keys=(wr.RunsetGroupKey(key='sac_atari', value='exp_name'),)): '#fbc531', wr.RunsetGroup(runset_name='SAC-discrete 0.88', keys=(wr.RunsetGroupKey(key='sac_atari', value='exp_name'),)): '#00a8ff'})])
+ report2_1.save()
+ return report1, report2_1
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Combine blocks into a new report
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, wr):
+ report_4 = wr.Report(PROJECT, title='Report with links', description='Use `wr.Link(text, url)` to add links inside normal text, or use normal markdown syntax in a MarkdownBlock', blocks=[wr.H1('This is a normal heading'), wr.P('And here is some normal text'), wr.H1(['This is a heading ', wr.Link('with a link!', url='https://wandb.ai/')]), wr.P(['Most text formats support ', wr.Link('adding links', url='https://wandb.ai/')]), wr.MarkdownBlock('You can also use markdown syntax for [links](https://wandb.ai/)')])
+ report_4.save()
+ return
+
+
+@app.cell
+def _(PROJECT, report1, report2_1, wr):
+ report3 = wr.Report(PROJECT, title='Combined blocks report', description='This report combines blocks from both Report 1 and Report 2', blocks=[*report1.blocks, *report2_1.blocks])
+ report3.save()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## I tried mutating an object in list but it didn't work!
+ tl;dr: It should always work if you assign a value to the attribute instead of mutating. If you really need to mutate, do it before assignment.
+
+ ---
+
+ This can happen in a few places that contain lists of wandb objects, e.g.:
+ - `report.blocks`
+ - `panel_grid.panels`
+ - `panel_grid.runsets`
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, wr):
+ report_5 = wr.Report(project=PROJECT)
+ return (report_5,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Good: Assign `b`
+ """)
+ return
+
+
+@app.cell
+def _(report_5, wr):
+ b = wr.H1(text=['Hello', ' World!'])
+ report_5.blocks = [b]
+ assert b.text == ['Hello', ' World!']
+ assert report_5.blocks[0].text == ['Hello', ' World!']
+ return (b,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Bad: Mutate `b` without reassigning
+ """)
+ return
+
+
+@app.cell
+def _(b, report_5):
+ b.text = ['Something', ' New']
+ assert b.text == ['Something', ' New']
+ # This will error!
+ assert report_5.blocks[0].text == ['Hello', ' World!']
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Good: Mutate `b` and then reassign it
+ """)
+ return
+
+
+@app.cell
+def _(b, report_5):
+ report_5.blocks = [b]
+ assert b.text == ['Something', ' New']
+ assert report_5.blocks[0].text == ['Something', ' New']
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/intro-run-quickstart/intro_run_quickstart.py b/marimo/convert/intro-run-quickstart/intro_run_quickstart.py
new file mode 100644
index 00000000..c7afaf6c
--- /dev/null
+++ b/marimo/convert/intro-run-quickstart/intro_run_quickstart.py
@@ -0,0 +1,113 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Use W&B to track, visualize, and manage machine learning experiments of any size.
+
+ Install W&B to track, visualize, and manage machine learning experiments of any size.
+
+ ## Install W&B Python SDK
+
+ Install the W&B Python SDK (`wandb`) with your preferred Python package installer. This notebook uses `pip`:
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, import the W&B Python SDK and other Python packages you will use in this notebook:
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from getpass import getpass
+ import random
+ import os
+
+ return getpass, os, random, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log in
+
+ To authenticate your machine with W&B, you need a W&B API key. Run the following cell and, when prompted, enter your API key:
+ """)
+ return
+
+
+@app.cell
+def _(getpass, os):
+ os.environ["WANDB_API_KEY"] = getpass("Enter your W&B API key: ")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Create a machine learning training experiment
+
+ The following example simulates a simple training experiment and logs metrics to W&B.
+
+ First, define the W&B project name and a `config` dictionary. The config stores the input values for the experiment, such as the number of epochs and the learning rate.
+
+ Next, initialize a W&B run with [`wandb.init()`](https://docs.wandb.ai/models/ref/python/functions/init). The run records the config, metrics, and other information from the training script.
+
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(random, wandb):
+ # Project that the run is recorded to
+ project = "my-awesome-project"
+
+ # Dictionary with hyperparameters
+ config = {
+ 'epochs' : 10,
+ 'lr' : 0.01
+ }
+
+ with wandb.init(project=project, config=config) as run:
+ offset = random.random() / 5
+ print(f"lr: {config['lr']}")
+
+ # Simulate a training run
+ for epoch in range(2, config['epochs']):
+ 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({"accuracy": acc, "loss": loss})
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/jax-simple-training-loop-in-jax-and-flax/jax_simple_training_loop_in_jax_and_flax.py b/marimo/convert/jax-simple-training-loop-in-jax-and-flax/jax_simple_training_loop_in_jax_and_flax.py
new file mode 100644
index 00000000..c218e050
--- /dev/null
+++ b/marimo/convert/jax-simple-training-loop-in-jax-and-flax/jax_simple_training_loop_in_jax_and_flax.py
@@ -0,0 +1,574 @@
+# /// script
+# dependencies = ["flax", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Writing a Simple Training Loop in JAX and FLAX
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Packages 📦 and Basic Setup
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### ❤️ Install Packages
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb flax !pip install -q wandb flax
+ return
+
+
+@app.cell
+def _():
+ import jax
+ import jax.numpy as jnp
+
+ import optax
+
+ from flax import linen as nn
+ from flax.training import train_state
+ from flax.serialization import (
+ to_state_dict, msgpack_serialize, from_bytes
+ )
+
+ import os
+ import wandb
+ import numpy as np
+ from typing import Callable
+ from tqdm.auto import tqdm
+
+ import tensorflow as tf
+ import tensorflow_datasets as tfds
+
+ return (
+ Callable,
+ from_bytes,
+ jax,
+ jnp,
+ msgpack_serialize,
+ nn,
+ np,
+ optax,
+ os,
+ tf,
+ tfds,
+ to_state_dict,
+ tqdm,
+ train_state,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### ⚙️ Project Configuration using **`wandb.config`**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can now call [**`wandb.init`**](https://docs.wandb.ai/guides/track/launch) to initialize a new job. This creates a new run in [**Weights & Biases**](https://wandb.ai/site) and launches a background process to sync data. We will also sync all the configs of our experiments with the W&B run, which makes it far easier for us to reproduce the results of the experiment later.
+ """)
+ return
+
+
+@app.cell
+def _(nn, wandb):
+ wandb.init(
+ project="simple-training-loop",
+ entity="jax-series",
+ job_type="simple-train-loop"
+ )
+
+ config = wandb.config
+ config.seed = 42
+ config.batch_size = 64
+ config.validation_split = 0.2
+ config.pooling = "avg"
+ config.learning_rate = 1e-4
+ config.epochs = 15
+
+ MODULE_DICT = {
+ "avg": nn.avg_pool,
+ "max": nn.max_pool,
+ }
+ return MODULE_DICT, config
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💿 The Dataset
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ most JAX practitioners prefer to use the **`tf.data`** API for building data loading pipelines for JAX and Flax-based machine learning workflow. In this notebook, we will build a simple data loading pipeline for the CIFAR-10 dataset using Tensorflow Datasets for Image Classification.
+ """)
+ return
+
+
+@app.cell
+def _(config, tf, tfds):
+ (full_train_set, test_dataset), ds_info = tfds.load(
+ 'cifar10',
+ split=['train', 'test'],
+ shuffle_files=True,
+ as_supervised=True,
+ with_info=True,
+ )
+
+ def normalize_img(image, label):
+ image = tf.cast(image, tf.float32) / 255.
+ return image, label
+
+ full_train_set = full_train_set.map(
+ normalize_img, num_parallel_calls=tf.data.AUTOTUNE
+ )
+
+ num_data = tf.data.experimental.cardinality(
+ full_train_set
+ ).numpy()
+ print("Total number of data points:", num_data)
+ train_dataset = full_train_set.take(
+ num_data * (1 - config.validation_split)
+ )
+ val_dataset = full_train_set.take(
+ num_data * (config.validation_split)
+ )
+ print(
+ "Number of train data points:",
+ tf.data.experimental.cardinality(train_dataset).numpy()
+ )
+ print(
+ "Number of val data points:",
+ tf.data.experimental.cardinality(val_dataset).numpy()
+ )
+
+ train_dataset = train_dataset.cache()
+ train_dataset = train_dataset.shuffle(
+ tf.data.experimental.cardinality(train_dataset).numpy()
+ )
+ train_dataset = train_dataset.batch(config.batch_size)
+
+ val_dataset = val_dataset.cache()
+ val_dataset = val_dataset.shuffle(
+ tf.data.experimental.cardinality(val_dataset).numpy()
+ )
+ val_dataset = val_dataset.batch(config.batch_size)
+
+
+ test_dataset = test_dataset.map(
+ normalize_img, num_parallel_calls=tf.data.AUTOTUNE
+ )
+ print(
+ "Number of test data points:",
+ tf.data.experimental.cardinality(test_dataset).numpy()
+ )
+ test_dataset = test_dataset.cache()
+ test_dataset = test_dataset.batch(config.batch_size)
+ return test_dataset, train_dataset, val_dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ✍️ Model Architecture
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let us now define a very simple classification convolution based neural network. Instead of some famous architecture we'll create a simple custom architecture by subclassing [**`linen.Module`**](https://flax.readthedocs.io/en/latest/_modules/flax/linen/module.html#Module).
+ """)
+ return
+
+
+@app.cell
+def _(Callable, nn):
+ class CNN(nn.Module):
+ pool_module: Callable = nn.avg_pool
+
+ def setup(self):
+ self.conv_1 = nn.Conv(features=32, kernel_size=(3, 3))
+ self.conv_2 = nn.Conv(features=32, kernel_size=(3, 3))
+ self.conv_3 = nn.Conv(features=64, kernel_size=(3, 3))
+ self.conv_4 = nn.Conv(features=64, kernel_size=(3, 3))
+ self.conv_5 = nn.Conv(features=128, kernel_size=(3, 3))
+ self.conv_6 = nn.Conv(features=128, kernel_size=(3, 3))
+ self.dense_1 = nn.Dense(features=1024)
+ self.dense_2 = nn.Dense(features=512)
+ self.dense_output = nn.Dense(features=10)
+
+ @nn.compact
+ def __call__(self, x):
+ x = nn.relu(self.conv_1(x))
+ x = nn.relu(self.conv_2(x))
+ x = self.pool_module(x, window_shape=(2, 2), strides=(2, 2))
+ x = nn.relu(self.conv_3(x))
+ x = nn.relu(self.conv_4(x))
+ x = self.pool_module(x, window_shape=(2, 2), strides=(2, 2))
+ x = nn.relu(self.conv_5(x))
+ x = nn.relu(self.conv_6(x))
+ x = self.pool_module(x, window_shape=(2, 2), strides=(2, 2))
+ x = x.reshape((x.shape[0], -1))
+ x = nn.relu(self.dense_1(x))
+ x = nn.relu(self.dense_2(x))
+ return self.dense_output(x)
+
+ return (CNN,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now that we have defined, the CNN Module, we would need to initialize it. However, unlike Tensorflow or PyTorch, the parameters of a Flax Module are not stored with the models themselves. We would need to initialize parameters by calling the init function, using a PRNG Key and a dummy input parameter with the same shape as the expected input.
+ """)
+ return
+
+
+@app.cell
+def _(CNN, MODULE_DICT, config, jax, jnp):
+ rng = jax.random.PRNGKey(config.seed)
+ x = jnp.ones(shape=(config.batch_size, 32, 32, 3))
+ model = CNN(pool_module=MODULE_DICT[config.pooling])
+ params = model.init(rng, x)
+ jax.tree_map(lambda x: x.shape, params)
+ return model, rng, x
+
+
+@app.cell
+def _(model, nn, rng, x):
+ nn.tabulate(model, rng)(x)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ After we initialize the model we'll use the variables to create a [**TrainState**](https://flax.readthedocs.io/en/latest/flax.training.html#flax.training.train_state.TrainState), a utility class for handling parameter and gradient updates. This is a key feature of the new Flax version. Instead of initializing the model again and again with new variables we just update the "state" of the model and pass this as inputs to functions.
+ """)
+ return
+
+
+@app.cell
+def _(config, jnp, model, optax, rng, train_state):
+ def init_train_state(
+ model, random_key, shape, learning_rate
+ ) -> train_state.TrainState:
+ # Initialize the Model
+ variables = model.init(random_key, jnp.ones(shape))
+ # Create the optimizer
+ optimizer = optax.adam(learning_rate)
+ # Create a State
+ return train_state.TrainState.create(
+ apply_fn = model.apply,
+ tx=optimizer,
+ params=variables['params']
+ )
+
+
+ state = init_train_state(
+ model, rng, (config.batch_size, 32, 32, 3), config.learning_rate
+ )
+ print(type(state))
+ return (state,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### ⚙️ Utility Functions
+ """)
+ return
+
+
+@app.cell
+def _(jax, optax):
+ def cross_entropy_loss(*, logits, labels):
+ one_hot_encoded_labels = jax.nn.one_hot(labels, num_classes=10)
+ return optax.softmax_cross_entropy(
+ logits=logits, labels=one_hot_encoded_labels
+ ).mean()
+
+ return (cross_entropy_loss,)
+
+
+@app.cell
+def _(cross_entropy_loss, jnp):
+ def compute_metrics(*, logits, labels):
+ loss = cross_entropy_loss(logits=logits, labels=labels)
+ accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
+ metrics = {
+ 'loss': loss,
+ 'accuracy': accuracy,
+ }
+ return metrics
+
+ return (compute_metrics,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🧱 + 🏗 = 🏠 Training
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ * Any train step should take in two basic parameters; the state and the batch (or whatever format the input is) in question.
+
+ * We usually define the loss function within this function as best practice. We get the logits from the model, using the `apply_fn` from the TrainState (which is just the apply method of the model) by passing the parameters and the input. We then compute the loss by using the logits and input and return the loss as well as the logits (this is key).
+
+ * We then transform the function using `jax.value_and_grad()` transformation. Instead of `jax.grad()` which just creates a function which returns the derivative of the function. We use `jax.value_and_grad()` which returns the gradient as well as the evaluation of the function. (Notice the `has_aux` parameter, we set this to True because the loss function returns the loss as well as the logits, an auxiliary object)
+
+ * We then calculate the gradients and obtain the logits by passing in the parameters of the state. Notice how the function returns both the gradients and the logits (because we used `jax.value_and_grad()` instead of `jax.grad()`) we'll later need these logits to calculate metrics after the step
+
+ * We then essentially perform backpropagation by updating the TrainState using the calculated gradients by using the `.apply_gradients()` method
+
+ * Calculate the metrics using the utility `compute_metrics` function.
+ """)
+ return
+
+
+@app.cell
+def _(compute_metrics, cross_entropy_loss, jax, jnp, train_state):
+ @jax.jit
+ def train_step(
+ state: train_state.TrainState, batch: jnp.ndarray
+ ):
+ image, label = batch
+
+ def loss_fn(params):
+ logits = state.apply_fn({'params': params}, image)
+ loss = cross_entropy_loss(logits=logits, labels=label)
+ return loss, logits
+
+ gradient_fn = jax.value_and_grad(loss_fn, has_aux=True)
+ (_, logits), grads = gradient_fn(state.params)
+ state = state.apply_gradients(grads=grads)
+ metrics = compute_metrics(logits=logits, labels=label)
+ return state, metrics
+
+ return (train_step,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Similar to our `train_step` this function also takes the state and the batch. We simply perform a forward pass using the data and obtain the logits and then compute the corresponding metrics. As this is the `eval_step` we don't compute the gradients or update the parameters of the TrainState.
+ """)
+ return
+
+
+@app.cell
+def _(compute_metrics, jax):
+ @jax.jit
+ def eval_step(state, batch):
+ image, label = batch
+ logits = state.apply_fn({'params': state.params}, image)
+ return compute_metrics(logits=logits, labels=label)
+
+ return (eval_step,)
+
+
+@app.cell
+def _(from_bytes, jax, msgpack_serialize, np, os, to_state_dict, wandb):
+ def save_checkpoint(ckpt_path, state, epoch):
+ with open(ckpt_path, "wb") as outfile:
+ outfile.write(msgpack_serialize(to_state_dict(state)))
+ artifact = wandb.Artifact(
+ f'{wandb.run.name}-checkpoint', type='dataset'
+ )
+ artifact.add_file(ckpt_path)
+ wandb.log_artifact(artifact, aliases=["latest", f"epoch_{epoch}"])
+
+
+ def load_checkpoint(ckpt_file, state):
+ artifact = wandb.use_artifact(
+ f'{wandb.run.name}-checkpoint:latest'
+ )
+ artifact_dir = artifact.download()
+ ckpt_path = os.path.join(artifact_dir, ckpt_file)
+ with open(ckpt_path, "rb") as data_file:
+ byte_data = data_file.read()
+ return from_bytes(state, byte_data)
+
+
+ def accumulate_metrics(metrics):
+ metrics = jax.device_get(metrics)
+ return {
+ k: np.mean([metric[k] for metric in metrics])
+ for k in metrics[0]
+ }
+
+ return accumulate_metrics, load_checkpoint, save_checkpoint
+
+
+@app.cell
+def _(
+ accumulate_metrics,
+ eval_step,
+ load_checkpoint,
+ save_checkpoint,
+ tf,
+ tfds,
+ tqdm,
+ train_state,
+ train_step,
+ wandb,
+):
+ def train_and_evaluate(
+ train_dataset,
+ eval_dataset,
+ test_dataset,
+ state: train_state.TrainState,
+ epochs: int,
+ ):
+ num_train_batches = tf.data.experimental.cardinality(train_dataset)
+ num_eval_batches = tf.data.experimental.cardinality(eval_dataset)
+ num_test_batches = tf.data.experimental.cardinality(test_dataset)
+
+ for epoch in tqdm(range(1, epochs + 1)):
+
+ best_eval_loss = 1e6
+
+ train_batch_metrics = []
+ train_datagen = iter(tfds.as_numpy(train_dataset))
+ for batch_idx in range(num_train_batches):
+ batch = next(train_datagen)
+ state, metrics = train_step(state, batch)
+ train_batch_metrics.append(metrics)
+
+ train_batch_metrics = accumulate_metrics(train_batch_metrics)
+ print(
+ 'TRAIN (%d/%d): Loss: %.4f, accuracy: %.2f' % (
+ epoch, epochs, train_batch_metrics['loss'],
+ train_batch_metrics['accuracy'] * 100
+ )
+ )
+
+ eval_batch_metrics = []
+ eval_datagen = iter(tfds.as_numpy(eval_dataset))
+ for batch_idx in range(num_eval_batches):
+ batch = next(eval_datagen)
+ metrics = eval_step(state, batch)
+ eval_batch_metrics.append(metrics)
+
+ eval_batch_metrics = accumulate_metrics(eval_batch_metrics)
+ print(
+ 'EVAL (%d/%d): Loss: %.4f, accuracy: %.2f\n' % (
+ epoch, epochs, eval_batch_metrics['loss'],
+ eval_batch_metrics['accuracy'] * 100
+ )
+ )
+
+ wandb.log({
+ "Train Loss": train_batch_metrics['loss'],
+ "Train Accuracy": train_batch_metrics['accuracy'],
+ "Validation Loss": eval_batch_metrics['loss'],
+ "Validation Accuracy": eval_batch_metrics['accuracy']
+ }, step=epoch)
+
+ if eval_batch_metrics['loss'] < best_eval_loss:
+ save_checkpoint("checkpoint.msgpack", state, epoch)
+
+ restored_state = load_checkpoint("checkpoint.msgpack", state)
+ test_batch_metrics = []
+ test_datagen = iter(tfds.as_numpy(test_dataset))
+ for batch_idx in range(num_test_batches):
+ batch = next(test_datagen)
+ metrics = eval_step(restored_state, batch)
+ test_batch_metrics.append(metrics)
+
+ test_batch_metrics = accumulate_metrics(test_batch_metrics)
+ print(
+ 'Test: Loss: %.4f, accuracy: %.2f' % (
+ test_batch_metrics['loss'],
+ test_batch_metrics['accuracy'] * 100
+ )
+ )
+
+ wandb.log({
+ "Test Loss": test_batch_metrics['loss'],
+ "Test Accuracy": test_batch_metrics['accuracy']
+ })
+
+ return state, restored_state
+
+ return (train_and_evaluate,)
+
+
+@app.cell
+def _(
+ config,
+ state,
+ test_dataset,
+ train_and_evaluate,
+ train_dataset,
+ val_dataset,
+):
+ state_1, best_state = train_and_evaluate(train_dataset, val_dataset, test_dataset, state, epochs=config.epochs)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/jax-training-with-tfrecords-in-jax-imagenette/jax_training_with_tfrecords_in_jax_imagenette.py b/marimo/convert/jax-training-with-tfrecords-in-jax-imagenette/jax_training_with_tfrecords_in_jax_imagenette.py
new file mode 100644
index 00000000..834557c8
--- /dev/null
+++ b/marimo/convert/jax-training-with-tfrecords-in-jax-imagenette/jax_training_with_tfrecords_in_jax_imagenette.py
@@ -0,0 +1,495 @@
+# /// script
+# dependencies = ["flax", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb flax !pip install -q wandb flax
+ return
+
+
+@app.cell
+def _():
+ import os
+ import wandb
+ import numpy as np
+ from glob import glob
+ from typing import Callable
+ from tqdm.auto import tqdm
+ import matplotlib.pyplot as plt
+
+ import jax
+ import jax.numpy as jnp
+
+ import optax
+
+ from flax import linen as nn
+ from flax.training import train_state
+ from flax.serialization import (
+ to_state_dict, msgpack_serialize, from_bytes
+ )
+
+ import tensorflow as tf
+ import tensorflow_datasets as tfds
+ AUTOTUNE = tf.data.AUTOTUNE
+ return (
+ AUTOTUNE,
+ Callable,
+ from_bytes,
+ glob,
+ jax,
+ jnp,
+ msgpack_serialize,
+ nn,
+ np,
+ optax,
+ os,
+ plt,
+ tf,
+ tfds,
+ to_state_dict,
+ tqdm,
+ train_state,
+ wandb,
+ )
+
+
+@app.cell
+def _(nn, wandb):
+ wandb.init(
+ project="simple-training-loop",
+ entity="jax-series",
+ job_type="tfrecord"
+ )
+
+ config = wandb.config
+ config.seed = 42
+ config.image_size = 227
+ config.batch_size = 64
+ config.pooling = "max"
+ config.learning_rate = 1e-4
+ config.epochs = 15
+ config.artifact_address = 'jax-series/simple-training-loop/imagenette-tfrecords:v3'
+ config.labels = [
+ 'tench', 'english_springer', 'english_springer', 'chain_saw',
+ 'church', 'french_horn', 'grabage_truck', 'gas_pump',
+ 'golf_ball', 'parachute'
+ ]
+
+ MODULE_DICT = {
+ "avg": nn.avg_pool,
+ "max": nn.max_pool,
+ }
+ return MODULE_DICT, config
+
+
+@app.cell
+def _(AUTOTUNE, tf):
+ def decode_image(image_data):
+ image = tf.image.decode_jpeg(image_data, channels=3)
+ image = tf.cast(image, tf.float32) / 255.0
+ return image
+
+ def read_labeled_tfrecord(example):
+ feature = {
+ "image": tf.io.FixedLenFeature([], tf.string),
+ "label": tf.io.FixedLenFeature([], tf.int64),
+ }
+
+ example = tf.io.parse_single_example(example, feature)
+ image = decode_image(example['image'])
+ label = tf.cast(example['label'], tf.int32)
+ return image, label
+
+ def load_dataset(filenames, ordered = False):
+
+ ignore_order = tf.data.Options()
+ if not ordered:
+ ignore_order.experimental_deterministic = False
+
+ dataset = tf.data.TFRecordDataset(
+ filenames, num_parallel_reads=AUTOTUNE
+ )
+ dataset_len = sum(1 for _ in dataset)
+ dataset = dataset.with_options(ignore_order)
+ dataset = dataset.map(
+ read_labeled_tfrecord, num_parallel_calls=AUTOTUNE
+ )
+ return dataset, dataset_len
+
+ return (load_dataset,)
+
+
+@app.cell
+def _(config, glob, os, wandb):
+ artifact = wandb.use_artifact(
+ config.artifact_address, type='dataset'
+ )
+ artifact_dir = artifact.download()
+ train_files = glob(os.path.join(artifact_dir, "train", "*.tfrec"))
+ val_files = glob(os.path.join(artifact_dir, "val", "*.tfrec"))
+ return train_files, val_files
+
+
+@app.cell
+def _(load_dataset, train_files):
+ sample_dataset, _ = load_dataset(train_files)
+ sample_dataset = sample_dataset.shuffle(1024)
+ sample_dataset.element_spec
+ return (sample_dataset,)
+
+
+@app.cell
+def _(config, plt, sample_dataset):
+ plt.figure(figsize=(16, 16))
+ for i in range(16):
+ x, y = next(iter(sample_dataset))
+ x, y = x.numpy(), y.numpy().tolist()
+ ax = plt.subplot(4, 4, i + 1)
+ plt.imshow(x)
+ plt.axis("off")
+ name = config.labels[y]
+ ax.set_title(name, fontsize=20)
+ return
+
+
+@app.cell
+def _(AUTOTUNE, config, load_dataset, tf):
+ def resize_image(image, label):
+ image = tf.image.resize(
+ image, [config.image_size, config.image_size]
+ )
+ return image, label
+
+
+ def data_augment(image, label):
+ image = tf.image.random_flip_left_right(image)
+ image = tf.image.random_hue(image, 0.01)
+ image = tf.image.random_saturation(image, 0.70, 1.30)
+ image = tf.image.random_contrast(image, 0.80, 1.20)
+ image = tf.image.random_brightness(image, 0.10)
+ return image, label
+
+ def get_training_dataset(filenames, batch_size):
+ dataset, dataset_len = load_dataset(filenames, ordered = False)
+ dataset = dataset.map(
+ resize_image, num_parallel_calls=AUTOTUNE
+ )
+ dataset = dataset.map(
+ data_augment, num_parallel_calls=AUTOTUNE
+ )
+ dataset = dataset.repeat()
+ dataset = dataset.shuffle(2048)
+ dataset = dataset.batch(batch_size)
+ dataset = dataset.prefetch(AUTOTUNE)
+ return dataset, dataset_len // batch_size
+
+ def get_val_dataset(filenames, batch_size):
+ dataset, dataset_len = load_dataset(filenames, ordered = True)
+ dataset = dataset.map(
+ resize_image, num_parallel_calls=AUTOTUNE
+ )
+ dataset = dataset.map(
+ data_augment, num_parallel_calls=AUTOTUNE
+ )
+ dataset = dataset.batch(batch_size)
+ dataset = dataset.prefetch(AUTOTUNE)
+ return dataset, dataset_len // batch_size
+
+ return (get_training_dataset,)
+
+
+@app.cell
+def _(config, get_training_dataset, train_files, val_files):
+ train_dataset, num_train_batches = get_training_dataset(train_files, config.batch_size)
+ val_dataset, num_val_batches = get_training_dataset(val_files, config.batch_size)
+ return num_train_batches, num_val_batches, train_dataset, val_dataset
+
+
+@app.cell
+def _(Callable, nn):
+ class AlexNet(nn.Module):
+ num_classes: int
+ pool_module: Callable = nn.avg_pool
+
+ def setup(self):
+ self.conv_1 = nn.Conv(
+ features=96, kernel_size=(11, 11), strides=4, padding="VALID"
+ )
+ self.conv_2 = nn.Conv(
+ features=256, kernel_size=(5, 5), strides=1, padding="VALID"
+ )
+ self.conv_3 = nn.Conv(
+ features=384, kernel_size=(3, 3), strides=1, padding="VALID"
+ )
+ self.conv_4 = nn.Conv(
+ features=384, kernel_size=(3, 3), strides=1, padding="VALID"
+ )
+ self.conv_5 = nn.Conv(
+ features=256, kernel_size=(3, 3), strides=1, padding="VALID"
+ )
+ self.dense_1 = nn.Dense(features=1024)
+ self.dense_2 = nn.Dense(features=512)
+ self.dense_output = nn.Dense(features=self.num_classes)
+
+ def __call__(self, x):
+ x = nn.relu(self.conv_1(x))
+ x = self.pool_module(x, window_shape=(3, 3), strides=(2, 2))
+ x = nn.relu(self.conv_2(x))
+ x = self.pool_module(x, window_shape=(3, 3), strides=(2, 2))
+ x = nn.relu(self.conv_3(x))
+ x = nn.relu(self.conv_4(x))
+ x = nn.relu(self.conv_5(x))
+ x = self.pool_module(x, window_shape=(3, 3), strides=(2, 2))
+ x = x.reshape((x.shape[0], -1))
+ x = nn.relu(self.dense_1(x))
+ x = nn.relu(self.dense_2(x))
+ return self.dense_output(x)
+
+ return (AlexNet,)
+
+
+@app.cell
+def _(AlexNet, MODULE_DICT, config, jax, jnp):
+ rng = jax.random.PRNGKey(config.seed)
+ x_1 = jnp.ones(shape=(config.batch_size, config.image_size, config.image_size, 3))
+ model = AlexNet(num_classes=len(config.labels), pool_module=MODULE_DICT[config.pooling])
+ params = model.init(rng, x_1)
+ jax.tree_map(lambda x: x.shape, params)
+ return model, rng, x_1
+
+
+@app.cell
+def _(model, nn, rng, x_1):
+ nn.tabulate(model, rng)(x_1)
+ return
+
+
+@app.cell
+def _(config, jnp, model, optax, rng, train_state):
+ def init_train_state(
+ model, random_key, shape, learning_rate
+ ) -> train_state.TrainState:
+ variables = model.init(random_key, jnp.ones(shape))
+ optimizer = optax.adam(learning_rate)
+ return train_state.TrainState.create(
+ apply_fn = model.apply,
+ tx=optimizer,
+ params=variables['params']
+ )
+
+
+ state = init_train_state(
+ model=model,
+ random_key=rng,
+ shape=(config.batch_size, config.image_size, config.image_size, 3),
+ learning_rate=config.learning_rate
+ )
+ print(type(state))
+ return (state,)
+
+
+@app.cell
+def _(config, jax, optax):
+ def cross_entropy_loss(*, logits, labels):
+ one_hot_encoded_labels = jax.nn.one_hot(
+ labels, num_classes=len(config.labels)
+ )
+ return optax.softmax_cross_entropy(
+ logits=logits, labels=one_hot_encoded_labels
+ ).mean()
+
+ return (cross_entropy_loss,)
+
+
+@app.cell
+def _(cross_entropy_loss, jnp):
+ def compute_metrics(*, logits, labels):
+ loss = cross_entropy_loss(logits=logits, labels=labels)
+ accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
+ metrics = {
+ 'loss': loss,
+ 'accuracy': accuracy,
+ }
+ return metrics
+
+ return (compute_metrics,)
+
+
+@app.cell
+def _(compute_metrics, cross_entropy_loss, jax, jnp, train_state):
+ @jax.jit
+ def train_step(
+ state: train_state.TrainState, batch: jnp.ndarray
+ ):
+ image, label = batch
+
+ def loss_fn(params):
+ logits = state.apply_fn({'params': params}, image)
+ loss = cross_entropy_loss(logits=logits, labels=label)
+ return loss, logits
+
+ gradient_fn = jax.value_and_grad(loss_fn, has_aux=True)
+ (_, logits), grads = gradient_fn(state.params)
+ state = state.apply_gradients(grads=grads)
+ metrics = compute_metrics(logits=logits, labels=label)
+ return state, metrics
+
+ return (train_step,)
+
+
+@app.cell
+def _(compute_metrics, jax):
+ @jax.jit
+ def eval_step(state, batch):
+ image, label = batch
+ logits = state.apply_fn({'params': state.params}, image)
+ return compute_metrics(logits=logits, labels=label)
+
+ return (eval_step,)
+
+
+@app.cell
+def _(from_bytes, jax, msgpack_serialize, np, os, to_state_dict, wandb):
+ def save_checkpoint(ckpt_path, state, epoch):
+ with open(ckpt_path, "wb") as outfile:
+ outfile.write(msgpack_serialize(to_state_dict(state)))
+ artifact = wandb.Artifact(
+ f'{wandb.run.name}-checkpoint', type='dataset'
+ )
+ artifact.add_file(ckpt_path)
+ wandb.log_artifact(artifact, aliases=["latest", f"epoch_{epoch}"])
+
+
+ def load_checkpoint(ckpt_file, state):
+ artifact = wandb.use_artifact(
+ f'{wandb.run.name}-checkpoint:latest'
+ )
+ artifact_dir = artifact.download()
+ ckpt_path = os.path.join(artifact_dir, ckpt_file)
+ with open(ckpt_path, "rb") as data_file:
+ byte_data = data_file.read()
+ return from_bytes(state, byte_data)
+
+
+ def accumulate_metrics(metrics):
+ metrics = jax.device_get(metrics)
+ return {
+ k: np.mean([metric[k] for metric in metrics])
+ for k in metrics[0]
+ }
+
+ return accumulate_metrics, save_checkpoint
+
+
+@app.cell
+def _(
+ accumulate_metrics,
+ eval_step,
+ save_checkpoint,
+ tfds,
+ tqdm,
+ train_state,
+ train_step,
+ wandb,
+):
+ def train_and_evaluate(
+ train_dataset,
+ eval_dataset,
+ num_train_batches,
+ num_eval_batches,
+ state: train_state.TrainState,
+ epochs: int,
+ ):
+ for epoch in tqdm(range(1, epochs + 1)):
+
+ best_eval_loss = 1e6
+
+ train_batch_metrics = []
+ train_datagen = iter(tfds.as_numpy(train_dataset))
+ for batch_idx in range(num_train_batches):
+ batch = next(train_datagen)
+ state, metrics = train_step(state, batch)
+ train_batch_metrics.append(metrics)
+
+ train_batch_metrics = accumulate_metrics(train_batch_metrics)
+ print(
+ 'TRAIN (%d/%d): Loss: %.4f, accuracy: %.2f' % (
+ epoch, epochs, train_batch_metrics['loss'],
+ train_batch_metrics['accuracy'] * 100
+ )
+ )
+
+ eval_batch_metrics = []
+ eval_datagen = iter(tfds.as_numpy(eval_dataset))
+ for batch_idx in range(num_eval_batches):
+ batch = next(eval_datagen)
+ metrics = eval_step(state, batch)
+ eval_batch_metrics.append(metrics)
+
+ eval_batch_metrics = accumulate_metrics(eval_batch_metrics)
+ print(
+ 'EVAL (%d/%d): Loss: %.4f, accuracy: %.2f\n' % (
+ epoch, epochs, eval_batch_metrics['loss'],
+ eval_batch_metrics['accuracy'] * 100
+ )
+ )
+
+ wandb.log({
+ "Train Loss": train_batch_metrics['loss'],
+ "Train Accuracy": train_batch_metrics['accuracy'],
+ "Validation Loss": eval_batch_metrics['loss'],
+ "Validation Accuracy": eval_batch_metrics['accuracy']
+ }, step=epoch)
+
+ if eval_batch_metrics['loss'] < best_eval_loss:
+ save_checkpoint("checkpoint.msgpack", state, epoch)
+
+ return state
+
+ return (train_and_evaluate,)
+
+
+@app.cell
+def _(
+ config,
+ num_train_batches,
+ num_val_batches,
+ state,
+ train_and_evaluate,
+ train_dataset,
+ val_dataset,
+):
+ state_1 = train_and_evaluate(train_dataset, val_dataset, num_train_batches, num_val_batches, state, epochs=config.epochs)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/jupyter-interactive-w-b-charts-inside-jupyter/jupyter_interactive_w_b_charts_inside_jupyter.py b/marimo/convert/jupyter-interactive-w-b-charts-inside-jupyter/jupyter_interactive_w_b_charts_inside_jupyter.py
new file mode 100644
index 00000000..a30566dc
--- /dev/null
+++ b/marimo/convert/jupyter-interactive-w-b-charts-inside-jupyter/jupyter_interactive_w_b_charts_inside_jupyter.py
@@ -0,0 +1,380 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Use W&B without leaving Jupyter
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Jupyter is the preferred development environment for many ML practitioners
+ because it supports rapid experimentation and
+ highly visual workflows (including creating charts).
+ Plus tools like Google Colab, Kaggle Kernels, and Paperspace Gradient
+ make it easy to share and collaborate on notebooks.
+
+ Quick experiments, visualization, and collaboration
+ are core values of W&B,
+ so we've made it easy to use W&B inside Jupyter.
+
+ In a nutshell, the steps are:
+
+ 1. Use one of two methods to get hold of a `Run`, `Sweep`, or `Report` object, depending on whether you're logging to a new experiment or analyzing an old one.
+ 2. `.display` the object to get a live dashboard beneath a cell.
+ 3. Interact with the dashboard: log new results, create charts, or review metadata.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's a (static) preview of one such dashboard:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Import, install, and log in
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Method 1: `display` and log to a live W&B `Run`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The result of the last line of each cell in a Jupyter notebook is "displayed" automatically.
+
+ Our W&B pages hook into this system:
+ they are rendered as an interactive window.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ First we need to kick the run off with
+ [`wandb.init`](https://docs.wandb.ai/guides/track/launch).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ run = wandb.init(project="jupyter-projo",
+ config={"batch_size": 128,
+ "learning_rate": 0.01,
+ "dataset": "CIFAR-100"})
+ return (run,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Then we create an interactive dashboard of the size we want for the run and display it.
+ """)
+ return
+
+
+@app.cell
+def _(run):
+ run.display(height=720)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Anything logged as part of this experiment (until you call `wandb.finish`)
+ will be added to that chart.
+
+ Run the cell below to watch the metrics stream in live!
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ import time
+
+ for ii in range(30):
+ wandb.log({"acc": 1 - 2 ** -ii, "loss": 2 ** -ii})
+ time.sleep(0.5)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > **Anything else you can do from the
+ [Run Page](https://docs.wandb.ai/ref/app/pages/run-page)
+ can be done here** --
+ [edit a chart](https://docs.wandb.ai/ref/app/pages/run-page#charts-tab),
+ create a shareable link to it,
+ and send it to collaborator;
+ review your [system metrics](https://docs.wandb.ai/ref/app/pages/run-page#system-tabs)
+ or the
+ [logs from the standard out](https://docs.wandb.ai/ref/app/pages/run-page#logs-tab)
+ or the
+ [datasets and models you've logged](https://docs.wandb.ai/ref/app/pages/run-page#artifacts-tab);
+ check the
+ [configuration metadata](https://docs.wandb.ai/ref/app/pages/run-page#overview-tab).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ `wandb` also prints a URL. That URL points to [the webpage](https://docs.wandb.ai/ref/app/pages/run-page)
+ where your run's results are stored -- nothing to worry about if your notebook crashes or your kernel dies, it's all there!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Finishing the run
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ When you are done with your experiment,
+ call `wandb.finish` to let us know there's nothing more to log.
+
+ We'll print out a handy summary and history of your run,
+ plus links to the webpages where all your run's information is stored.
+
+ > **Hot Tip!** If you turn on [code saving](https://docs.wandb.ai/ref/app/features/panels/code) in your W&B [settings](https://wandb.ai/settings),
+ we'll also save a copy of the notebook and its "session history": all the cells you ran, in order, in the state that you ran them in, with their outputs. Handy!
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ if wandb.run is not None:
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Method 2: `display` and analyze a finished W&B `Run`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Interaction with W&B dashboards for training runs
+ isn't limited to watching information come in live
+ from the comfort of a notebook interface.
+
+ All of the information you log to or create within W&B
+ is available in perpetuity and programmatically via the W&B
+ [Public API](https://docs.wandb.ai/guides/track/public-api-guide).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ api = wandb.Api()
+ return (api,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this example, we'll take a look at the training run for a chess piece detector
+ created using [YOLOv5](https://ultralytics.com/yolov5),
+ which includes a [W&B integration](https://docs.wandb.ai/guides/integrations/yolov5).
+
+ You can train your own with [this colab](http://wandb.me/yolo-chess).
+ """)
+ return
+
+
+@app.cell
+def _(api):
+ team, _project, _run_id = ('wandb', 'yolo-chess', '33fp7u8d')
+ run_1 = api.run(f'{team}/{_project}/{_run_id}')
+ run_1.display(height=1080)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # But it's not just about `Run`s
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Anything you can do in a W&B workspace can be done from inside Jupyter
+ if you have the URL for the workspace.
+
+ That means that, without leaving Jupyter, you can use W&B to:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Interactively analyze data in [Tables](https://docs.wandb.ai/guides/data-vis)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ And it doesn't have to be your own work -- it can be a `coworker`'s page as well.
+ """)
+ return
+
+
+@app.cell
+def _(api):
+ coworker, _project, _run_id = ('stacey', 'model_iterz', '10x1nnh2')
+ run_2 = api.run(f'{coworker}/{_project}/{_run_id}')
+ run_2.display(height=720)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Analyze the results of hyperparameter [Sweeps](https://docs.wandb.ai/guides/sweeps)
+ """)
+ return
+
+
+@app.cell
+def _(api):
+ _entity, _project, sweep_id = ('charlesfrye', 'mnist-sweeps', 'n60n6wv1')
+ sweep = api.sweep(f'{_entity}/{_project}/{sweep_id}')
+ sweep.display(height=1080) # you may need to zoom out to see the whole window!
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Share results in [Reports](https://docs.wandb.ai/guides/reports)
+ """)
+ return
+
+
+@app.cell
+def _():
+ _entity, _project = ('charlesfrye', 'mnist-sweeps')
+ report_name = 'Third-Pass-Trying-Different-Shapes--VmlldzoxNjY1NDk'
+ # magic command not supported in marimo; please file an issue to add support
+ # %wandb {entity}/{project}/reports/{report_name} -h 1024
+ url = f'https://wandb.ai/{_entity}/{_project}/reports/{report_name}'
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/kaolin-wisp-vqad/kaolin_wisp_vqad.py b/marimo/convert/kaolin-wisp-vqad/kaolin_wisp_vqad.py
new file mode 100644
index 00000000..5624fa80
--- /dev/null
+++ b/marimo/convert/kaolin-wisp-vqad/kaolin_wisp_vqad.py
@@ -0,0 +1,163 @@
+# /// script
+# dependencies = ["requirements-txt", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Kaolin-Wisp + WandB Demo 🪄🐝
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install Kaolin Core and Kaolin Wisp
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # Install OpenEXR
+ #! sudo apt-get update
+ subprocess.call(['sudo', 'apt-get', 'update'])
+ #! sudo apt-get install libopenexr-dev
+ subprocess.call(['sudo', 'apt-get', 'install', 'libopenexr-dev'])
+ subprocess.call(['git', 'clone', '--recursive', 'https://github.com/NVIDIAGameWorks/kaolin'])
+ # Install Kaolin
+ #! git clone --recursive https://github.com/NVIDIAGameWorks/kaolin
+ import os
+ os.chdir('kaolin')
+ subprocess.call(['python', 'setup.py', 'develop'])
+ #! python setup.py develop
+ subprocess.call(['python', '-c', 'import kaolin; print(kaolin.__version__)'])
+ #! python -c "import kaolin; print(kaolin.__version__)"
+ os.chdir('..')
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/NVIDIAGameWorks/kaolin-wisp'])
+ os.chdir('kaolin-wisp')
+ subprocess.call(['python', 'setup.py', 'develop'])
+ # Install Kaolin-Wisp
+ #! git clone --depth 1 https://github.com/NVIDIAGameWorks/kaolin-wisp
+ # packages added via marimo's package management: requirements.txt !pip install -q -r requirements.txt
+ # packages added via marimo's package management: wandb !pip install -q --upgrade wandb
+ os.chdir('..')
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Download Sample Data for a V8 Model Engine
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # Download Dataset
+ #! gdown https://drive.google.com/uc?id=18hY0DpX2bK-q9iY_cog5Q0ZI7YEjephE
+ subprocess.call(['gdown', 'https://drive.google.com/uc?id=18hY0DpX2bK-q9iY_cog5Q0ZI7YEjephE'])
+ #! unzip -q V8.zip
+ subprocess.call(['unzip', '-q', 'V8.zip'])
+ #! rm V8.zip
+ subprocess.call(['rm', 'V8.zip'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train VQAD
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ A great aspect of Kaolin Wisp is that it comes with the goodness of [Weights & Biases](https://wandb.ai/site) integrated with itself!!!
+
+ To track training and validation metrics, render 3D interactive plots, reproduce your configurations and results, and many more features in your Weights & Biases workspace just add the additional flag `--wandb_project ` when initializing the training script.
+
+ The complete list of features supported by Weights & Biases:
+
+ - Log training and validation metrics in real time.
+
+ - Log system metrics in real time.
+
+ - Log RGB, RGBA, Depth renderings etc. during training.
+
+ - Log interactive 360 degree renderings post training
+ in all levels of detail.
+
+ - Log model checkpoints as [Weights & Biases artifacts](https://wandb.ai/site/artifacts).
+
+ - Sync experiment configs for reproducibility.
+
+ - Host Tensorboard instance inside Weights & Biases run.
+
+ The full list of optional arguments related to logging on Weights & Biases include:
+
+ - `--wandb_project`: Name of Weights & Biases project
+
+ - `--wandb_run_name`: Name of Weights & Biases run [Optional]
+ - `--wandb_entity`: Name of Weights & Biases entity under which your project resides [Optional]
+
+ - `--wandb_viz_nerf_angles`: Number of angles in the 360 degree renderings [Optional, default set to 20]
+
+ - `--wandb_viz_nerf_distance`: Camera distance to visualize Scene from for 360 degree renderings on Weights & Biases [Optional, default set to 3]
+ """)
+ return
+
+
+@app.cell
+def _(os, subprocess):
+ os.chdir('kaolin-wisp')
+ #! WISP_HEADLESS=1 python3 app/main.py --config configs/vqad_nerf.yaml --dataset-path ../V8_/ --dataset-num-workers 4 --wandb_project "vector-quantized-auto-decoder" --wandb_run_name test-vqad-nerf/V8 --wandb_viz_nerf_distance 5
+ subprocess.call(['WISP_HEADLESS=1', 'python3', 'app/main.py', '--config', 'configs/vqad_nerf.yaml', '--dataset-path', '../V8_/', '--dataset-num-workers', '4', '--wandb_project', 'vector-quantized-auto-decoder', '--wandb_run_name', 'test-vqad-nerf/V8', '--wandb_viz_nerf_distance', '5'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If you wish to train using one of the numerous scenes from the [RTMV Dataset](http://www.cs.umd.edu/~mmeshry/projects/rtmv/), you can replace the gdown URL with one of the tar files from [here](https://drive.google.com/drive/folders/1cc5ArA16pEznMd92z7pwgD1Z4uBqafUN). You also need to change the `--dataset-path` paramter while training to the respective path of the model that you wish to train on.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-cosine-decay-using-keras/keras_cosine_decay_using_keras.py b/marimo/convert/keras-cosine-decay-using-keras/keras_cosine_decay_using_keras.py
new file mode 100644
index 00000000..fa7a8cb8
--- /dev/null
+++ b/marimo/convert/keras-cosine-decay-using-keras/keras_cosine_decay_using_keras.py
@@ -0,0 +1,202 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using Cosine Decay with Keras
+
+ This notebook demonstrates how to use the Cosine Decay learning rate schedule with Keras.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qq wandb
+ return
+
+
+@app.cell
+def _():
+ import tensorflow as tf
+ from tensorflow.keras import layers
+ from tensorflow.keras import models
+ import tensorflow_datasets as tfds
+
+ # Weights and Biases related imports
+ import wandb
+ from wandb.keras import WandbMetricsLogger
+
+ return WandbMetricsLogger, layers, models, tf, tfds, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ configs = dict(
+ num_classes = 10,
+ shuffle_buffer = 1024,
+ batch_size = 64,
+ image_size = 28,
+ image_channels = 1,
+ earlystopping_patience = 3,
+ learning_rate = 1e-3,
+ epochs = 10,
+ num_steps = 0.7,
+ )
+ return (configs,)
+
+
+@app.cell
+def _(configs, tf, tfds):
+ AUTOTUNE = tf.data.AUTOTUNE
+
+
+ def parse_data(example):
+ # Get image
+ image = example["image"]
+
+ # Get label
+ label = example["label"]
+ label = tf.one_hot(label, depth=configs["num_classes"])
+
+ return image, label
+
+
+ def get_dataloader(ds, configs, dataloader_type="train"):
+ dataloader = ds.map(parse_data, num_parallel_calls=AUTOTUNE)
+
+ if dataloader_type=="train":
+ dataloader = dataloader.shuffle(configs["shuffle_buffer"])
+
+ dataloader = (
+ dataloader
+ .batch(configs["batch_size"])
+ .prefetch(AUTOTUNE)
+ )
+
+ return dataloader
+
+ train_ds, valid_ds = tfds.load('fashion_mnist', split=['train', 'test'])
+
+ trainloader = get_dataloader(train_ds, configs)
+ validloader = get_dataloader(valid_ds, configs, dataloader_type="valid")
+ return trainloader, validloader
+
+
+@app.cell
+def _(configs, layers, models, tf):
+ def get_model(configs):
+ backbone = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet', include_top=False)
+ backbone.trainable = True
+
+ inputs = layers.Input(shape=(configs["image_size"], configs["image_size"], configs["image_channels"]))
+ resize = layers.Resizing(32, 32)(inputs)
+ neck = layers.Conv2D(3, (3,3), padding="same")(resize)
+ preprocess_input = tf.keras.applications.mobilenet.preprocess_input(neck)
+ x = backbone(preprocess_input)
+ x = layers.GlobalAveragePooling2D()(x)
+ outputs = layers.Dense(configs["num_classes"], activation="softmax")(x)
+
+ return models.Model(inputs=inputs, outputs=outputs)
+
+
+ tf.keras.backend.clear_session()
+ model = get_model(configs)
+ model.summary()
+ return (model,)
+
+
+@app.cell
+def _(configs, tf, trainloader):
+ # Learning Rate
+ total_steps = len(trainloader)*configs["epochs"]
+ decay_steps = total_steps * configs["num_steps"]
+
+ cosine_decay_scheduler = tf.keras.optimizers.schedules.CosineDecay(
+ initial_learning_rate = configs["learning_rate"],
+ decay_steps = decay_steps,
+ alpha=0.1
+ )
+ return (cosine_decay_scheduler,)
+
+
+@app.cell
+def _(cosine_decay_scheduler, model, tf):
+ model.compile(
+ optimizer = tf.keras.optimizers.Adam(cosine_decay_scheduler),
+ loss = "categorical_crossentropy",
+ metrics = ["accuracy"]
+ )
+ return
+
+
+@app.cell
+def _(WandbMetricsLogger, configs, model, trainloader, validloader, wandb):
+ # Initialize a W&B run
+ run = wandb.init(
+ project = "cosine_decay",
+ config = configs,
+ )
+
+ # Train your model
+ model.fit(
+ trainloader,
+ epochs = configs["epochs"],
+ validation_data = validloader,
+ callbacks = [
+ WandbMetricsLogger(log_freq=2),
+ ]
+ )
+ return (run,)
+
+
+@app.cell
+def _(model, validloader, wandb):
+ eval_loss, eval_acc = model.evaluate(validloader)
+
+ wandb.log({
+ "eval_loss": eval_loss,
+ "eval_acc": eval_acc
+ })
+ return
+
+
+@app.cell
+def _(run):
+ # Close the W&B run
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-dreambooth-inference/keras_dreambooth_inference.py b/marimo/convert/keras-dreambooth-inference/keras_dreambooth_inference.py
new file mode 100644
index 00000000..a528524b
--- /dev/null
+++ b/marimo/convert/keras-dreambooth-inference/keras_dreambooth_inference.py
@@ -0,0 +1,159 @@
+# /// script
+# dependencies = ["dreambooth-keras @ git+https://github.com/soumik12345/dreambooth-keras.git"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🧨 Dreambooth-Keras + WandB 🪄🐝
+
+ [](https://colab.research.google.com/github/soumik12345/dreambooth-keras/blob/main/notebooks/inference_wandb.ipynb)
+
+
+
+ This notebook shows how to perform inference with a DreamBooth fine-tuned Stable Diffusion model.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🌈 Install Dreambooth-Keras
+
+ We would use [soumik12345/dreambooth-keras](https://github.com/soumik12345/dreambooth-keras) which is a fork of [sayakpaul/dreambooth-keras](https://github.com/sayakpaul/dreambooth-keras) developed by [**Sayak Paul**](https://github.com/sayakpaul) and [**Chansung Park**](https://github.com/deep-diver).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: git+https://github.com/soumik12345/dreambooth-keras.git !pip install -q git+https://github.com/soumik12345/dreambooth-keras.git
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from PIL import Image
+ from dreambooth_keras.utils import load_model_from_wandb_artifact
+
+ return Image, load_model_from_wandb_artifact, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🐝 Initialize WandB run
+
+ We initialize a [Weights & Biases run](https://docs.wandb.ai/guides/runs) for storing generated images to a [Weights & Biases table](https://docs.wandb.ai/guides/data-vis).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project="dreambooth-keras", job_type="inference")
+
+ config = wandb.config
+ config.model_artifact_address = "geekyrakshit/dreambooth-keras/run_n5oakq7c_model:v0"
+ config.image_resolution = 512
+ config.num_diffusion_steps = 500
+ config.batch_size = 5
+ config.unique_id = "sks"
+ config.class_category = "monkey"
+ config.prompt = "a painting of sks monkey in the style of Michelangelo"
+ config.unconditional_guidance_scale = 15
+
+
+ wandb_table = wandb.Table(columns=[
+ "prompt", "images", "unique-id", "class-category","image-resolution", "num-diffusion-steps"
+ ])
+ return config, wandb_table
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🧑🎨 Perform Inference
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ First we load our model from Weights & Biases artifacts created using the [`dreambooth_keras.utils.DreamBoothCheckpointCallback`](https://github.com/soumik12345/dreambooth-keras/blob/main/dreambooth_keras/utils.py#L93) which automatically logs model checkpoints as [Weights & Biases artifacts](https://docs.wandb.ai/guides/data-and-model-versioning) at the end of each epoch during training. We load these checkpoint using the simple utility [`dreambooth_keras.utils.load_model_from_wandb_artifact`](https://github.com/soumik12345/dreambooth-keras/blob/main/dreambooth_keras/utils.py#L23).
+ """)
+ return
+
+
+@app.cell
+def _(config, load_model_from_wandb_artifact):
+ dreambooth_model = load_model_from_wandb_artifact(
+ artifact_address=config.model_artifact_address,
+ image_resolution=config.image_resolution
+ )
+ return (dreambooth_model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we perform inference on our *dreamboothed* stable-diffusion model.
+ """)
+ return
+
+
+@app.cell
+def _(config, dreambooth_model):
+ _images = dreambooth_model.text_to_image(config.prompt, batch_size=config.batch_size, num_steps=config.num_diffusion_steps, unconditional_guidance_scale=config.unconditional_guidance_scale)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next we log our images to a [Weights & Biases table](https://docs.wandb.ai/guides/data-vis) that not only makes ut easier to visualize but also easily accessible for future reference.
+ """)
+ return
+
+
+@app.cell
+def _(Image, config, wandb, wandb_table):
+ _images = [wandb.Image(Image.fromarray(image), caption=f'{i}: {config.prompt}') for i, image in enumerate(_images)]
+ wandb_table.add_data(config.prompt, _images, config.unique_id, config.class_category, config.image_resolution, config.num_diffusion_steps)
+ wandb.log({'Inference-Results': wandb_table})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-fine-tune-vision-transformer-using-kerascv/keras_fine_tune_vision_transformer_using_kerascv.py b/marimo/convert/keras-fine-tune-vision-transformer-using-kerascv/keras_fine_tune_vision_transformer_using_kerascv.py
new file mode 100644
index 00000000..b26806cb
--- /dev/null
+++ b/marimo/convert/keras-fine-tune-vision-transformer-using-kerascv/keras_fine_tune_vision_transformer_using_kerascv.py
@@ -0,0 +1,342 @@
+# /// script
+# dependencies = ["keras-cv", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Installations and Imports
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: keras-cv !pip install -qq keras-cv
+ # packages added via marimo's package management: wandb !pip install -qq wandb
+ return
+
+
+@app.cell
+def _():
+ import numpy as np
+ from argparse import Namespace
+
+ import tensorflow as tf
+ import tensorflow_datasets as tfds
+ from tensorflow.keras import layers
+ from tensorflow.keras import models
+
+ import keras_cv as kcv
+ from keras_cv.models import ViTTiny16
+ from keras_cv.layers import preprocessing
+
+ import wandb
+ from wandb.keras import WandbMetricsLogger
+ from wandb.keras import WandbEvalCallback
+
+ return (
+ Namespace,
+ ViTTiny16,
+ WandbEvalCallback,
+ WandbMetricsLogger,
+ np,
+ preprocessing,
+ tf,
+ tfds,
+ wandb,
+ )
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Hyperparameters
+ """)
+ return
+
+
+@app.cell
+def _(Namespace):
+ configs = Namespace(
+ learning_rate = 1e-4,
+ batch_size = 64,
+ num_epochs = 10,
+ image_size = 224,
+ num_classes = 120,
+ num_steps = 1.0,
+ )
+ return (configs,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Dataset and Dataloaders
+ """)
+ return
+
+
+@app.cell
+def _(configs, preprocessing, tf, tfds):
+ AUTOTUNE = tf.data.AUTOTUNE
+
+
+ def parse_data(example):
+ "Apply preprocessing to one data sample at a time."
+ image = example["image"]
+ image = tf.image.convert_image_dtype(image, tf.float32)
+ image = tf.image.resize(image, (configs.image_size, configs.image_size))
+
+ label = example["label"]
+ label = tf.one_hot(label, configs.num_classes)
+
+ return image, label
+
+
+ base_augmentations = tf.keras.Sequential(
+ [
+ tf.keras.layers.RandomFlip("horizontal"),
+ tf.keras.layers.RandomRotation(factor=0.02),
+ tf.keras.layers.RandomZoom(height_factor=0.2, width_factor=0.2),
+ ],
+ name="base_augmentation",
+ )
+
+ mixup = preprocessing.MixUp(alpha=0.8)
+
+
+ def apply_base_augmentations(images, labels):
+ images = base_augmentations(images)
+ return images, labels
+
+
+ ds_train, ds_test = tfds.load('stanford_dogs', split=['train', 'test'])
+
+ trainloader = (
+ ds_train
+ .map(parse_data, num_parallel_calls=AUTOTUNE)
+ .batch(configs.batch_size)
+ .map(apply_base_augmentations, num_parallel_calls=AUTOTUNE)
+ .map(lambda images, labels: mixup({"images": images, "labels": labels}), num_parallel_calls=AUTOTUNE)
+ .map(lambda x: (x["images"], x["labels"]), num_parallel_calls=AUTOTUNE)
+ .shuffle(1024)
+ .prefetch(AUTOTUNE)
+ )
+
+ testloader = (
+ ds_test
+ .map(parse_data, num_parallel_calls=AUTOTUNE)
+ .batch(configs.batch_size)
+ .prefetch(AUTOTUNE)
+ )
+ return testloader, trainloader
+
+
+@app.cell
+def _(ViTTiny16, configs, tf):
+ def get_model():
+ inputs = tf.keras.layers.Input(shape=(configs.image_size, configs.image_size, 3))
+
+ vit = ViTTiny16(
+ include_rescaling=False,
+ include_top=False,
+ name="ViTTiny32",
+ weights="imagenet",
+ input_tensor=inputs,
+ pooling="token_pooling",
+ activation=tf.keras.activations.gelu,
+ )
+
+ vit.trainable = True
+
+ outputs = tf.keras.layers.Dense(configs.num_classes, activation="softmax")(vit.output)
+ model = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+ return model
+
+ model = get_model()
+ model.summary()
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Compile the Model
+
+ We will use `CosineDecay` learning rate scheduler.
+ """)
+ return
+
+
+@app.cell
+def _(configs, model, tf, trainloader):
+ total_steps = len(trainloader)*configs.num_epochs
+ decay_steps = total_steps * configs.num_steps
+
+ cosine_decay_scheduler = tf.keras.optimizers.schedules.CosineDecay(
+ configs.learning_rate, decay_steps, alpha=0.1
+ )
+
+ model.compile(
+ optimizer=tf.keras.optimizers.Adam(learning_rate=cosine_decay_scheduler),
+ loss=tf.keras.losses.CategoricalCrossentropy(),
+ metrics=["accuracy"],
+ )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # [OPTIONAL] Model Prediction Visualization
+
+ We will build a custom Keras callback by subclassing `WandbEvalCallback` for model prediction visualization.
+ """)
+ return
+
+
+@app.cell
+def _(WandbEvalCallback, np, tf, wandb):
+ class WandbClfEvalCallback(WandbEvalCallback):
+ def __init__(
+ self, validloader, data_table_columns, pred_table_columns, num_samples=100
+ ):
+ super().__init__(data_table_columns, pred_table_columns)
+
+ self.val_data = validloader.unbatch().take(num_samples)
+
+ def add_ground_truth(self, logs=None):
+ for idx, (image, label) in enumerate(self.val_data):
+ self.data_table.add_data(
+ idx,
+ wandb.Image(image),
+ np.argmax(label, axis=-1)
+ )
+
+ def add_model_predictions(self, epoch, logs=None):
+ # Get predictions
+ preds = self._inference()
+ table_idxs = self.data_table_ref.get_index()
+
+ for idx in table_idxs:
+ pred = preds[idx]
+ self.pred_table.add_data(
+ epoch,
+ self.data_table_ref.data[idx][0],
+ self.data_table_ref.data[idx][1],
+ self.data_table_ref.data[idx][2],
+ pred
+ )
+
+ def _inference(self):
+ preds = []
+ for image, label in self.val_data:
+ pred = self.model(tf.expand_dims(image, axis=0))
+ argmax_pred = tf.argmax(pred, axis=-1).numpy()[0]
+ preds.append(argmax_pred)
+
+ return preds
+
+ return (WandbClfEvalCallback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train the model with W&B
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbClfEvalCallback,
+ WandbMetricsLogger,
+ configs,
+ model,
+ testloader,
+ trainloader,
+ wandb,
+):
+ # Initialize a W&B run
+ run = wandb.init(
+ project="keras_cv_vit",
+ save_code=False,
+ config=vars(configs),
+ )
+
+ # Fine-tune the model
+ model.fit(
+ trainloader,
+ epochs=configs.num_epochs,
+ validation_data=testloader,
+ callbacks=[
+ WandbMetricsLogger(log_freq=2),
+ WandbClfEvalCallback(
+ validloader = testloader,
+ data_table_columns = ["idx", "image", "label"],
+ pred_table_columns = ["epoch", "idx", "image", "label", "pred"]
+ )
+ ],
+ )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Model Evaluation
+ """)
+ return
+
+
+@app.cell
+def _(model, testloader, wandb):
+ eval_loss, eval_acc = model.evaluate(testloader)
+ wandb.log({
+ "eval_loss": eval_loss,
+ "eval_acc": eval_acc
+ })
+ return
+
+
+@app.cell
+def _(wandb):
+ # Close the W&B run
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-image-segmentation-with-keras/keras_image_segmentation_with_keras.py b/marimo/convert/keras-image-segmentation-with-keras/keras_image_segmentation_with_keras.py
new file mode 100644
index 00000000..b2b6b519
--- /dev/null
+++ b/marimo/convert/keras-image-segmentation-with-keras/keras_image_segmentation_with_keras.py
@@ -0,0 +1,332 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Imports and Setups
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qq wandb
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.keras import WandbMetricsLogger
+ from wandb.keras import WandbEvalCallback
+
+ return WandbEvalCallback, WandbMetricsLogger, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import tensorflow as tf
+ from tensorflow.keras import layers
+ from tensorflow.keras import models
+
+ import tensorflow_datasets as tfds
+
+ import os
+ import numpy as np
+ from argparse import Namespace
+ import matplotlib.pyplot as plt
+
+ return Namespace, layers, models, np, tf, tfds
+
+
+@app.cell
+def _(Namespace):
+ configs = Namespace(
+ img_size = 128,
+ batch_size = 32,
+ num_classes = 3,
+ )
+ configs
+ return (configs,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Dataloader
+
+ We will be using Oxford Pets Dataset which we can directly get from TensorFlow Datasets.
+ """)
+ return
+
+
+@app.cell
+def _(tfds):
+ train_ds, valid_ds = tfds.load('oxford_iiit_pet', split=["train", "test"])
+ return train_ds, valid_ds
+
+
+@app.cell
+def _(configs, tf, train_ds, valid_ds):
+ AUTOTUNE = tf.data.experimental.AUTOTUNE
+
+
+ def parse_data(example):
+ # Parse image
+ image = example["image"]
+ image = tf.image.convert_image_dtype(image, tf.float32)
+ image = tf.image.resize(image, size=(configs.img_size, configs.img_size))
+
+ # Parse mask
+ mask = example["segmentation_mask"] - 1 # ground truth labels are [1,2,3].
+ mask = tf.image.resize(mask, size=(configs.img_size, configs.img_size), method='nearest')
+ mask = tf.one_hot(tf.squeeze(mask, axis=-1), depth=configs.num_classes)
+
+ return image, mask
+
+ trainloader = (
+ train_ds
+ .shuffle(1024)
+ .map(parse_data, num_parallel_calls=AUTOTUNE)
+ .batch(configs.batch_size)
+ .prefetch(AUTOTUNE)
+ )
+
+ validloader = (
+ valid_ds
+ .map(parse_data, num_parallel_calls=AUTOTUNE)
+ .batch(configs.batch_size)
+ .prefetch(AUTOTUNE)
+ )
+ return trainloader, validloader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Model
+ """)
+ return
+
+
+@app.cell
+def _(configs, layers, models):
+ # ref: https://github.com/ayulockin/deepimageinpainting/blob/master/Image_Inpainting_Autoencoder_Decoder_v2_0.ipynb
+ class SegmentationModel:
+ '''
+ Build UNET based model for segmentation task.
+ '''
+ def prepare_model(self, OUTPUT_CHANNEL, input_size=(configs.img_size, configs.img_size, 3)):
+ inputs = layers.Input(input_size)
+
+ conv1, pool1 = self.__ConvBlock(32, (3,3), (2,2), 'relu', 'same', inputs)
+ conv2, pool2 = self.__ConvBlock(64, (3,3), (2,2), 'relu', 'same', pool1)
+ conv3, pool3 = self.__ConvBlock(128, (3,3), (2,2), 'relu', 'same', pool2)
+ conv4, pool4 = self.__ConvBlock(256, (3,3), (2,2), 'relu', 'same', pool3)
+
+ conv5, up6 = self.__UpConvBlock(512, 256, (3,3), (2,2), (2,2), 'relu', 'same', pool4, conv4)
+ conv6, up7 = self.__UpConvBlock(256, 128, (3,3), (2,2), (2,2), 'relu', 'same', up6, conv3)
+ conv7, up8 = self.__UpConvBlock(128, 64, (3,3), (2,2), (2,2), 'relu', 'same', up7, conv2)
+ conv8, up9 = self.__UpConvBlock(64, 32, (3,3), (2,2), (2,2), 'relu', 'same', up8, conv1)
+
+ conv9 = self.__ConvBlock(32, (3,3), (2,2), 'relu', 'same', up9, False)
+
+ outputs = layers.Conv2D(OUTPUT_CHANNEL, (3, 3), activation='softmax', padding='same')(conv9)
+
+ return models.Model(inputs=[inputs], outputs=[outputs])
+
+ def __ConvBlock(self, filters, kernel_size, pool_size, activation, padding, connecting_layer, pool_layer=True):
+ conv = layers.Conv2D(filters=filters, kernel_size=kernel_size, activation=activation, padding=padding)(connecting_layer)
+ conv = layers.Conv2D(filters=filters, kernel_size=kernel_size, activation=activation, padding=padding)(conv)
+ if pool_layer:
+ pool = layers.MaxPooling2D(pool_size)(conv)
+ return conv, pool
+ else:
+ return conv
+
+ def __UpConvBlock(self, filters, up_filters, kernel_size, up_kernel, up_stride, activation, padding, connecting_layer, shared_layer):
+ conv = layers.Conv2D(filters=filters, kernel_size=kernel_size, activation=activation, padding=padding)(connecting_layer)
+ conv = layers.Conv2D(filters=filters, kernel_size=kernel_size, activation=activation, padding=padding)(conv)
+ up = layers.Conv2DTranspose(filters=up_filters, kernel_size=up_kernel, strides=up_stride, padding=padding)(conv)
+ up = layers.concatenate([up, shared_layer], axis=3)
+
+ return conv, up
+
+ return (SegmentationModel,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Initialize Model and Compile
+ """)
+ return
+
+
+@app.cell
+def _(SegmentationModel, configs, tf):
+ # output channel is 3 because we have three classes in our mask
+ tf.keras.backend.clear_session()
+ model = SegmentationModel().prepare_model(configs.num_classes)
+
+ model.compile(
+ optimizer="adam",
+ loss="categorical_crossentropy",
+ )
+
+ model.summary()
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Callback
+ """)
+ return
+
+
+@app.cell
+def _():
+ segmentation_classes = ['pet', 'pet_outline', 'background']
+
+ # returns a dictionary of labels
+ def labels():
+ l = {}
+ for i, label in enumerate(segmentation_classes):
+ l[i] = label
+ return l
+
+ return (labels,)
+
+
+@app.cell
+def _(WandbEvalCallback, labels, np, tf, wandb):
+ class WandbSemanticLogger(WandbEvalCallback):
+ def __init__(
+ self,
+ validloader,
+ data_table_columns=["index", "image"],
+ pred_table_columns=["epoch", "index", "image", "prediction"],
+ num_samples=100,
+ ):
+ super().__init__(
+ data_table_columns,
+ pred_table_columns,
+ )
+
+ self.val_data = validloader.unbatch().take(num_samples)
+
+ def add_ground_truth(self, logs):
+ for idx, (image, mask) in enumerate(self.val_data):
+ self.data_table.add_data(
+ idx,
+ self._prepare_wandb_mask(
+ image.numpy(),
+ np.argmax(mask.numpy(), axis=-1),
+ "ground_truth"
+ )
+ )
+
+ def add_model_predictions(self, epoch, logs):
+ data_table_ref = self.data_table_ref
+ table_idxs = data_table_ref.get_index()
+
+ for idx, (image, mask) in enumerate(self.val_data):
+ prediction = self.model.predict(tf.expand_dims(image, axis=0), verbose=0)
+ prediction = np.argmax(tf.squeeze(prediction, axis=0).numpy(), axis=-1)
+
+ self.pred_table.add_data(
+ epoch,
+ data_table_ref.data[idx][0],
+ self._prepare_wandb_mask(
+ data_table_ref.data[idx][1],
+ np.argmax(mask.numpy(), axis=-1),
+ "ground_truth"
+ ),
+ self._prepare_wandb_mask(
+ data_table_ref.data[idx][1],
+ prediction,
+ "prediction"
+ )
+ )
+
+ def _prepare_wandb_mask(self, image, mask, mask_type):
+ return wandb.Image(
+ image,
+ masks = {
+ "ground_truth": {
+ "mask_data": mask,
+ "class_labels": labels()
+ }})
+
+ return (WandbSemanticLogger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ WandbSemanticLogger,
+ configs,
+ model,
+ trainloader,
+ validloader,
+ wandb,
+):
+ run = wandb.init(project='image-segmentation', config=configs)
+
+ _ = model.fit(
+ trainloader,
+ epochs=10,
+ validation_data=validloader,
+ callbacks=[
+ WandbMetricsLogger(log_freq=2),
+ WandbSemanticLogger(validloader)
+ ]
+ )
+
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-keras-core-monai-medmnist-keras/keras_keras_core_monai_medmnist_keras.py b/marimo/convert/keras-keras-core-monai-medmnist-keras/keras_keras_core_monai_medmnist_keras.py
new file mode 100644
index 00000000..d034e8f3
--- /dev/null
+++ b/marimo/convert/keras-keras-core-monai-medmnist-keras/keras_keras_core_monai_medmnist_keras.py
@@ -0,0 +1,460 @@
+# /// script
+# dependencies = ["monai-weekly", "namex", "wandb-addons @ git+https://github.com/soumik12345/wandb-addons"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🩺 Medical Image Classification Tutorial using MonAI and Keras
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/keras/keras_core/monai_medmnist_keras.ipynb)
+
+ This notebook demonstrates
+ - an end-to-end training using [MonAI](https://github.com/Project-MONAI/MONAI) and [KerasCore](https://github.com/keras-team/keras-core).
+ - how we can use the backend-agnostic Keras callbacks for [Weights & Biases](https://wandb.ai/site) to manage and track our experiment.
+
+ Original Notebook: https://github.com/Project-MONAI/tutorials/blob/main/2d_classification/mednist_tutorial.ipynb
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installing and Importing the Dependencies
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ - We install the `main` branch of [KerasCore](https://github.com/keras-team/keras-core), this lets us use the latest feature merged in KerasCore.
+ - We install [monai](https://github.com/Project-MONAI/MONAI), a PyTorch-based, open-source framework for deep learning in healthcare imaging, part of the PyTorch Ecosystem.
+ - We also install [wandb-addons](https://github.com/soumik12345/wandb-addons), a library that hosts the backend-agnostic callbacks compatible with KerasCore
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # install the `main` branch of KerasCore
+ # packages added via marimo's package management: namex !pip install -qq namex
+ #! apt install python3.10-venv
+ subprocess.call(['apt', 'install', 'python3.10-venv'])
+ #! git clone --depth 1 https://github.com/soumik12345/keras-core.git && cd keras-core && python pip_build.py --install
+ subprocess.call(['pip_build.py', '--install'])
+
+ # install monai and wandb-addons
+ # packages added via marimo's package management: git+https://github.com/soumik12345/wandb-addons !pip install -qq git+https://github.com/soumik12345/wandb-addons
+ # packages added via marimo's package management: monai-weekly[pillow, tqdm] !pip install -q "monai-weekly[pillow, tqdm]"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We specify the Keras backend to be using `torch` by explicitly specifying the environment variable `KERAS_BACKEND`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.environ["KERAS_BACKEND"] = "torch"
+
+ import shutil
+ import tempfile
+ import matplotlib.pyplot as plt
+ import PIL
+ import torch
+ import numpy as np
+ from sklearn.metrics import classification_report
+
+ import keras_core as keras
+ from keras_core.utils import TorchModuleWrapper
+
+ from monai.apps import download_and_extract
+ from monai.config import print_config
+ from monai.data import decollate_batch, DataLoader
+ from monai.metrics import ROCAUCMetric
+ from monai.networks.nets import DenseNet121
+ from monai.transforms import (
+ Activations,
+ EnsureChannelFirst,
+ AsDiscrete,
+ Compose,
+ LoadImage,
+ RandFlip,
+ RandRotate,
+ RandZoom,
+ ScaleIntensity,
+ )
+ from monai.utils import set_determinism
+
+ import wandb
+ from wandb_addons.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+ return (
+ Activations,
+ AsDiscrete,
+ Compose,
+ DataLoader,
+ DenseNet121,
+ EnsureChannelFirst,
+ LoadImage,
+ PIL,
+ RandFlip,
+ RandRotate,
+ RandZoom,
+ ScaleIntensity,
+ TorchModuleWrapper,
+ WandbMetricsLogger,
+ download_and_extract,
+ keras,
+ np,
+ os,
+ plt,
+ tempfile,
+ torch,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We initialize a [wandb run](https://docs.wandb.ai/guides/runs) and set the configs for the experiment.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project="keras-torch")
+
+ config = wandb.config
+ config.batch_size = 128
+ config.num_epochs = 1
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup data directory
+
+ You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.
+ This allows you to save results and reuse downloads.
+ If not specified a temporary directory will be used.
+ """)
+ return
+
+
+@app.cell
+def _(os, tempfile):
+ directory = os.environ.get("MONAI_DATA_DIRECTORY")
+ root_dir = tempfile.mkdtemp() if directory is None else directory
+ print(root_dir)
+ return (root_dir,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Download dataset
+
+ The MedNIST dataset was gathered from several sets from [TCIA](https://wiki.cancerimagingarchive.net/display/Public/Data+Usage+Policies+and+Restrictions),
+ [the RSNA Bone Age Challenge](http://rsnachallenges.cloudapp.net/competitions/4),
+ and [the NIH Chest X-ray dataset](https://cloud.google.com/healthcare/docs/resources/public-datasets/nih-chest).
+
+ The dataset is kindly made available by [Dr. Bradley J. Erickson M.D., Ph.D.](https://www.mayo.edu/research/labs/radiology-informatics/overview) (Department of Radiology, Mayo Clinic)
+ under the Creative Commons [CC BY-SA 4.0 license](https://creativecommons.org/licenses/by-sa/4.0/).
+
+ If you use the MedNIST dataset, please acknowledge the source.
+ """)
+ return
+
+
+@app.cell
+def _(download_and_extract, os, root_dir):
+ resource = "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/MedNIST.tar.gz"
+ md5 = "0bc7306e7427e00ad1c5526a6677552d"
+
+ compressed_file = os.path.join(root_dir, "MedNIST.tar.gz")
+ data_dir = os.path.join(root_dir, "MedNIST")
+ if not os.path.exists(data_dir):
+ download_and_extract(resource, compressed_file, root_dir, md5)
+ return (data_dir,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Read image filenames from the dataset folders
+
+ First of all, check the dataset files and show some statistics.
+ There are 6 folders in the dataset: Hand, AbdomenCT, CXR, ChestCT, BreastMRI, HeadCT,
+ which should be used as the labels to train our classification model.
+ """)
+ return
+
+
+@app.cell
+def _(PIL, data_dir, os):
+ class_names = sorted((x for x in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, x))))
+ num_class = len(class_names)
+ image_files = [[os.path.join(data_dir, class_names[_i], x) for x in os.listdir(os.path.join(data_dir, class_names[_i]))] for _i in range(num_class)]
+ num_each = [len(image_files[_i]) for _i in range(num_class)]
+ image_files_list = []
+ image_class = []
+ for _i in range(num_class):
+ image_files_list.extend(image_files[_i])
+ image_class.extend([_i] * num_each[_i])
+ num_total = len(image_class)
+ image_width, image_height = PIL.Image.open(image_files_list[0]).size
+ print(f'Total image count: {num_total}')
+ print(f'Image dimensions: {image_width} x {image_height}')
+ print(f'Label names: {class_names}')
+ print(f'Label counts: {num_each}')
+ return class_names, image_class, image_files_list, num_class, num_total
+
+
+@app.cell
+def _(PIL, class_names, image_class, image_files_list, np, num_total, plt):
+ plt.subplots(3, 3, figsize=(8, 8))
+ for _i, k in enumerate(np.random.randint(num_total, size=9)):
+ im = PIL.Image.open(image_files_list[k])
+ arr = np.array(im)
+ plt.subplot(3, 3, _i + 1)
+ plt.xlabel(class_names[image_class[k]])
+ plt.imshow(arr, cmap='gray', vmin=0, vmax=255)
+ plt.tight_layout()
+ plt.show()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Prepare training, validation and test data lists
+
+ Randomly select 10% of the dataset as validation and 10% as test.
+ """)
+ return
+
+
+@app.cell
+def _(image_class, image_files_list, np):
+ val_frac = 0.1
+ test_frac = 0.1
+ length = len(image_files_list)
+ indices = np.arange(length)
+ np.random.shuffle(indices)
+ test_split = int(test_frac * length)
+ val_split = int(val_frac * length) + test_split
+ test_indices = indices[:test_split]
+ val_indices = indices[test_split:val_split]
+ train_indices = indices[val_split:]
+ train_x = [image_files_list[_i] for _i in train_indices]
+ train_y = [image_class[_i] for _i in train_indices]
+ val_x = [image_files_list[_i] for _i in val_indices]
+ val_y = [image_class[_i] for _i in val_indices]
+ test_x = [image_files_list[_i] for _i in test_indices]
+ test_y = [image_class[_i] for _i in test_indices]
+ print(f'Training count: {len(train_x)}, Validation count: {len(val_x)}, Test count: {len(test_x)}')
+ return test_x, test_y, train_x, train_y, val_x, val_y
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Define MONAI transforms, Dataset and Dataloader to pre-process data
+ """)
+ return
+
+
+@app.cell
+def _(
+ Activations,
+ AsDiscrete,
+ Compose,
+ EnsureChannelFirst,
+ LoadImage,
+ RandFlip,
+ RandRotate,
+ RandZoom,
+ ScaleIntensity,
+ np,
+ num_class,
+):
+ train_transforms = Compose(
+ [
+ LoadImage(image_only=True),
+ EnsureChannelFirst(),
+ ScaleIntensity(),
+ RandRotate(range_x=np.pi / 12, prob=0.5, keep_size=True),
+ RandFlip(spatial_axis=0, prob=0.5),
+ RandZoom(min_zoom=0.9, max_zoom=1.1, prob=0.5),
+ ]
+ )
+
+ val_transforms = Compose([LoadImage(image_only=True), EnsureChannelFirst(), ScaleIntensity()])
+
+ y_pred_trans = Compose([Activations(softmax=True)])
+ y_trans = Compose([AsDiscrete(to_onehot=num_class)])
+ return train_transforms, val_transforms
+
+
+@app.cell
+def _(
+ DataLoader,
+ config,
+ test_x,
+ test_y,
+ torch,
+ train_transforms,
+ train_x,
+ train_y,
+ val_transforms,
+ val_x,
+ val_y,
+):
+ class MedNISTDataset(torch.utils.data.Dataset):
+ def __init__(self, image_files, labels, transforms):
+ self.image_files = image_files
+ self.labels = labels
+ self.transforms = transforms
+
+ def __len__(self):
+ return len(self.image_files)
+
+ def __getitem__(self, index):
+ return self.transforms(self.image_files[index]), self.labels[index]
+
+
+ train_ds = MedNISTDataset(train_x, train_y, train_transforms)
+ train_loader = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True, num_workers=2)
+
+ val_ds = MedNISTDataset(val_x, val_y, val_transforms)
+ val_loader = DataLoader(val_ds, batch_size=config.batch_size, num_workers=2)
+
+ test_ds = MedNISTDataset(test_x, test_y, val_transforms)
+ test_loader = DataLoader(test_ds, batch_size=config.batch_size, num_workers=2)
+ return train_loader, val_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We typically define a model in PyTorch using [`torch.nn.Module`s](https://pytorch.org/docs/stable/notes/modules.html) which act as the building blocks of stateful computation. Even though Keras supports PyTorch as a backend, it does not mean that we can nest torch modules inside a [`keras_core.Model`](https://keras.io/keras_core/api/models/), because trainable variables inside a Keras Model is tracked exclusively via [Keras Layers](https://keras.io/keras_core/api/layers/).
+
+ KerasCore provides us with a feature called `TorchModuleWrapper` which enables us to do exactly this. The `TorchModuleWrapper` is a Keras Layer that accepts a torch module and tracks its trainable variables, essentially converting the torch module into a Keras Layer. This enables us to put any torch modules inside a Keras Model and train them with a single `model.fit()`!
+
+ The idea of the `TorchModuleWrapper` was proposed by Keras' creator [François Chollet](https://github.com/fchollet) on [this issue thread](https://github.com/keras-team/keras-core/issues/604).
+ """)
+ return
+
+
+@app.cell
+def _(DenseNet121, TorchModuleWrapper, keras, num_class, torch, train_loader):
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ inputs = keras.Input(shape=(1, 64, 64))
+ outputs = TorchModuleWrapper(
+ DenseNet121(
+ spatial_dims=2, in_channels=1, out_channels=num_class
+ )
+ )(inputs)
+ model = keras.Model(inputs, outputs)
+
+ # model = MedMnistModel()
+ model(next(iter(train_loader))[0].to(device)).shape
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Note:** It is actually possible to use torch modules inside a Keras Model without having to explicitly have them wrapped with the `TorchModuleWrapper` as evident by [this tweet](https://twitter.com/fchollet/status/1697381832164290754) from François Chollet. However, this doesn't seem to work at the point of time this example was created, as reported in [this issue](https://github.com/keras-team/keras-core/issues/834).
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ config,
+ keras,
+ model,
+ train_loader,
+ val_loader,
+ wandb,
+):
+ # Compile the model
+ model.compile(
+ loss="sparse_categorical_crossentropy",
+ optimizer=keras.optimizers.Adam(1e-5),
+ metrics=["accuracy"],
+ )
+
+ # Define the backend-agnostic WandB callbacks for KerasCore
+ callbacks = [
+ # Track experiment metrics
+ WandbMetricsLogger(log_freq="batch")
+ ]
+
+ # Train the model by calling model.fit
+ model.fit(
+ train_loader,
+ validation_data=val_loader,
+ epochs=config.num_epochs,
+ callbacks=callbacks,
+ )
+
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-keras-core-timm-keras/keras_keras_core_timm_keras.py b/marimo/convert/keras-keras-core-timm-keras/keras_keras_core_timm_keras.py
new file mode 100644
index 00000000..65963ca0
--- /dev/null
+++ b/marimo/convert/keras-keras-core-timm-keras/keras_keras_core_timm_keras.py
@@ -0,0 +1,416 @@
+# /// script
+# dependencies = ["namex", "wandb-addons @ git+https://github.com/soumik12345/wandb-addons"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥 Fine-tune a [Timm](https://huggingface.co/docs/timm/index) Model with Keras and WandB 🦄
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/keras/keras_core/timm_keras.ipynb)
+
+ This notebook demonstrates
+ - how we can fine-tune a pre-trained model from timm using [KerasCore](https://github.com/keras-team/keras-core).
+ - how we can use the backend-agnostic Keras callbacks for [Weights & Biases](https://wandb.ai/site) to manage and track our experiment.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installing and Importing the Dependencies
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ - We install the `main` branch of [KerasCore](https://github.com/keras-team/keras-core), this lets us use the latest feature merged in KerasCore.
+ - We install [timm](https://huggingface.co/docs/timm/index), a library containing SOTA computer vision models, layers, utilities, optimizers, schedulers, data-loaders, augmentations, and training/evaluation scripts.
+ - We also install [wandb-addons](https://github.com/soumik12345/wandb-addons), a library that hosts the backend-agnostic callbacks compatible with KerasCore
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # install the `main` branch of KerasCore
+ # packages added via marimo's package management: namex !pip install -qq namex
+ #! apt install python3.10-venv
+ subprocess.call(['apt', 'install', 'python3.10-venv'])
+ #! git clone --depth 1 https://github.com/soumik12345/keras-core.git && cd keras-core && python pip_build.py --install
+ subprocess.call(['pip_build.py', '--install'])
+
+ # install timm and wandb-addons
+ # packages added via marimo's package management: git+https://github.com/soumik12345/wandb-addons !pip install -qq git+https://github.com/soumik12345/wandb-addons
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We specify the Keras backend to be using `torch` by explicitly specifying the environment variable `KERAS_BACKEND`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.environ["KERAS_BACKEND"] = "torch"
+
+ import numpy as np
+ import matplotlib.pyplot as plt
+
+ import torch
+ import torch.nn as nn
+ import torch.nn.functional as F
+
+ import timm
+ from timm.data import resolve_data_config
+
+ import torchvision
+ from torchvision import datasets, models, transforms
+ from torchvision.transforms.functional import InterpolationMode
+
+ import wandb
+ from wandb_addons.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+ return (
+ InterpolationMode,
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ datasets,
+ np,
+ os,
+ plt,
+ resolve_data_config,
+ timm,
+ torch,
+ torchvision,
+ transforms,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We initialize a [wandb run](https://docs.wandb.ai/guides/runs) and set the configs for the experiment.
+ """)
+ return
+
+
+@app.cell
+def _(resolve_data_config, wandb):
+ wandb.init(project="keras-torch")
+
+ config = wandb.config
+ config.model_name = "xception41"
+ config.freeze_backbone = False
+ config.preprocess_config = resolve_data_config({}, model=config.model_name)
+ config.dropout_rate = 0.5
+ config.batch_size = 4
+ config.num_epochs = 25
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## A PyTorch-based Input Pipeline
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will be using the [ImageNette](https://github.com/fastai/imagenette) dataset for this experiment. Imagenette is a subset of 10 easily classified classes from [Imagenet](https://www.image-net.org/) (tench, English springer, cassette player, chain saw, church, French horn, garbage truck, gas pump, golf ball, parachute).
+
+ First, let's download this dataset.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz -P imagenette
+ subprocess.call(['wget', 'https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz', '-P', 'imagenette'])
+ #! tar zxf imagenette/imagenette2-320.tgz -C imagenette
+ subprocess.call(['tar', 'zxf', 'imagenette/imagenette2-320.tgz', '-C', 'imagenette'])
+ #! gzip -d imagenette/imagenette2-320.tgz
+ subprocess.call(['gzip', '-d', 'imagenette/imagenette2-320.tgz'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we create our standard torch-based data loading pipeline.
+ """)
+ return
+
+
+@app.cell
+def _(InterpolationMode, config, datasets, os, torch, transforms):
+ # Define pre-processing and augmentation transforms for the train and validation sets
+ data_transforms = {
+ 'train': transforms.Compose([
+ transforms.RandomResizedCrop(
+ size=config.preprocess_config["input_size"][1],
+ interpolation=InterpolationMode.BICUBIC,
+ ),
+ transforms.RandomHorizontalFlip(),
+ transforms.ToTensor(),
+ transforms.Normalize(
+ config.preprocess_config["mean"],
+ config.preprocess_config["std"]
+ )
+ ]),
+ 'val': transforms.Compose([
+ transforms.Resize(256),
+ transforms.CenterCrop(config.preprocess_config["input_size"][1]),
+ transforms.ToTensor(),
+ transforms.Normalize(
+ config.preprocess_config["mean"],
+ config.preprocess_config["std"]
+ )
+ ]),
+ }
+
+ # Define the train and validation datasets
+ data_dir = 'imagenette/imagenette2-320'
+ image_datasets = {
+ x: datasets.ImageFolder(
+ os.path.join(data_dir, x), data_transforms[x]
+ )
+ for x in ['train', 'val']
+ }
+
+ # Define the torch dataloaders corresponding to the train and validation dataset
+ dataloaders = {
+ x: torch.utils.data.DataLoader(
+ image_datasets[x],
+ batch_size=config.batch_size,
+ shuffle=True,
+ num_workers=4
+ )
+ for x in ['train', 'val']
+ }
+ dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
+ class_names = image_datasets['train'].classes
+
+ # Specify the global device
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ return class_names, dataloaders, device
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's take a look at a few of the samples.
+ """)
+ return
+
+
+@app.cell
+def _(class_names, config, dataloaders, np, plt, torchvision):
+ def imshow(inp, title=None):
+ """Display image for Tensor."""
+ inp = inp.numpy().transpose((1, 2, 0))
+ mean = np.array(config.preprocess_config["mean"])
+ std = np.array(config.preprocess_config["std"])
+ inp = std * inp + mean
+ inp = np.clip(inp, 0, 1)
+ plt.imshow(inp)
+ if title is not None:
+ plt.title(title)
+ plt.pause(0.001)
+
+
+ # Get a batch of training data
+ inputs, classes = next(iter(dataloaders['train']))
+ print(inputs.shape, classes.shape)
+
+ # Make a grid from batch
+ out = torchvision.utils.make_grid(inputs)
+
+ imshow(out, title=[class_names[x] for x in classes])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Creating and Training our Classifier
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We typically define a model in PyTorch using [`torch.nn.Module`s](https://pytorch.org/docs/stable/notes/modules.html) which act as the building blocks of stateful computation. Even though Keras supports PyTorch as a backend, it does not mean that we can nest torch modules inside a [`keras_core.Model`](https://keras.io/keras_core/api/models/), because trainable variables inside a Keras Model is tracked exclusively via [Keras Layers](https://keras.io/keras_core/api/layers/).
+
+ KerasCore provides us with a feature called `TorchModuleWrapper` which enables us to do exactly this. The `TorchModuleWrapper` is a Keras Layer that accepts a torch module and tracks its trainable variables, essentially converting the torch module into a Keras Layer. This enables us to put any torch modules inside a Keras Model and train them with a single `model.fit()`!
+
+ The idea of the `TorchModuleWrapper` was proposed by Keras' creator [François Chollet](https://github.com/fchollet) on [this issue thread](https://github.com/keras-team/keras-core/issues/604).
+ """)
+ return
+
+
+@app.cell
+def _(timm):
+ import keras_core as keras
+ from keras_core.utils import TorchModuleWrapper
+
+
+ class TimmClassifier(keras.Model):
+
+ def __init__(self, model_name, freeze_backbone, dropout_rate, num_classes, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # Define the pre-trained module from timm
+ self.backbone = TorchModuleWrapper(
+ timm.create_model(model_name, pretrained=True)
+ )
+ self.backbone.trainable = not freeze_backbone
+
+ # Build the classification head using keras layers
+ self.global_average_pooling = keras.layers.GlobalAveragePooling2D()
+ self.dropout = keras.layers.Dropout(dropout_rate)
+ self.classification_head = keras.layers.Dense(num_classes)
+
+ def call(self, inputs):
+ # We get the unpooled features from the timm backbone by calling `forward_features`
+ # on the torch module corresponding to the backbone.
+ x = self.backbone.module.forward_features(inputs)
+ x = self.global_average_pooling(x)
+ x = self.dropout(x)
+ x = self.classification_head(x)
+ return keras.activations.softmax(x, axis=1)
+
+ return TimmClassifier, keras
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Note:** It is actually possible to use torch modules inside a Keras Model without having to explicitly have them wrapped with the `TorchModuleWrapper` as evident by [this tweet](https://twitter.com/fchollet/status/1697381832164290754) from François Chollet. However, this doesn't seem to work at the point of time this example was created, as reported in [this issue](https://github.com/keras-team/keras-core/issues/834).
+ """)
+ return
+
+
+@app.cell
+def _(TimmClassifier, class_names, config, device, torch):
+ # Now, we define the model and pass a random tensor to check the output shape
+ model = TimmClassifier(
+ model_name=config.model_name,
+ freeze_backbone=config.freeze_backbone,
+ dropout_rate=config.dropout_rate,
+ num_classes=len(class_names)
+ )
+ model(torch.ones(1, *config.preprocess_config["input_size"]).to(device)).shape
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, in standard Keras fashion, all we need to do is compile the model and call `model.fit()`!
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ config,
+ dataloaders,
+ keras,
+ model,
+):
+ # Create exponential decay learning rate scheduler
+ decay_steps = config.num_epochs * len(dataloaders["train"]) // config.batch_size
+ lr_scheduler = keras.optimizers.schedules.ExponentialDecay(
+ initial_learning_rate=1e-3, decay_steps=decay_steps, decay_rate=0.1,
+ )
+
+ # Compile the model
+ model.compile(
+ loss="sparse_categorical_crossentropy",
+ optimizer=keras.optimizers.Adam(lr_scheduler),
+ metrics=["accuracy"],
+ )
+
+ # Define the backend-agnostic WandB callbacks for KerasCore
+ callbacks = [
+ # Track experiment metrics
+ WandbMetricsLogger(log_freq="batch"),
+ # Track and version model checkpoints
+ WandbModelCheckpoint("model.keras")
+ ]
+
+ # Train the model by calling model.fit
+ model.fit(
+ dataloaders["train"],
+ validation_data=dataloaders["val"],
+ epochs=config.num_epochs,
+ callbacks=callbacks,
+ )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In order to know more about the backend-agnostic Keras callbacks for Weights & Biases, check out the [docs for wandb-addons](https://geekyrakshit.dev/wandb-addons/keras/keras_core/).
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-keras-core-torchvision-keras/keras_keras_core_torchvision_keras.py b/marimo/convert/keras-keras-core-torchvision-keras/keras_keras_core_torchvision_keras.py
new file mode 100644
index 00000000..12ab1e30
--- /dev/null
+++ b/marimo/convert/keras-keras-core-torchvision-keras/keras_keras_core_torchvision_keras.py
@@ -0,0 +1,441 @@
+# /// script
+# dependencies = ["namex", "wandb-addons @ git+https://github.com/soumik12345/wandb-addons"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # 🔥 Fine-tune a TorchVision Model with Keras and WandB 🦄
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/keras/keras_core/torchvision_keras.ipynb)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Introduction
+
+ [TorchVision](https://pytorch.org/vision/stable/index.html) is a library part of the [PyTorch](http://pytorch.org/) project that consists of popular datasets, model architectures, and common image transformations for computer vision. This example demonstrates how we can perform transfer learning for image classification using a pre-trained backbone model from TorchVision on the [Imagenette dataset](https://github.com/fastai/imagenette) using KerasCore. We will also demonstrate the compatibility of KerasCore with an input system consisting of [Torch Datasets and Dataloaders](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html).
+
+ ### References:
+
+ - [Customizing what happens in `fit()` with PyTorch](https://keras.io/keras_core/guides/custom_train_step_in_torch/)
+ - [PyTorch Datasets and Dataloaders](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)
+ - [Transfer learning for Computer Vision using PyTorch](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html)
+
+ ## Setup
+
+ - We install the `main` branch of [KerasCore](https://github.com/keras-team/keras-core), this lets us use the latest feature merged in KerasCore.
+ - We also install [wandb-addons](https://github.com/soumik12345/wandb-addons), a library that hosts the backend-agnostic callbacks compatible with KerasCore
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # install the `main` branch of KerasCore
+ # packages added via marimo's package management: namex !pip install -qq namex
+ #! apt install python3.10-venv
+ subprocess.call(['apt', 'install', 'python3.10-venv'])
+ #! git clone --depth 1 https://github.com/soumik12345/keras-core.git && cd keras-core && python pip_build.py --install
+ subprocess.call(['pip_build.py', '--install'])
+
+ # install wandb-addons
+ # packages added via marimo's package management: git+https://github.com/soumik12345/wandb-addons !pip install -qq git+https://github.com/soumik12345/wandb-addons
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.environ["KERAS_BACKEND"] = "torch"
+
+ import numpy as np
+ from tqdm.auto import tqdm
+ import matplotlib.pyplot as plt
+
+ import torch
+ import torch.nn as nn
+ import torch.nn.functional as F
+
+ import torchvision
+ from torchvision import datasets, models, transforms
+
+ import keras_core as keras
+ from keras_core.utils import TorchModuleWrapper
+
+ import wandb
+ from wandb_addons.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+ return (
+ TorchModuleWrapper,
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ datasets,
+ keras,
+ models,
+ nn,
+ np,
+ os,
+ plt,
+ torch,
+ tqdm,
+ transforms,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Define the Hyperparameters
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project="keras-torch", entity="ml-colabs", job_type="torchvision/train")
+
+ config = wandb.config
+ config.batch_size = 32
+ config.image_size = 224
+ config.freeze_backbone = True
+ config.initial_learning_rate = 1e-3
+ config.num_epochs = 5
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Creating the Torch Datasets and Dataloaders
+
+ In this example, we would train an image classification model on the [Imagenette dataset](https://github.com/fastai/imagenette). Imagenette is a subset of 10 easily classified classes from [Imagenet](https://www.image-net.org/) (tench, English springer, cassette player, chain saw, church, French horn, garbage truck, gas pump, golf ball, parachute).
+ """)
+ return
+
+
+@app.cell
+def _(keras):
+ # Fetch the imagenette dataset
+ data_dir = keras.utils.get_file(
+ fname="imagenette2-320.tgz",
+ origin="https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz",
+ extract=True,
+ )
+ data_dir = data_dir.replace(".tgz", "")
+ return (data_dir,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we define pre-processing and augmentation transforms from TorchVision for the train and validation sets.
+ """)
+ return
+
+
+@app.cell
+def _(config, transforms):
+ data_transforms = {
+ 'train': transforms.Compose([
+ transforms.RandomResizedCrop(config.image_size),
+ transforms.RandomHorizontalFlip(),
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ]),
+ 'val': transforms.Compose([
+ transforms.Resize(256),
+ transforms.CenterCrop(config.image_size),
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ]),
+ }
+ return (data_transforms,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Finally, we will use TorchVision and the [`torch.utils.data`](https://pytorch.org/docs/stable/data.html) packages for creating the dataloaders for trainig and validation.
+ """)
+ return
+
+
+@app.cell
+def _(config, data_dir, data_transforms, datasets, os, torch):
+ # Define the train and validation datasets
+ image_datasets = {
+ x: datasets.ImageFolder(
+ os.path.join(data_dir, x), data_transforms[x]
+ )
+ for x in ['train', 'val']
+ }
+
+ # Define the torch dataloaders corresponding to the
+ # train and validation dataset
+ dataloaders = {
+ x: torch.utils.data.DataLoader(
+ image_datasets[x],
+ batch_size=config.batch_size,
+ shuffle=True,
+ num_workers=4
+ )
+ for x in ['train', 'val']
+ }
+ dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
+ class_names = image_datasets['train'].classes
+ return class_names, dataloaders
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let us visualize a few samples from the training dataloader.
+ """)
+ return
+
+
+@app.cell
+def _(class_names, dataloaders, np, plt):
+ plt.figure(figsize=(10, 10))
+ _sample_images, _sample_labels = next(iter(dataloaders['train']))
+ _sample_images = _sample_images.numpy()
+ _sample_labels = _sample_labels.numpy()
+ for _idx in range(9):
+ ax = plt.subplot(3, 3, _idx + 1)
+ _image = _sample_images[_idx].transpose((1, 2, 0))
+ _mean = np.array([0.485, 0.456, 0.406])
+ _std = np.array([0.229, 0.224, 0.225])
+ _image = _std * _image + _mean
+ _image = np.clip(_image, 0, 1)
+ plt.imshow(_image)
+ plt.title('Ground Truth Label: ' + class_names[int(_sample_labels[_idx])])
+ plt.axis('off')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## The Image Classification Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We typically define a model in PyTorch using [`torch.nn.Module`s](https://pytorch.org/docs/stable/notes/modules.html) which act as the building blocks of stateful computation. Let us define the ResNet18 model from the TorchVision package as a `torch.nn.Module` pre-trained on the [Imagenet1K dataset](https://huggingface.co/datasets/imagenet-1k).
+ """)
+ return
+
+
+@app.cell
+def _(models, nn):
+ # Define the pre-trained resnet18 module from TorchVision
+ resnet_18 = models.resnet18(weights='IMAGENET1K_V1')
+
+ # We set the classification head of the pre-trained ResNet18
+ # module to an identity module
+ resnet_18.fc = nn.Identity()
+ return (resnet_18,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ven though Keras supports PyTorch as a backend, it does not mean that we can nest torch modules inside a [`keras_core.Model`](https://keras.io/keras_core/api/models/), because trainable variables inside a Keras Model is tracked exclusively via [Keras Layers](https://keras.io/keras_core/api/layers/).
+
+ KerasCore provides us with a feature called `TorchModuleWrapper` which enables us to do exactly this. The `TorchModuleWrapper` is a Keras Layer that accepts a torch module and tracks its trainable variables, essentially converting the torch module into a Keras Layer. This enables us to put any torch modules inside a Keras Model and train them with a single `model.fit()`!
+ """)
+ return
+
+
+@app.cell
+def _(TorchModuleWrapper, config, resnet_18):
+ # We set the trainable ResNet18 backbone to be a Keras Layer
+ # using `TorchModuleWrapper`
+ backbone = TorchModuleWrapper(resnet_18)
+
+ # We set this to `False` if you want to freeze the backbone
+ backbone.trainable = config.freeze_backbone
+ return (backbone,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we will build a Keras functional model with the backbone layer.
+ """)
+ return
+
+
+@app.cell
+def _(backbone, class_names, config, keras):
+ inputs = keras.Input(shape=(3, config.image_size, config.image_size))
+ x = backbone(inputs)
+ x = keras.layers.Dropout(0.5)(x)
+ x = keras.layers.Dense(len(class_names))(x)
+ outputs = keras.activations.softmax(x, axis=1)
+ model = keras.Model(inputs, outputs, name="ResNet18_Classifier")
+
+ model.summary()
+ return (model,)
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ config,
+ dataloaders,
+ keras,
+ model,
+):
+ # Create exponential decay learning rate scheduler
+ decay_steps = config.num_epochs * len(dataloaders["train"]) // config.batch_size
+ lr_scheduler = keras.optimizers.schedules.ExponentialDecay(
+ initial_learning_rate=config.initial_learning_rate,
+ decay_steps=decay_steps,
+ decay_rate=0.1,
+ )
+
+ # Compile the model
+ model.compile(
+ loss="sparse_categorical_crossentropy",
+ optimizer=keras.optimizers.Adam(lr_scheduler),
+ metrics=["accuracy"],
+ )
+
+ # Define the backend-agnostic WandB callbacks for KerasCore
+ callbacks = [
+ # Track experiment metrics with WandB
+ WandbMetricsLogger(log_freq="batch"),
+ # Save best model checkpoints to WandB
+ WandbModelCheckpoint(
+ filepath="model.weights.h5",
+ monitor="val_loss",
+ save_best_only=True,
+ save_weights_only=True,
+ )
+ ]
+
+ # Train the model by calling model.fit
+ history = model.fit(
+ dataloaders["train"],
+ validation_data=dataloaders["val"],
+ epochs=config.num_epochs,
+ callbacks=callbacks,
+ )
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Evaluation and Inference
+
+ Now, we let us load the best model weights checkpoint and evaluate the model.
+ """)
+ return
+
+
+@app.cell
+def _(dataloaders, model, os, wandb):
+ wandb.init(
+ project="keras-torch", entity="ml-colabs", job_type="torchvision/eval"
+ )
+ artifact = wandb.use_artifact(
+ 'ml-colabs/keras-torch/run_hiceci7f_model:latest', type='model'
+ )
+ artifact_dir = artifact.download()
+
+ model.load_weights(os.path.join(artifact_dir, "model.weights.h5"))
+
+ _, val_accuracy = model.evaluate(dataloaders["val"])
+ wandb.log({"Validation-Accuracy": val_accuracy})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Finally, let us visualize the some predictions of the model
+ """)
+ return
+
+
+@app.cell
+def _(class_names, dataloaders, keras, model, np, tqdm, wandb):
+ table = wandb.Table(columns=['Image', 'Ground-Truth', 'Prediction'] + ['Confidence-' + cls for cls in class_names])
+ _sample_images, _sample_labels = next(iter(dataloaders['train']))
+ sample_pred_probas = model(_sample_images.to('cuda')).detach()
+ sample_pred_logits = keras.ops.argmax(sample_pred_probas, axis=1)
+ sample_pred_logits = sample_pred_logits.to('cpu').numpy()
+ sample_pred_probas = sample_pred_probas.to('cpu').numpy()
+ _sample_images = _sample_images.numpy()
+ _sample_labels = _sample_labels.numpy()
+ # We perform inference and detach the predicted probabilities from the Torch
+ # computation graph with a tensor that does not require gradient computation.
+ for _idx in tqdm(range(_sample_images.shape[0])):
+ _image = _sample_images[_idx].transpose((1, 2, 0))
+ _mean = np.array([0.485, 0.456, 0.406])
+ _std = np.array([0.229, 0.224, 0.225])
+ _image = _std * _image + _mean
+ _image = np.clip(_image, 0, 1)
+ table.add_data(wandb.Image(_image), class_names[int(_sample_labels[_idx])], class_names[int(sample_pred_logits[_idx])], *sample_pred_probas[_idx].tolist())
+ wandb.log({'Evaluation-Table': table})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-keras-nsynth-instrument-prediction/keras_keras_nsynth_instrument_prediction.py b/marimo/convert/keras-keras-nsynth-instrument-prediction/keras_keras_nsynth_instrument_prediction.py
new file mode 100644
index 00000000..73b4d030
--- /dev/null
+++ b/marimo/convert/keras-keras-nsynth-instrument-prediction/keras_keras_nsynth_instrument_prediction.py
@@ -0,0 +1,610 @@
+# /// script
+# dependencies = ["soundfile", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # NSynth Instrument Prediction using Keras
+
+
+
+ Based on the [Medium post](https://bit.ly/2UaNKQp) made by David Schwertfeger
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb soundfile !pip install -Uq wandb soundfile
+ return
+
+
+@app.cell
+def _():
+ import tensorflow as tf
+ import tensorflow_datasets as tfds
+ import wandb
+
+ return tf, tfds, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(tf):
+ tf.__version__
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Set Configuration Values
+ """)
+ return
+
+
+@app.cell
+def _():
+ _N_CLASSES = 11
+ _SAMPLE_RATE = 16000
+ _DURATION = 4 #seconds?
+ return
+
+
+@app.cell
+def _():
+ _FFT_SIZE = 1024
+ _HOP_SIZE = 512
+ _N_MEL_BINS = 64
+ _N_SPECTROGRAM_BINS = (_FFT_SIZE // 2) + 1
+ _F_MIN = 0.0
+ _F_MAX = _SAMPLE_RATE / 2
+ return
+
+
+@app.cell
+def _():
+ _TRAIN_DS_SIZE = 289205 # Adjust this to reduce the amount of data during training
+ _TRAIN_EPOCHS = 2
+ _TRAIN_BATCH_SIZE = 256
+ _TRAIN_STEPS = 40000 // _TRAIN_BATCH_SIZE
+ return
+
+
+@app.cell
+def _():
+ _VAL_DS_SIZE = 12678 # Adjust this to reduce the amount of data during validation
+ _VAL_BATCH_SIZE = 256
+ _VAL_STEPS = _VAL_DS_SIZE / _VAL_BATCH_SIZE
+ return
+
+
+@app.cell
+def _():
+ _TEST_DS_SIZE = 4096
+ _TEST_BATCH_SIZE = 256
+ _TEST_STEPS = _TEST_DS_SIZE / _TEST_BATCH_SIZE
+ return
+
+
+@app.cell
+def _():
+ model_config_defaults = {
+ #Dataset Specific
+ "n_classes": _N_CLASSES,
+ "sample_rate" : _SAMPLE_RATE,
+ "duration": _DURATION,
+
+ #Model specific
+ "fft_size" : _FFT_SIZE,
+ "hop_size" : _HOP_SIZE,
+ "n_mels" : _N_MEL_BINS,
+ "f_min" : _F_MIN,
+ "f_max" : _F_MAX,
+
+ #Training data
+ "train_ds_size": _TRAIN_DS_SIZE,
+ "train_epochs": _TRAIN_EPOCHS,
+ "train_batch_size": _TRAIN_BATCH_SIZE,
+ "train_steps": _TRAIN_STEPS,
+
+ #Validation data
+ "val_ds_size": _VAL_DS_SIZE,
+ "val_batch_size": _VAL_BATCH_SIZE,
+ "val_steps": _VAL_STEPS,
+
+ #Testing data
+ "test_ds_size": _TEST_DS_SIZE,
+ "test_batch_size": _TEST_BATCH_SIZE,
+ "test_steps": _TEST_STEPS,
+ }
+ return (model_config_defaults,)
+
+
+@app.cell
+def _(model_config_defaults, wandb):
+ run = wandb.init(config = model_config_defaults, project="keras_nsynth_instrument_prediction-test")
+ model_config = run.config
+ return model_config, run
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Load ~~and Save Data~~
+ """)
+ return
+
+
+@app.cell
+def _(tfds):
+ # Load NSynth's test split as a tf.data.Dataset
+ # https://www.tensorflow.org/datasets/catalog/nsynth
+ (raw_train_ds, raw_validation_ds, raw_test_ds), ds_info = tfds.load(name='nsynth/full',
+ split=['train', 'valid', "test"],
+ try_gcs=True,
+ with_info=True)
+ return ds_info, raw_test_ds, raw_train_ds, raw_validation_ds
+
+
+@app.cell
+def _(ds_info):
+ ds_info
+ return
+
+
+@app.cell
+def _(tf):
+ def prep_data(raw_ds, batch_size, data_type):
+ # Let's train a model to predict the instrument family from audio
+ # https://magenta.tensorflow.org/datasets/nsynth#instrument-families
+ ds = raw_ds.map(lambda x: (x['audio'], x['instrument']['family']))
+
+ # Build your input pipeline
+ if data_type in ["train", "validation"]:
+ prepped_ds = (ds
+ .shuffle(1000, reshuffle_each_iteration=True) #is having 2 shuffles redundant?
+ .batch(batch_size)
+ .prefetch(tf.data.AUTOTUNE)
+ .repeat()
+ )
+ else:
+ prepped_ds = (ds
+ .batch(batch_size)
+ .prefetch(tf.data.AUTOTUNE)
+ )
+ return prepped_ds
+
+ return (prep_data,)
+
+
+@app.cell
+def _(model_config, prep_data, raw_test_ds, raw_train_ds, raw_validation_ds):
+ train_ds = prep_data(raw_train_ds, model_config["train_batch_size"], "train")
+ validation_ds = prep_data(raw_validation_ds, model_config["val_batch_size"], "validation")
+ test_ds = prep_data(raw_test_ds, model_config["test_batch_size"], "test")
+ return test_ds, train_ds, validation_ds
+
+
+@app.cell
+def _():
+ #Causes disk error
+ # tf.data.experimental.save(train_ds, "./train")
+ # tf.data.experimental.save(validation_ds, "./val")
+ # tf.data.experimental.save(test_ds, "./test")
+ return
+
+
+@app.cell
+def _():
+ # train_data_artifact = wandb.Artifact(name="nsynth_train", type="dataset")
+ # train_data_artifact.add_dir("./train")
+ # run.log_artifact(train_data_artifact)
+ return
+
+
+@app.cell
+def _():
+ # val_data_artifact = wandb.Artifact(name="nsynth_val", type="dataset")
+ # val_data_artifact.add_dir("./val")
+ # run.log_artifact(val_data_artifact)
+ return
+
+
+@app.cell
+def _():
+ # test_data_artifact = wandb.Artifact(name="nsynth_test", type="dataset")
+ # test_data_artifact.add_dir("./test")
+ # run.log_artifact(test_data_artifact)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Create Keras Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Define Custom LogMel Layer
+ """)
+ return
+
+
+@app.cell
+def _(tf):
+ class LogMelSpectrogram(tf.keras.layers.Layer):
+ """Compute log-magnitude mel-scaled spectrograms."""
+
+ def __init__(self, sample_rate, fft_size, hop_size, n_mels,
+ f_min=0.0, f_max=None, **kwargs):
+ super(LogMelSpectrogram, self).__init__(**kwargs)
+ self.sample_rate = sample_rate
+ self.fft_size = fft_size
+ self.hop_size = hop_size
+ self.n_mels = n_mels
+ self.f_min = f_min
+ self.f_max = f_max if f_max else sample_rate / 2
+ self.mel_filterbank = tf.signal.linear_to_mel_weight_matrix(
+ num_mel_bins=self.n_mels,
+ num_spectrogram_bins=fft_size // 2 + 1,
+ sample_rate=self.sample_rate,
+ lower_edge_hertz=self.f_min,
+ upper_edge_hertz=self.f_max)
+
+ def build(self, input_shape):
+ self.non_trainable_weights.append(self.mel_filterbank)
+ super(LogMelSpectrogram, self).build(input_shape)
+
+ def call(self, waveforms):
+ """Forward pass.
+
+ Parameters
+ ----------
+ waveforms : tf.Tensor, shape = (None, n_samples)
+ A Batch of mono waveforms.
+
+ Returns
+ -------
+ log_mel_spectrograms : (tf.Tensor), shape = (None, time, freq, ch)
+ The corresponding batch of log-mel-spectrograms
+ """
+ def _tf_log10(x):
+ numerator = tf.math.log(x)
+ denominator = tf.math.log(tf.constant(10, dtype=numerator.dtype))
+ return numerator / denominator
+
+ def power_to_db(magnitude, amin=1e-16, top_db=80.0):
+ """
+ https://librosa.github.io/librosa/generated/librosa.core.power_to_db.html
+ """
+ ref_value = tf.reduce_max(magnitude)
+ log_spec = 10.0 * _tf_log10(tf.maximum(amin, magnitude))
+ log_spec -= 10.0 * _tf_log10(tf.maximum(amin, ref_value))
+ log_spec = tf.maximum(log_spec, tf.reduce_max(log_spec) - top_db)
+
+ return log_spec
+
+ spectrograms = tf.signal.stft(waveforms,
+ frame_length=self.fft_size,
+ frame_step=self.hop_size,
+ pad_end=False)
+
+ magnitude_spectrograms = tf.abs(spectrograms)
+
+ mel_spectrograms = tf.matmul(tf.square(magnitude_spectrograms),
+ self.mel_filterbank)
+
+ log_mel_spectrograms = power_to_db(mel_spectrograms)
+
+ # add channel dimension
+ log_mel_spectrograms = tf.expand_dims(log_mel_spectrograms, 3)
+
+ return log_mel_spectrograms
+
+ def get_config(self):
+ config = {
+ 'fft_size': self.fft_size,
+ 'hop_size': self.hop_size,
+ 'n_mels': self.n_mels,
+ 'sample_rate': self.sample_rate,
+ 'f_min': self.f_min,
+ 'f_max': self.f_max,
+ }
+ config.update(super(LogMelSpectrogram, self).get_config())
+
+ return config
+
+ return (LogMelSpectrogram,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Use LogMel Layer in Keras Model
+ """)
+ return
+
+
+@app.cell
+def _():
+ from tensorflow.keras.layers import (BatchNormalization, Conv2D, Dense,
+ Dropout, Flatten, Input, MaxPool2D)
+ from tensorflow.keras.models import Model
+
+ return (
+ BatchNormalization,
+ Conv2D,
+ Dense,
+ Dropout,
+ Flatten,
+ Input,
+ MaxPool2D,
+ Model,
+ )
+
+
+@app.cell
+def _(
+ BatchNormalization,
+ Conv2D,
+ Dense,
+ Dropout,
+ Flatten,
+ Input,
+ LogMelSpectrogram,
+ MaxPool2D,
+ Model,
+):
+ def ConvModel(n_classes, sample_rate, duration, fft_size, hop_size, n_mels, f_min=0.0, f_max=None, **kwargs):
+ n_samples = sample_rate * duration
+ input_shape = (n_samples,)
+
+ x = Input(shape=input_shape, name='input', dtype='float32')
+ y = LogMelSpectrogram(sample_rate, fft_size, hop_size, n_mels, f_min, f_max)(x)
+
+ # data normalization (on frequency axis)
+ y = BatchNormalization(axis=2)(y)
+
+ # effectively 1D convolution, since kernel spans entire frequency-axis
+ y = Conv2D(32, (3, n_mels), activation='relu')(y)
+ y = BatchNormalization()(y)
+ y = MaxPool2D((1, y.shape[2]))(y)
+
+ y = Conv2D(32, (3, 1), activation='relu')(y)
+ y = BatchNormalization()(y)
+ y = MaxPool2D(pool_size=(2, 1))(y)
+
+ y = Flatten()(y)
+ y = Dense(64, activation='relu')(y)
+ y = Dropout(0.25)(y)
+ y = Dense(n_classes, activation='softmax')(y)
+
+ return Model(inputs=x, outputs=y)
+
+ return (ConvModel,)
+
+
+@app.cell
+def _(ConvModel, model_config):
+ model = ConvModel(**model_config)
+ model.compile(optimizer='adam',
+ loss='sparse_categorical_crossentropy',
+ metrics=['sparse_categorical_accuracy'])
+ model.summary()
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Aside: Visualize Keras Model
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install visualkeras
+ return
+
+
+@app.cell
+def _():
+ import visualkeras
+
+ return (visualkeras,)
+
+
+@app.cell
+def _(model, visualkeras):
+ visualkeras.layered_view(model, to_file='model.png', legend=True)
+ return
+
+
+@app.cell
+def _(run, wandb):
+ run.log({"model_image": wandb.Image("model.png", caption="Visualized Keras Model")})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Define Callbacks
+ """)
+ return
+
+
+@app.cell
+def _():
+ from wandb.keras import WandbMetricsLogger
+
+ return (WandbMetricsLogger,)
+
+
+@app.cell
+def _():
+ labels = ["bass", "brass", "flute", "guitar", "keyboard", "mallet", "organ", "reed", "string", "synth_lead", "vocal"]
+ return (labels,)
+
+
+@app.cell
+def _(WandbMetricsLogger):
+ wandb_callback = WandbMetricsLogger(log_freq=2)
+ return (wandb_callback,)
+
+
+@app.cell
+def _():
+ import numpy as np
+ from tensorflow.keras.callbacks import Callback
+ from sklearn.metrics import accuracy_score
+
+ return Callback, accuracy_score, np
+
+
+@app.cell
+def _(Callback, accuracy_score, np, run, wandb):
+ class AudioPredictionCallback(Callback):
+ def __init__(self, labels, prediction_data, sr):
+
+ super(AudioPredictionCallback, self).__init__()
+ self.labels = labels
+ self.prediction_data = prediction_data
+ self.sr = sr
+
+ def on_epoch_end(self, epoch, logs=None):
+ id_list = []
+ input_audio = []
+ true_index = []
+ true_labels = []
+ prediction_probs_list = []
+ predicted_index = []
+ predicted_labels = []
+
+ for batch_x, batch_y in self.prediction_data:
+ prediction_probs = self.model.predict(batch_x)
+ predictions = prediction_probs.argmax(axis=1)
+
+ for x in batch_x:
+ wandb_audio = wandb.Audio(x, sample_rate=self.sr)
+ input_audio.append(wandb_audio)
+
+ for y in batch_y:
+ true_index.append(y.numpy())
+ true_labels.append(self.labels[y])
+
+ for prediction_prob in prediction_probs:
+ prediction_probs_list.append(prediction_prob)
+
+ for pred in predictions:
+ predicted_index.append(pred)
+ predicted_labels.append(self.labels[pred])
+
+ #All ids should match on repeate calls as the data is assumed to be never shuffled
+ id_list = list(range(len(input_audio)))
+
+ table_data = np.array([id_list, input_audio, true_labels, predicted_labels]).T
+ columns = ["id", "audio", "true", "prediction"]
+ prediction_table = wandb.Table(data=table_data, columns=columns)
+
+ prediction_table_artifact = wandb.Artifact(name="audio_table", type="prediction")
+ prediction_table_artifact.add(prediction_table, "audio_table")
+
+ acc = accuracy_score(true_labels, predicted_labels)
+
+ cm = wandb.plot.confusion_matrix(
+ y_true=true_index,
+ preds=predicted_index,
+ class_names = self.labels,
+ title="Confusion Matrix")
+ pr_curve = wandb.plot.pr_curve(true_index, prediction_probs_list, labels=self.labels, title="Precision vs. Recall")
+
+ run.log_artifact(prediction_table_artifact)
+ run.log({
+ "test_cm": cm,
+ "test_pr_curve": pr_curve,
+ "test_acc": acc
+ })
+
+ return (AudioPredictionCallback,)
+
+
+@app.cell
+def _(AudioPredictionCallback, labels, model_config, test_ds):
+ audio_prediction_callback = AudioPredictionCallback(labels, test_ds, model_config["sample_rate"])
+ return (audio_prediction_callback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Train Model
+ """)
+ return
+
+
+@app.cell
+def _(
+ audio_prediction_callback,
+ model,
+ model_config,
+ train_ds,
+ validation_ds,
+ wandb_callback,
+):
+ model.fit(train_ds,
+ epochs=model_config["train_epochs"],
+ steps_per_epoch=model_config["train_steps"],
+ validation_data = validation_ds,
+ validation_steps = model_config["val_steps"],
+ callbacks=[wandb_callback, audio_prediction_callback],
+ verbose="auto",
+ shuffle=True #Do i need this?
+ )
+ return
+
+
+@app.cell
+def _(run):
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-keras-param-opti-using-sweeps/keras_keras_param_opti_using_sweeps.py b/marimo/convert/keras-keras-param-opti-using-sweeps/keras_keras_param_opti_using_sweeps.py
new file mode 100644
index 00000000..b3bcd016
--- /dev/null
+++ b/marimo/convert/keras-keras-param-opti-using-sweeps/keras_keras_param_opti_using_sweeps.py
@@ -0,0 +1,565 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🧹 Introduction to Hyperparameter Sweeps using W&B and Keras
+
+ 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 pick the most accurate model. They enable this by automatically searching through combinations of hyperparameter values (e.g. learning rate, batch size, number of hidden layers, optimizer type) to find the most optimal values.
+
+ In this tutorial we'll see how you can run sophisticated hyperparameter sweeps in 3 easy steps using Weights and Biases.
+
+ 
+
+ ## 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.com/library/sweeps/configuration) that specifies the parameters to search through, the search strategy, the optimization metric et all.
+
+ 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)`
+
+ 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 so you can tweak the parameters,
+ try out different models,
+ or try a Sweep with your own dataset!
+
+ ## Resources
+ - [Sweeps docs →](https://docs.wandb.ai/sweeps)
+ - [Launching from the command line →](https://www.wandb.com/articles/hyperparameter-tuning-as-easy-as-1-2-3)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🚀 Install, Import, and Log in
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 0️⃣: Install W&B
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -Uq wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 1️⃣: Import W&B and Login
+ """)
+ return
+
+
+@app.cell
+def _():
+ import numpy as np
+ import tensorflow as tf
+ from tensorflow import keras
+ from keras import layers
+
+ return keras, layers, np, tf
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+ return WandbMetricsLogger, WandbModelCheckpoint, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > Side note: If this is your first time using W&B or you are not logged in, the link that appears after running `wandb.login()` will take you to sign-up/login page. Signing up is as easy as a few clicks.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 👩🍳 Prepare Dataset
+
+ We will use MNIST directly from `keras.datasets`
+ """)
+ return
+
+
+@app.cell
+def _(keras, np):
+ # Get the dataset
+ (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
+ num_classes = len(np.unique(y_train))
+ input_shape = x_train.shape[-2:] + (1,)
+ return input_shape, num_classes, x_test, x_train, y_test, y_train
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ prepare data for training:
+ - scale in [0,1]
+ - transform targets to cateforical
+ """)
+ return
+
+
+@app.cell
+def _(keras, np, num_classes, x_test, x_train, y_test, y_train):
+ # Scale
+ x_train_1 = x_train / 255.0
+ x_test_1 = x_test / 255.0
+ x_train_1 = np.expand_dims(x_train_1, -1)
+ # Make sure images have shape (28, 28, 1)
+ x_test_1 = np.expand_dims(x_test_1, -1)
+ y_train_1 = keras.utils.to_categorical(y_train, num_classes)
+ # convert class vectors to binary class matrices
+ y_test_1 = keras.utils.to_categorical(y_test, num_classes)
+ return x_test_1, x_train_1, y_test_1, y_train_1
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🧠 Define the Model and Training Loop
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 2️⃣ 🏗️ Build a Simple Classifier
+ """)
+ return
+
+
+@app.cell
+def _(input_shape, keras, layers, num_classes):
+ def ConvNet(dropout=0.2):
+ return keras.Sequential(
+ [
+ keras.Input(shape=input_shape),
+ layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
+ layers.MaxPooling2D(pool_size=(2, 2)),
+ layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
+ layers.MaxPooling2D(pool_size=(2, 2)),
+ layers.Flatten(),
+ layers.Dropout(dropout),
+ layers.Dense(num_classes, activation="softmax"),
+ ]
+ )
+
+ model = ConvNet()
+
+ model.summary()
+ return (ConvNet,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 3️⃣ Write a Training script
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ tf,
+ x_test_1,
+ x_train_1,
+ y_test_1,
+ y_train_1,
+):
+ def get_optimizer(lr=0.001, optimizer='adam'):
+ """Select optmizer between adam and sgd with momentum"""
+ if optimizer.lower() == 'adam':
+ return tf.keras.optimizers.Adam(learning_rate=lr)
+ if optimizer.lower() == 'sgd':
+ return tf.keras.optimizers.SGD(learning_rate=lr, momentum=0.1)
+
+ def train(model, batch_size=64, epochs=10, lr=0.001, optimizer='adam', log_freq=10):
+ tf.keras.backend.clear_session()
+ model.compile(loss='categorical_crossentropy', optimizer=get_optimizer(lr, optimizer), metrics=['accuracy']) # Compile model like you usually do.
+ wandb_callbacks = [WandbMetricsLogger(log_freq=log_freq), WandbModelCheckpoint(filepath='my_model_{epoch:02d}')]
+ model.fit(x_train_1, y_train_1, batch_size=batch_size, epochs=epochs, validation_data=(x_test_1, y_test_1), callbacks=wandb_callbacks) # callback setup
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 4️⃣ Define the Sweep
+
+ Fundamentally, a Sweep combines a strategy for trying out a bunch of hyperparameter values with the code that evalutes them.
+ Whether that strategy is as simple as trying every option
+ or as complex as [BOHB](https://arxiv.org/abs/1807.01774),
+ Weights & Biases Sweeps have you covered.
+ You just need to _define your strategy_
+ in the form of a [configuration](https://docs.wandb.com/sweeps/configuration).
+
+ When you're setting up a Sweep in a notebook like this,
+ that config object is a nested dictionary.
+ When you run a Sweep via the command line,
+ the config object is a
+ [YAML file](https://docs.wandb.com/sweeps/quickstart#2-sweep-config).
+
+ Let's walk through the definition of a Sweep config together.
+ We'll do it slowly, so we get a chance to explain each component.
+ In a typical Sweep pipeline,
+ this step would be done in a single assignment.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 👈 Pick a `method`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The first thing we need to define is the `method`
+ for choosing new parameter values.
+
+ We provide the following search `methods`:
+ * **`grid` Search** – Iterate over every combination of hyperparameter values.
+ Very effective, but can be computationally costly.
+ * **`random` Search** – Select each new combination at random according to provided `distribution`s. Surprisingly effective!
+ * **`bayes`ian Search** – Create a probabilistic model of metric score as a function of the hyperparameters, and choose parameters with high probability of improving the metric. Works well for small numbers of continuous parameters but scales poorly.
+
+ We'll stick with `random`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ sweep_config = {
+ 'method': 'bayes'
+ }
+ return (sweep_config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For `bayes`ian Sweeps,
+ you also need to tell us a bit about your `metric`.
+ We need to know its `name`, so we can find it in the model outputs
+ and we need to know whether your `goal` is to `minimize` it
+ (e.g. if it's the squared error)
+ or to `maximize` it
+ (e.g. if it's the accuracy).
+ """)
+ return
+
+
+@app.cell
+def _(sweep_config):
+ metric = {
+ 'name': 'val_loss',
+ 'goal': 'minimize'
+ }
+
+ sweep_config['metric'] = metric
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If you're not running a `bayes`ian Sweep, you don't have to,
+ but it's not a bad idea to include this in your `sweep_config` anyway,
+ in case you change your mind later.
+ It's also good reproducibility practice to keep note of things like this,
+ in case you, or someone else,
+ come back to your Sweep in 6 months or 6 years
+ and don't know whether `val_G_batch` is supposed to be high or low.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 📃 Name the hyper`parameters`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once you've picked a `method` to try out new values of the hyperparameters,
+ you need to define what those `parameters` are.
+
+ Most of the time, this step is straightforward:
+ you just give the `parameter` a name
+ and specify a list of legal `values`
+ of the parameter.
+
+ For example, when we choose the `optimizer` for our network,
+ there's only a finite number of options.
+ Here we stick with the two most popular choices, `adam` and `sgd`.
+ Even for hyperparameters that have potentially infinite options,
+ it usually only makes sense to try out
+ a few select `values`,
+ as we do here with `dropout`.
+ """)
+ return
+
+
+@app.cell
+def _(sweep_config):
+ parameters_dict = {
+ 'optimizer': {
+ 'values': ['adam', 'sgd']
+ },
+ 'dropout': {
+ 'values': [0.1, 0.3, 0.5]
+ },
+ }
+
+ sweep_config['parameters'] = parameters_dict
+ return (parameters_dict,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ It's often the case that there are hyperparameters
+ that we don't want to vary in this Sweep,
+ but which we still want to set in our `sweep_config`.
+
+ In that case, we just set the `value` directly:
+ """)
+ return
+
+
+@app.cell
+def _(parameters_dict):
+ parameters_dict.update({
+ 'epochs': {
+ 'value': 1}
+ })
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For a `grid` search, that's all you ever need.
+
+ For a `random` search,
+ all the `values` of a parameter are equally likely to be chosen on a given run.
+
+ If that just won't do,
+ you can instead specify a named `distribution`,
+ plus its parameters, like the mean `mu`
+ and standard deviation `sigma` of a `normal` distribution.
+
+ See more on how to set the distributions of your random variables [here](https://docs.wandb.com/sweeps/configuration#distributions).
+ """)
+ return
+
+
+@app.cell
+def _(parameters_dict):
+ import math
+
+ parameters_dict.update({
+ 'learning_rate': {
+ # a flat distribution between 0 and 0.1
+ 'distribution': 'uniform',
+ 'min': 0.001,
+ 'max': 0.1
+ },
+ 'batch_size': {
+ 'values': [64, 128]
+ }
+ })
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ When we're finished, `sweep_config` is a nested dictionary
+ that specifies exactly which `parameters` we're interested in trying
+ and what `method` we're going to use to try them.
+ """)
+ return
+
+
+@app.cell
+def _(sweep_config):
+ import pprint
+
+ pprint.pprint(sweep_config)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ But that's not all of the configuration options!
+
+ For example, we also offer the option to `early_terminate` your runs with the [HyperBand](https://arxiv.org/pdf/1603.06560.pdf) scheduling algorithm. See more [here](https://docs.wandb.com/sweeps/configuration#stopping-criteria).
+
+ You can find a list of all configuration options [here](https://docs.wandb.com/library/sweeps/configuration)
+ and a big collection of examples in YAML format [here](https://github.com/wandb/examples/tree/master/examples/keras/keras-cnn-fashion).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 5️⃣: Wrap the Training Loop
+
+ You'll need a function, like `sweep_train` below,
+ that uses `wandb.config` to set the hyperparameters
+ before `train` gets called.
+ """)
+ return
+
+
+@app.cell
+def _(ConvNet, train, wandb):
+ def sweep_train(config_defaults=None):
+ # Initialize wandb with a sample project name
+ with wandb.init(config=config_defaults): # this gets over-written in the Sweep
+
+ # Specify the other hyperparameters to the configuration, if any
+ wandb.config.architecture_name = "ConvNet"
+ wandb.config.dataset_name = "MNIST"
+
+ # initialize model
+ model = ConvNet(wandb.config.dropout)
+
+ train(model,
+ wandb.config.batch_size,
+ wandb.config.epochs,
+ wandb.config.learning_rate,
+ wandb.config.optimizer)
+
+ return (sweep_train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 6️⃣: Initialize Sweep and Run Agent
+ """)
+ return
+
+
+@app.cell
+def _(sweep_config, wandb):
+ sweep_id = wandb.sweep(sweep_config, project="sweeps-keras")
+ return (sweep_id,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can limit the number of total runs with the `count` parameter, we will limit a 10 to make the script run fast, feel free to increase the number of runs and see what happens.
+ """)
+ return
+
+
+@app.cell
+def _(sweep_id, sweep_train, wandb):
+ wandb.agent(sweep_id, function=sweep_train, count=10)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👀 Visualize Results
+
+ Click on the **Sweep URL** link above to see your live results.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🤓 Advanced Setup
+ 1. [Environment variables](https://docs.wandb.com/library/environment-variables): Set API keys in environment variables so you can run training on a managed cluster.
+ 2. [Offline mode](https://docs.wandb.com/library/technical-faq#can-i-run-wandb-offline): Use `dryrun` mode to train offline and sync results later.
+ 3. [On-prem](https://docs.wandb.com/self-hosted): 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.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-restorers-evaluation-low-light/keras_restorers_evaluation_low_light.py b/marimo/convert/keras-restorers-evaluation-low-light/keras_restorers_evaluation_low_light.py
new file mode 100644
index 00000000..52001c51
--- /dev/null
+++ b/marimo/convert/keras-restorers-evaluation-low-light/keras_restorers_evaluation_low_light.py
@@ -0,0 +1,88 @@
+# /// script
+# dependencies = ["pip", "restorers @ git+https://github.com/soumik12345/restorers.git", "setuptools"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌈 Restorers + WandB 🪄🐝
+
+
+
+ This notebook shows how to perform inference with a low-light enhancement using [**restorers**](https://github.com/soumik12345/restorers) and [**wandb**](https://wandb.ai/site). For more details regarding usage of restorers, refer to the following report:
+
+ [](https://wandb.ai/ml-colabs/low-light-enhancement/reports/Lighting-up-Images-in-the-Deep-Learning-Era--VmlldzozNzE4Njkz)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: pip setuptools !pip install -q --upgrade pip setuptools
+ # packages added via marimo's package management: git+https://github.com/soumik12345/restorers.git !pip install git+https://github.com/soumik12345/restorers.git
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from restorers.evaluation import LoLEvaluator
+ from restorers.metrics import PSNRMetric, SSIMMetric
+
+ return LoLEvaluator, PSNRMetric, SSIMMetric, wandb
+
+
+@app.cell
+def _(wandb):
+ # initialize a wandb run for inference
+ wandb.init(project="low-light-enhancement", job_type="evaluation")
+ return
+
+
+@app.cell
+def _(LoLEvaluator, PSNRMetric, SSIMMetric):
+ # Define the Evaluator for LoL dataset
+ evaluator = LoLEvaluator(
+ # pass the list of Keras metrics to be evaluated for
+ metrics=[PSNRMetric(max_val=1.0), SSIMMetric(max_val=1.0)],
+ # pass the wandb artifact for the LoL dataset
+ dataset_artifact_address="ml-colabs/dataset/LoL:v0",
+ input_size=256,
+ )
+ # initialize model from wandb artifacts
+ evaluator.initialize_model_from_wandb_artifact("artifact-address-of-your-model-checkpoint")
+ # evaluate
+ evaluator.evaluate()
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-restorers-inference-low-light/keras_restorers_inference_low_light.py b/marimo/convert/keras-restorers-inference-low-light/keras_restorers_inference_low_light.py
new file mode 100644
index 00000000..bb3be4df
--- /dev/null
+++ b/marimo/convert/keras-restorers-inference-low-light/keras_restorers_inference_low_light.py
@@ -0,0 +1,97 @@
+# /// script
+# dependencies = ["pip", "restorers @ git+https://github.com/soumik12345/restorers.git", "setuptools"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌈 Restorers + WandB 🪄🐝
+
+
+
+ This notebook shows how to perform inference with a low-light enhancement using [**restorers**](https://github.com/soumik12345/restorers) and [**wandb**](https://wandb.ai/site). For more details regarding usage of restorers, refer to the following report:
+
+ [](https://wandb.ai/ml-colabs/low-light-enhancement/reports/Lighting-up-Images-in-the-Deep-Learning-Era--VmlldzozNzE4Njkz)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: pip setuptools !pip install -q --upgrade pip setuptools
+ # packages added via marimo's package management: git+https://github.com/soumik12345/restorers.git !pip install git+https://github.com/soumik12345/restorers.git
+ return
+
+
+@app.cell
+def _():
+ import os
+ import wandb
+ from restorers.inference import LowLightInferer
+
+ return LowLightInferer, os, wandb
+
+
+@app.cell
+def _(wandb):
+ # initialize a wandb run for inference
+ wandb.init(project="low-light-enhancement", job_type="inference")
+ return
+
+
+@app.cell
+def _(os, wandb):
+ images_artifact = wandb.use_artifact('ml-colabs/low-light-enhancement/run-7ngsohcn-DarkImagesTable:v0', type='run_table')
+ images_artifact_dir = images_artifact.download()
+ sample_image = os.path.join(images_artifact_dir, "media/images/0b63c6b0cfdfd95675f7/image_9.png")
+ return (sample_image,)
+
+
+@app.cell
+def _(LowLightInferer, sample_image):
+ # initialize the inferer
+ inferer = LowLightInferer(
+ resize_factor=1, model_alias="Zero-DCE"
+ )
+ # intialize the model from wandb artifacts
+ inferer.initialize_model_from_wandb_artifact(
+ # This artifact address corresponds to a Zero-DCE model trained on the LoL dataset
+ "ml-colabs/low-light-enhancement/run_oaa25znm_model:v99"
+ )
+ # infer on a directory of images
+ # inferer.infer("./dark_images")
+ # or infer on a single image
+ inferer.infer(sample_image)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-restorers-train-mirnetv2-restorers/keras_restorers_train_mirnetv2_restorers.py b/marimo/convert/keras-restorers-train-mirnetv2-restorers/keras_restorers_train_mirnetv2_restorers.py
new file mode 100644
index 00000000..022791da
--- /dev/null
+++ b/marimo/convert/keras-restorers-train-mirnetv2-restorers/keras_restorers_train_mirnetv2_restorers.py
@@ -0,0 +1,172 @@
+# /// script
+# dependencies = ["pip", "restorers @ git+https://github.com/soumik12345/restorers.git", "setuptools"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌈 Restorers + WandB 🪄🐝
+
+
+
+ This notebook shows how to train a [MirNetv2](https://www.waqaszamir.com/publication/zamir-2022-mirnetv2/zamir-2022-mirnetv2.pdf) model for low-light enhancement using [**restorers**](https://github.com/soumik12345/restorers) and [**wandb**](https://wandb.ai/site). For more details regarding usage of restorers, refer to the following report:
+
+ [](https://wandb.ai/ml-colabs/low-light-enhancement/reports/Lighting-up-Images-in-the-Deep-Learning-Era--VmlldzozNzE4Njkz)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: pip setuptools !pip install -q --upgrade pip setuptools
+ # packages added via marimo's package management: git+https://github.com/soumik12345/restorers.git !pip install git+https://github.com/soumik12345/restorers.git
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import tensorflow as tf
+ from restorers.dataloader import LOLDataLoader
+
+ return LOLDataLoader, tf, wandb
+
+
+@app.cell
+def _(LOLDataLoader, wandb):
+ wandb.init(project="low-light-enhancement")
+
+ # define dataloader for the LoL dataset
+ data_loader = LOLDataLoader(
+ # size of image crops on which we will train
+ image_size=128,
+ # bit depth of the images
+ bit_depth=8,
+ # fraction of images for validation
+ val_split=0.2,
+ # visualize the dataset on WandB or not
+ visualize_on_wandb=True,
+ # the wandb artifact address of the dataset,
+ # this can be found from the `Usage` tab of
+ # the aforemenioned weave panel
+ dataset_artifact_address="ml-colabs/dataset/LoL:v0",
+ )
+
+ # call `get_datasets` on the `data_loader` to get
+ # the TensorFlow datasets corresponding to the
+ # training and validation splits
+ datasets = data_loader.get_datasets(batch_size=2)
+ train_dataset, val_dataset = datasets
+ return train_dataset, val_dataset
+
+
+@app.cell
+def _():
+ # import MirNetv2 from restorers
+ from restorers.model import MirNetv2
+
+
+ # define the MirNetv2 model; this gives us a `tf.keras.Model`
+ model = MirNetv2(
+ # number of channels in the feature map
+ channels=80,
+ # number of multi-scale residual blocks
+ channel_factor=1.5,
+ # factor by which number of the number of output channels vary
+ num_mrb_blocks=2,
+ # number of groups in which the input is split along the
+ # channel axis in the convolution layers.
+ add_residual_connection=True,
+ )
+ return (model,)
+
+
+@app.cell
+def _(model, tf):
+ from restorers.losses import CharbonnierLoss
+ # import Peak Signal-to-Noise Ratio and Structural Similarity metrics,
+ # implemented as part of restorers
+ from restorers.metrics import PSNRMetric, SSIMMetric
+
+
+ loss = CharbonnierLoss(
+ # a small constant to avoid division by zero
+ epsilon=1e-3,
+ # type of reduction applied to the loss, it needs to be
+ # explicitly specified in case of distributed training
+ reduction=tf.keras.losses.Reduction.SUM,
+ )
+
+
+ optimizer = tf.keras.optimizers.experimental.AdamW(learning_rate=2e-4,)
+
+ psnr_metric = PSNRMetric(max_val=1.0) # peak signal-to-noise ratio metric
+ ssim_metric = SSIMMetric(max_val=1.0) # structural similarity metric
+
+ model.compile(
+ optimizer=optimizer, loss=loss, metrics=[psnr_metric, ssim_metric]
+ )
+ return
+
+
+@app.cell
+def _(model, train_dataset, val_dataset):
+ # import the wandb callbacks for keras
+ from wandb.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+
+ callbacks = [
+ # define the metrics logger callback;
+ # we set the `log_freq="batch"` explicitly
+ # to the metrics are logged both batch-wise and epoch-wise
+ WandbMetricsLogger(log_freq="batch"),
+ # define the model checkpoint callback
+ WandbModelCheckpoint(
+ filepath="checkpoint",
+ monitor="val_loss",
+ save_best_only=False,
+ save_weights_only=False,
+ initial_value_threshold=None,
+ )
+ ]
+
+ # call model.fit()
+ model.fit(
+ train_dataset,
+ validation_data=val_dataset,
+ epochs=50,
+ callbacks=callbacks,
+ )
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-restorers-train-nafnet-restorers/keras_restorers_train_nafnet_restorers.py b/marimo/convert/keras-restorers-train-nafnet-restorers/keras_restorers_train_nafnet_restorers.py
new file mode 100644
index 00000000..fa70e29d
--- /dev/null
+++ b/marimo/convert/keras-restorers-train-nafnet-restorers/keras_restorers_train_nafnet_restorers.py
@@ -0,0 +1,167 @@
+# /// script
+# dependencies = ["pip", "restorers @ git+https://github.com/soumik12345/restorers.git", "setuptools"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌈 Restorers + WandB 🪄🐝
+
+
+
+ This notebook shows how to train a [NAFNet](https://arxiv.org/abs/2204.04676) model for low-light enhancement using [**restorers**](https://github.com/soumik12345/restorers) and [**wandb**](https://wandb.ai/site). For more details regarding usage of restorers, refer to the following report:
+
+ [](https://wandb.ai/ml-colabs/low-light-enhancement/reports/Lighting-up-Images-in-the-Deep-Learning-Era--VmlldzozNzE4Njkz)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: pip setuptools !pip install -q --upgrade pip setuptools
+ # packages added via marimo's package management: git+https://github.com/soumik12345/restorers.git !pip install git+https://github.com/soumik12345/restorers.git
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import tensorflow as tf
+ from restorers.dataloader import LOLDataLoader
+
+ return LOLDataLoader, tf, wandb
+
+
+@app.cell
+def _(LOLDataLoader, wandb):
+ wandb.init(project="low-light-enhancement")
+
+ # define dataloader for the LoL dataset
+ data_loader = LOLDataLoader(
+ # size of image crops on which we will train
+ image_size=128,
+ # bit depth of the images
+ bit_depth=8,
+ # fraction of images for validation
+ val_split=0.2,
+ # visualize the dataset on WandB or not
+ visualize_on_wandb=True,
+ # the wandb artifact address of the dataset,
+ # this can be found from the `Usage` tab of
+ # the aforemenioned weave panel
+ dataset_artifact_address="ml-colabs/dataset/LoL:v0",
+ )
+
+ # call `get_datasets` on the `data_loader` to get
+ # the TensorFlow datasets corresponding to the
+ # training and validation splits
+ datasets = data_loader.get_datasets(batch_size=2)
+ train_dataset, val_dataset = datasets
+ return train_dataset, val_dataset
+
+
+@app.cell
+def _():
+ # import MirNetv2 from restorers
+ from restorers.model import NAFNet
+
+
+ # define the MirNetv2 model; this gives us a `tf.keras.Model`
+ model = NAFNet(
+ filters=16,
+ middle_block_num=1,
+ encoder_block_nums=(1, 1, 1, 1),
+ decoder_block_nums=(1, 1, 1, 1)
+ )
+ return (model,)
+
+
+@app.cell
+def _(model, tf):
+ from restorers.losses import CharbonnierLoss
+ # import Peak Signal-to-Noise Ratio and Structural Similarity metrics,
+ # implemented as part of restorers
+ from restorers.metrics import PSNRMetric, SSIMMetric
+
+
+ loss = CharbonnierLoss(
+ # a small constant to avoid division by zero
+ epsilon=1e-3,
+ # type of reduction applied to the loss, it needs to be
+ # explicitly specified in case of distributed training
+ reduction=tf.keras.losses.Reduction.SUM,
+ )
+
+
+ optimizer = tf.keras.optimizers.experimental.AdamW(learning_rate=2e-4,)
+
+ psnr_metric = PSNRMetric(max_val=1.0) # peak signal-to-noise ratio metric
+ ssim_metric = SSIMMetric(max_val=1.0) # structural similarity metric
+
+ model.compile(
+ optimizer=optimizer, loss=loss, metrics=[psnr_metric, ssim_metric]
+ )
+ return
+
+
+@app.cell
+def _(model, train_dataset, val_dataset):
+ # import the wandb callbacks for keras
+ from wandb.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+
+ callbacks = [
+ # define the metrics logger callback;
+ # we set the `log_freq="batch"` explicitly
+ # to the metrics are logged both batch-wise and epoch-wise
+ WandbMetricsLogger(log_freq="batch"),
+ # define the model checkpoint callback
+ WandbModelCheckpoint(
+ filepath="checkpoint",
+ monitor="val_loss",
+ save_best_only=False,
+ save_weights_only=False,
+ initial_value_threshold=None,
+ )
+ ]
+
+ # call model.fit()
+ model.fit(
+ train_dataset,
+ validation_data=val_dataset,
+ epochs=50,
+ callbacks=callbacks,
+ )
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-restorers-train-zero-dce-restorers/keras_restorers_train_zero_dce_restorers.py b/marimo/convert/keras-restorers-train-zero-dce-restorers/keras_restorers_train_zero_dce_restorers.py
new file mode 100644
index 00000000..b661da56
--- /dev/null
+++ b/marimo/convert/keras-restorers-train-zero-dce-restorers/keras_restorers_train_zero_dce_restorers.py
@@ -0,0 +1,166 @@
+# /// script
+# dependencies = ["pip", "restorers @ git+https://github.com/soumik12345/restorers.git", "setuptools"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌈 Restorers + WandB 🪄🐝
+
+
+
+ This notebook shows how to train a [Zero-DCE](https://arxiv.org/abs/2001.06826) model for zero-reference low-light enhancement using [**restorers**](https://github.com/soumik12345/restorers) and [**wandb**](https://wandb.ai/site). For more details regarding usage of restorers, refer to the following report:
+
+ [](https://wandb.ai/ml-colabs/low-light-enhancement/reports/Lighting-up-Images-in-the-Deep-Learning-Era--VmlldzozNzE4Njkz)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: pip setuptools !pip install -q --upgrade pip setuptools
+ # packages added via marimo's package management: git+https://github.com/soumik12345/restorers.git !pip install git+https://github.com/soumik12345/restorers.git
+ return
+
+
+@app.cell
+def _():
+ import os
+ from glob import glob
+
+ import tensorflow as tf
+
+ import wandb
+ # import the wandb callbacks for keras
+ from wandb.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+ from restorers.model.zero_dce import ZeroDCE
+
+ return (
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ ZeroDCE,
+ glob,
+ os,
+ tf,
+ wandb,
+ )
+
+
+@app.cell
+def _(glob, os, tf, wandb):
+ wandb.init(project="low-light-enhancement", job_type="train")
+
+
+ def load_data(image_path):
+ image = tf.io.read_file(image_path)
+ image = tf.image.decode_png(image, channels=3)
+ image = tf.image.resize(
+ images=image,
+ size=[256, 256]
+ )
+ image = image / ((2 ** 8) - 1)
+ return image
+
+
+ def data_generator(low_light_images):
+ dataset = tf.data.Dataset.from_tensor_slices((low_light_images))
+ dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
+ dataset = dataset.batch(8, drop_remainder=True)
+ return dataset
+
+
+ artifact = wandb.use_artifact("ml-colabs/dataset/LoL:v0", type='dataset')
+ artifact_dir = artifact.download()
+
+ train_low_light_images = sorted(glob(os.path.join(artifact_dir, "our485", "low", "*")))
+ num_train_images = int((1 - 0.2) * len(train_low_light_images))
+ val_low_light_images = train_low_light_images[num_train_images:]
+ train_low_light_images = train_low_light_images[:num_train_images]
+
+ train_dataset = data_generator(train_low_light_images)
+ val_dataset = data_generator(val_low_light_images)
+ return train_dataset, val_dataset
+
+
+@app.cell
+def _(ZeroDCE, tf):
+ # define the ZeroDCE model; this gives us a `tf.keras.Model`
+ model = ZeroDCE(
+ num_intermediate_filters=32, # number of filters in the intermediate convolutional layers
+ num_iterations=8, # number of iterations of enhancement
+ decoder_channel_factor=1 # factor by which number filters in the decoder of deep curve estimation layer is multiplied
+ )
+
+ model.compile(
+ optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
+ weight_exposure_loss=1.0, # weight of the exposure control loss
+ weight_color_constancy_loss=0.5, # weight of the color constancy loss
+ weight_illumination_smoothness_loss=20, # weight of the illumination smoothness loss
+ )
+ return (model,)
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ model,
+ train_dataset,
+ val_dataset,
+):
+ callbacks = [
+ # define the metrics logger callback;
+ # we set the `log_freq="batch"` explicitly
+ # to the metrics are logged both batch-wise and epoch-wise
+ WandbMetricsLogger(log_freq="batch"),
+ # define the model checkpoint callback
+ WandbModelCheckpoint(
+ filepath="checkpoint",
+ monitor="val_loss",
+ save_best_only=False,
+ save_weights_only=False,
+ initial_value_threshold=None,
+ )
+ ]
+
+ # call model.fit()
+ model.fit(
+ train_dataset,
+ validation_data=val_dataset,
+ epochs=50,
+ callbacks=callbacks,
+ )
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-simple-keras-integration/keras_simple_keras_integration.py b/marimo/convert/keras-simple-keras-integration/keras_simple_keras_integration.py
new file mode 100644
index 00000000..19e33bf9
--- /dev/null
+++ b/marimo/convert/keras-simple-keras-integration/keras_simple_keras_integration.py
@@ -0,0 +1,609 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # W&B 💘 Keras
+ Use Weights & Biases for machine learning experiment tracking, model checkpointing, and project collaboration.
+
+
+
+ ## Keras Documentation
+
+ For a full guide of how to log to Weights & Biases using Keras, see **[our Keras documentation](https://docs.wandb.ai/guides/integrations/keras)**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## What this notebook covers:
+
+ We show you how to integrate Weights & Biases with your Keras code to add experiment tracking to your pipeline. That includes:
+
+ 1. Storing hyperparameters and metadata in a `config`.
+ 2. Passing the wandb Keras callbacks to `model.fit`. This will automatically log training metrics, like loss, and system metrics, like GPU and CPU utilization.
+ 3. Using the `wandb.log` API to log custom metrics.
+
+ all using the CIFAR-10 dataset.
+
+ Then, we'll show you how to catch your model making mistakes by logging both the output predictions and the input images the network used to generate them.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Follow along with a [video tutorial](http://wandb.me/keras-video)!
+ **Note**: Sections starting with _Step_ are all you need to integrate W&B in an existing pipeline. The rest just loads data and defines a model.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install, Import, and Log In
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import random
+
+ import matplotlib.pyplot as plt
+ import numpy as np
+ import pandas as pd
+
+ import tensorflow as tf
+ from tensorflow import keras
+ from tensorflow.keras import layers
+ from tensorflow.keras import models
+ from tensorflow.keras.datasets import cifar10
+ import tensorflow_datasets as tfds
+
+ return cifar10, keras, layers, models, np, tf, tfds
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 0: Install W&B
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 1: Import W&B and Login
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.keras import WandbMetricsLogger, WandbModelCheckpoint, WandbEvalCallback
+
+ return WandbEvalCallback, WandbMetricsLogger, WandbModelCheckpoint, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > Side note: If this is your first time using W&B or you are not logged in, the link that appears after running `wandb.login` will take you to sign-up/login page. Signing up is easy!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download and Prepare the Dataset
+ """)
+ return
+
+
+@app.cell
+def _(cifar10):
+ (x_train, y_train), (x_test, y_test) = cifar10.load_data()
+
+ # Subsetting train data and normalizing to [0., 1.]
+ x_train, x_test = x_train[::5] / 255., x_test / 255.
+ y_train = y_train[::5]
+
+ CLASS_NAMES = ["airplane", "automobile", "bird", "cat",
+ "deer", "dog", "frog", "horse", "ship", "truck"]
+
+ print('Shape of x_train: ', x_train.shape)
+ print('Shape of y_train: ', y_train.shape)
+ print('Shape of x_test: ', x_test.shape)
+ print('Shape of y_test: ', y_test.shape)
+ return CLASS_NAMES, x_test, x_train, y_test, y_train
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Define the Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here, we define a standard CNN (with convolution and max-pooling) in Keras.
+ """)
+ return
+
+
+@app.cell
+def _(CLASS_NAMES, keras):
+ def Model():
+ inputs = keras.layers.Input(shape=(32, 32, 3))
+
+ x = keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu')(inputs)
+ x = keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu')(x)
+ x = keras.layers.MaxPooling2D(pool_size=2)(x)
+
+ x = keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu')(x)
+ x = keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu')(x)
+
+ x = keras.layers.GlobalAveragePooling2D()(x)
+
+ x = keras.layers.Dense(128, activation='relu')(x)
+ x = keras.layers.Dense(32, activation='relu')(x)
+
+ outputs = keras.layers.Dense(len(CLASS_NAMES), activation='softmax')(x)
+
+ return keras.models.Model(inputs=inputs, outputs=outputs)
+
+ return (Model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train the Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 2: Give `wandb.init` your `config`
+
+ You first initialize your wandb run, letting us know some training is about to happen. [Check out the official documentation for `.init` here $\rightarrow$](https://docs.wandb.com/library/init)
+
+ That's when you need to set your hyperparameters.
+ They're passed in as a dictionary via the `config` argument,
+ and then become available as the `config` attribute of `wandb`.
+
+ Learn more about `config` in this [Colab Notebook $\rightarrow$](http://wandb.me/config-colab)
+ """)
+ return
+
+
+@app.cell
+def _(Model, tf, wandb):
+ # Initialize wandb with your project name
+ run = wandb.init(project='my-keras-project',
+ config={ # and include hyperparameters and metadata
+ "learning_rate": 0.005,
+ "epochs": 5,
+ "batch_size": 1024,
+ "loss_function": "sparse_categorical_crossentropy",
+ "architecture": "CNN",
+ "dataset": "CIFAR-10"
+ })
+ config = wandb.config # We'll use this to configure our experiment
+
+ # Initialize model like you usually do.
+ tf.keras.backend.clear_session()
+ model = Model()
+ model.summary()
+
+ # Compile model like you usually do.
+ # Notice that we use config, so our metadata matches what gets executed
+ optimizer = tf.keras.optimizers.Adam(config.learning_rate)
+ model.compile(optimizer, config.loss_function, metrics=['acc'])
+ return config, model, run
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 3: Pass `WandbMetricsLogger` and `WandbModelCheckpoint` to `model.fit`
+
+ Keras has a [robust callbacks system](https://keras.io/api/callbacks/) that
+ allows users to separate model definition and the core training logic
+ from other behaviors that occur during training and testing.
+
+ That includes, for example,
+
+ **Click on the Project page link above to see your results!**
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ config,
+ model,
+ x_test,
+ x_train,
+ y_test,
+ y_train,
+):
+ # Add WandbMetricsLogger to log metrics and WandbModelCheckpoint to log model checkpoints
+ _wandb_callbacks = [WandbMetricsLogger(), WandbModelCheckpoint(filepath='my_model_{epoch:02d}')]
+ model.fit(x_train, y_train, epochs=config.epochs, batch_size=config.batch_size, validation_data=(x_test, y_test), callbacks=_wandb_callbacks)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Use `wandb.log` for custom metrics
+
+ Here, we log the error rate on the test set.
+ """)
+ return
+
+
+@app.cell
+def _(model, run, wandb, x_test, y_test):
+ loss, accuracy = model.evaluate(x_test, y_test)
+ print('Test Error Rate: ', round((1 - accuracy) * 100, 2))
+
+ # With wandb.log, we can easily pass in metrics as key-value pairs.
+ wandb.log({'Test Error Rate': round((1 - accuracy) * 100, 2)})
+
+ run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log predictions on test data using `WandbEvalCallback`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The `WandbEvalCallback` is an abstract base class to build Keras callbacks for primarily model prediction visualization and secondarily dataset visualization.
+
+ This is a dataset and task agnostic abstract callback. To use this, inherit from this base callback class and implement the `add_ground_truth` and `add_model_prediction` methods.
+
+ The `WandbEvalCallback` is a utility class that provides helpful methods to:
+
+ - create data and prediction `wandb.Table` instances,
+ - log data and prediction Tables as `wandb.Artifact`,
+ - logs the data table `on_train_begin`,
+ - logs the prediction table `on_epoch_end`.
+
+ As an example, we have implemented `WandbClsEvalCallback` below for an image classification task. This example callback:
+ - logs the validation data (`data_table`) to W&B,
+ - performs inference and logs the prediction (`pred_table`) to W&B on every epoch end.
+
+ ## How the memory footprint is reduced?
+
+ We log the `data_table` to W&B when the `on_train_begin` method is ivoked. Once it's uploaded as a W&B Artifact, we get a reference to this table which can be accessed using `data_table_ref` class variable. The `data_table_ref` is a 2D list that can be indexed like `self.data_table_ref[idx][n]` where `idx` is the row number while `n` is the column number. Let's see the usage in the example below.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Sub-class `WandbEvalCallback`
+ """)
+ return
+
+
+@app.cell
+def _(WandbEvalCallback, logits, np, tf, wandb):
+ class WandbClsEvalCallback(WandbEvalCallback):
+ def __init__(
+ self, validloader, data_table_columns, pred_table_columns, num_samples=100
+ ):
+ super().__init__(data_table_columns, pred_table_columns)
+
+ self.val_data = validloader.unbatch().take(num_samples)
+
+ def add_ground_truth(self, logs=None):
+ for idx, (image, label) in enumerate(self.val_data):
+ self.data_table.add_data(
+ idx,
+ wandb.Image(image),
+ np.argmax(label, axis=-1)
+ )
+
+ def add_model_predictions(self, epoch, logs=None):
+ # Get predictions
+ preds = self._inference()
+ table_idxs = self.data_table_ref.get_index()
+
+ for idx in table_idxs:
+ pred = preds[idx]
+ logit = logits[idx]
+ self.pred_table.add_data(
+ epoch,
+ self.data_table_ref.data[idx][0],
+ self.data_table_ref.data[idx][1],
+ self.data_table_ref.data[idx][2],
+ pred
+ )
+
+ def _inference(self):
+ preds = []
+ for image, label in self.val_data:
+ pred = self.model(tf.expand_dims(image, axis=0))
+ argmax_pred = tf.argmax(pred, axis=-1).numpy()[0]
+ preds.append(argmax_pred)
+
+ return preds
+
+ return (WandbClsEvalCallback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Create Dataset processing and Dataloaders functions
+ """)
+ return
+
+
+@app.cell
+def _(configs, tf):
+ AUTOTUNE = tf.data.AUTOTUNE
+
+ def parse_data(example):
+ # Get image
+ image = example["image"]
+ # image = tf.image.convert_image_dtype(image, dtype=tf.float32)
+
+ # Get label
+ label = example["label"]
+ label = tf.one_hot(label, depth=configs["num_classes"])
+
+ return image, label
+
+
+ def get_dataloader(ds, configs, dataloader_type="train"):
+ dataloader = ds.map(parse_data, num_parallel_calls=AUTOTUNE)
+
+ if dataloader_type=="train":
+ dataloader = dataloader.shuffle(configs["shuffle_buffer"])
+
+ dataloader = (
+ dataloader
+ .batch(configs["batch_size"])
+ .prefetch(AUTOTUNE)
+ )
+
+ return dataloader
+
+ return (get_dataloader,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Define our model
+ """)
+ return
+
+
+@app.cell
+def _(layers, models, tf):
+ def get_model(configs):
+ backbone = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet', include_top=False)
+ backbone.trainable = False
+
+ inputs = layers.Input(shape=(configs["image_size"], configs["image_size"], configs["image_channels"]))
+ resize = layers.Resizing(32, 32)(inputs)
+ neck = layers.Conv2D(3, (3,3), padding="same")(resize)
+ preprocess_input = tf.keras.applications.mobilenet.preprocess_input(neck)
+ x = backbone(preprocess_input)
+ x = layers.GlobalAveragePooling2D()(x)
+ outputs = layers.Dense(configs["num_classes"], activation="softmax")(x)
+
+ return models.Model(inputs=inputs, outputs=outputs)
+
+ return (get_model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Set our config
+ """)
+ return
+
+
+@app.cell
+def _():
+ configs = dict(
+ num_classes = 10,
+ shuffle_buffer = 1024,
+ batch_size = 64,
+ image_size = 28,
+ image_channels = 1,
+ earlystopping_patience = 3,
+ learning_rate = 1e-3,
+ epochs = 10
+ )
+ return (configs,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Dataset
+
+ In this example, we will be using [CIFAR100](https://www.tensorflow.org/datasets/catalog/cifar100) dataset from TensorFlow Dataset catalog. We aim to build a simple image classification pipeline using TensorFlow/Keras.
+ """)
+ return
+
+
+@app.cell
+def _(tfds):
+ train_ds, valid_ds = tfds.load('fashion_mnist', split=['train', 'test'])
+ return train_ds, valid_ds
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Create our Dataloaders and Model
+ """)
+ return
+
+
+@app.cell
+def _(configs, get_dataloader, train_ds, valid_ds):
+ trainloader = get_dataloader(train_ds, configs)
+ validloader = get_dataloader(valid_ds, configs, dataloader_type="valid")
+ return trainloader, validloader
+
+
+@app.cell
+def _(configs, get_model, tf):
+ tf.keras.backend.clear_session()
+ model_1 = get_model(configs)
+ model_1.summary()
+ return (model_1,)
+
+
+@app.cell
+def _(model_1, tf):
+ model_1.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy', tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='top@5_accuracy')])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Train the model and log the predictions to a W&B Table
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbClsEvalCallback,
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ configs,
+ model_1,
+ trainloader,
+ validloader,
+ wandb,
+):
+ run_1 = wandb.init(project='my-keras-project', config=configs)
+ _wandb_callbacks = [WandbMetricsLogger(log_freq=10), WandbModelCheckpoint(filepath='my_model_{epoch:02d}'), WandbClsEvalCallback(validloader, data_table_columns=['idx', 'image', 'ground_truth'], pred_table_columns=['epoch', 'idx', 'image', 'ground_truth', 'prediction'])]
+ model_1.fit(trainloader, epochs=configs['epochs'], validation_data=validloader, callbacks=_wandb_callbacks)
+ run_1.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Click on the **W&B project page** link above to see your live results.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Whats Next? Hyperparameters with Sweeps
+
+ We tried out two different hyperparameter settings by hand. You can use Weights & Biases Sweeps to automate hyperparameter testing and explore the space of possible models and optimization strategies.
+
+ ## [Check out Hyperparameter Optimization in TensorFlow uisng W&B Sweep $\rightarrow$](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/tensorflow/Hyperparameter_Optimization_in_TensorFlow_using_W&B_Sweeps.ipynb)
+
+ 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.com/library/sweeps/configuration) that specifies the parameters to search through, the search strategy, the optimization metric et all.
+
+ 2. **Initialize the sweep:**
+ `sweep_id = wandb.sweep(sweep_config)`
+
+ 3. **Run the sweep agent:**
+ `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.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Documentation
+
+ For a full guide of how to log to Weights & Biases using Keras, see **[our Keras documentation](https://docs.wandb.ai/guides/integrations/keras)**
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-use-wandbevalcallback-in-your-keras-workflow/keras_use_wandbevalcallback_in_your_keras_workflow.py b/marimo/convert/keras-use-wandbevalcallback-in-your-keras-workflow/keras_use_wandbevalcallback_in_your_keras_workflow.py
new file mode 100644
index 00000000..48db1afb
--- /dev/null
+++ b/marimo/convert/keras-use-wandbevalcallback-in-your-keras-workflow/keras_use_wandbevalcallback_in_your_keras_workflow.py
@@ -0,0 +1,372 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using Keras Evaluation Callbacks with Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This colab notebook introduces the `WandbEvalCallback` which is an abstract callback that be inherited to build useful callbacks for model prediction visualization and dataset visualization. Refer to the [💫 `WandbEvalCallback`](https://colab.research.google.com/drive/107uB39vBulCflqmOWolu38noWLxAT6Be#scrollTo=u50GwKJ70WeJ&line=1&uniqifier=1) section for more details.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌴 Setup and Installation
+
+ First, let us install the latest version of Weights and Biases. We will then authenticate this colab instance to use W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qq -U wandb
+ return
+
+
+@app.cell
+def _():
+ import os
+ import numpy as np
+ import tensorflow as tf
+ from tensorflow.keras import layers
+ from tensorflow.keras import models
+ import tensorflow_datasets as tfds
+
+ # Weights and Biases related imports
+ import wandb
+ from wandb.integration.keras import WandbMetricsLogger
+ from wandb.integration.keras import WandbModelCheckpoint
+ from wandb.integration.keras import WandbEvalCallback
+
+ return (
+ WandbEvalCallback,
+ WandbMetricsLogger,
+ layers,
+ models,
+ np,
+ tf,
+ tfds,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If this is your first time using W&B or you are not logged in, the link that appears after running `wandb.login()` will take you to sign-up/login page. Signing up for a [free account](https://wandb.ai/signup) is as easy as a few clicks.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌳 Hyperparameters
+
+ Use of proper config system is a recommended best practice for reproducible machine learning. We can track the hyperparameters for every experiment using W&B. In this colab we will be using simple Python `dict` as our config system.
+ """)
+ return
+
+
+@app.cell
+def _():
+ configs = dict(
+ num_classes = 10,
+ shuffle_buffer = 1024,
+ batch_size = 64,
+ image_size = 28,
+ image_channels = 1,
+ earlystopping_patience = 3,
+ learning_rate = 1e-3,
+ epochs = 10
+ )
+ return (configs,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🍁 Dataset
+
+ In this colab, we will be using [Fashion-MNIST](https://www.tensorflow.org/datasets/catalog/fashion_mnist) dataset from TensorFlow Dataset catalog. We aim to build a simple image classification pipeline using TensorFlow/Keras.
+ """)
+ return
+
+
+@app.cell
+def _(tfds):
+ train_ds, valid_ds = tfds.load('fashion_mnist', split=['train', 'test'])
+ return train_ds, valid_ds
+
+
+@app.cell
+def _(configs, tf):
+ AUTOTUNE = tf.data.AUTOTUNE
+
+
+ def parse_data(example):
+ # Get image
+ image = example["image"]
+ # image = tf.image.convert_image_dtype(image, dtype=tf.float32)
+
+ # Get label
+ label = example["label"]
+ label = tf.one_hot(label, depth=configs["num_classes"])
+
+ return image, label
+
+
+ def get_dataloader(ds, configs, dataloader_type="train"):
+ dataloader = ds.map(parse_data, num_parallel_calls=AUTOTUNE)
+
+ if dataloader_type=="train":
+ dataloader = dataloader.shuffle(configs["shuffle_buffer"])
+
+ dataloader = (
+ dataloader
+ .batch(configs["batch_size"])
+ .prefetch(AUTOTUNE)
+ )
+
+ return dataloader
+
+ return (get_dataloader,)
+
+
+@app.cell
+def _(configs, get_dataloader, train_ds, valid_ds):
+ trainloader = get_dataloader(train_ds, configs)
+ validloader = get_dataloader(valid_ds, configs, dataloader_type="valid")
+ return trainloader, validloader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🎄 Model
+ """)
+ return
+
+
+@app.cell
+def _(layers, models, tf):
+ def get_model(configs):
+ backbone = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet', include_top=False)
+ backbone.trainable = False
+
+ inputs = layers.Input(shape=(configs["image_size"], configs["image_size"], configs["image_channels"]))
+ resize = layers.Resizing(32, 32)(inputs)
+ neck = layers.Conv2D(3, (3,3), padding="same")(resize)
+ preprocess_input = tf.keras.applications.mobilenet.preprocess_input(neck)
+ x = backbone(preprocess_input)
+ x = layers.GlobalAveragePooling2D()(x)
+ outputs = layers.Dense(configs["num_classes"], activation="softmax")(x)
+
+ return models.Model(inputs=inputs, outputs=outputs)
+
+ return (get_model,)
+
+
+@app.cell
+def _(configs, get_model, tf):
+ tf.keras.backend.clear_session()
+ model = get_model(configs)
+ model.summary()
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌿 Compile Model
+ """)
+ return
+
+
+@app.cell
+def _(model, tf):
+ model.compile(
+ optimizer = "adam",
+ loss = "categorical_crossentropy",
+ metrics = ["accuracy", tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='top@5_accuracy')]
+ )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 💫 `WandbEvalCallback`
+
+ The `WandbEvalCallback` is an abstract base class to build Keras callbacks for primarily model prediction visualization and secondarily dataset visualization.
+
+ This is a dataset and task agnostic abstract callback. To use this, inherit from this base callback class and implement the `add_ground_truth` and `add_model_prediction` methods.
+
+ The `WandbEvalCallback` is a utility class that provides helpful methods to:
+
+ - create data and prediction `wandb.Table` instances,
+ - log data and prediction Tables as `wandb.Artifact`,
+ - logs the data table `on_train_begin`,
+ - logs the prediction table `on_epoch_end`.
+
+ As an example, we have implemented `WandbClfEvalCallback` below for an image classification task. This example callback:
+ - logs the validation data (`data_table`) to W&B,
+ - performs inference and logs the prediction (`pred_table`) to W&B on every epoch end.
+
+ ## How the memory footprint is reduced?
+
+ We log the `data_table` to W&B when the `on_train_begin` method is ivoked. Once it's uploaded as a W&B Artifact, we get a reference to this table which can be accessed using `data_table_ref` class variable. The `data_table_ref` is a 2D list that can be indexed like `self.data_table_ref[idx][n]` where `idx` is the row number while `n` is the column number. Let's see the usage in the example below.
+ """)
+ return
+
+
+@app.cell
+def _(WandbEvalCallback, np, tf, wandb):
+ class WandbClfEvalCallback(WandbEvalCallback):
+ def __init__(
+ self, validloader, data_table_columns, pred_table_columns, num_samples=100
+ ):
+ super().__init__(data_table_columns, pred_table_columns)
+
+ self.val_data = validloader.unbatch().take(num_samples)
+
+ def add_ground_truth(self, logs=None):
+ for idx, (image, label) in enumerate(self.val_data):
+ self.data_table.add_data(
+ idx,
+ wandb.Image(image),
+ np.argmax(label, axis=-1)
+ )
+
+ def add_model_predictions(self, epoch, logs=None):
+ # Get predictions
+ preds = self._inference()
+ table_idxs = self.data_table_ref.get_index()
+
+ for idx in table_idxs:
+ pred = preds[idx]
+ self.pred_table.add_data(
+ epoch,
+ self.data_table_ref.data[idx][0],
+ self.data_table_ref.data[idx][1],
+ self.data_table_ref.data[idx][2],
+ pred
+ )
+
+ def _inference(self):
+ preds = []
+ for image, label in self.val_data:
+ pred = self.model(tf.expand_dims(image, axis=0))
+ argmax_pred = tf.argmax(pred, axis=-1).numpy()[0]
+ preds.append(argmax_pred)
+
+ return preds
+
+ return (WandbClfEvalCallback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌻 Train
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbClfEvalCallback,
+ WandbMetricsLogger,
+ configs,
+ model,
+ trainloader,
+ validloader,
+ wandb,
+):
+ # Initialize a W&B run
+ run = wandb.init(
+ project = "intro-keras",
+ config = configs
+ )
+
+ # Train your model
+ model.fit(
+ trainloader,
+ epochs = configs["epochs"],
+ validation_data = validloader,
+ callbacks = [
+ WandbMetricsLogger(log_freq=10),
+ WandbClfEvalCallback(
+ validloader,
+ data_table_columns=["idx", "image", "ground_truth"],
+ pred_table_columns=["epoch", "idx", "image", "ground_truth", "prediction"]
+ ) # Notice the use of WandbEvalCallback here
+ ]
+ )
+
+ # Close the W&B run
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-use-wandbmetriclogger-in-your-keras-workflow/keras_use_wandbmetriclogger_in_your_keras_workflow.py b/marimo/convert/keras-use-wandbmetriclogger-in-your-keras-workflow/keras_use_wandbmetriclogger_in_your_keras_workflow.py
new file mode 100644
index 00000000..79347d67
--- /dev/null
+++ b/marimo/convert/keras-use-wandbmetriclogger-in-your-keras-workflow/keras_use_wandbmetriclogger_in_your_keras_workflow.py
@@ -0,0 +1,273 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using Keras MetricsLogger in your Keras workflow
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This colab notebook introduces the `WandbMetricsLogger` callback. Use this callback for [Experiment Tracking](https://docs.wandb.ai/guides/track). It will log your training and validation metrics along with system metrics to Weights and Biases.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌴 Setup and Installation
+
+ First, let us install the latest version of Weights and Biases. We will then authenticate this colab instance to use W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qq -U wandb
+ return
+
+
+@app.cell
+def _():
+ import os
+ import tensorflow as tf
+ from tensorflow.keras import layers
+ from tensorflow.keras import models
+ import tensorflow_datasets as tfds
+
+ # Weights and Biases related imports
+ import wandb
+ from wandb.integration.keras import WandbMetricsLogger
+
+ return WandbMetricsLogger, layers, models, tf, tfds, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If this is your first time using W&B or you are not logged in, the link that appears after running `wandb.login()` will take you to sign-up/login page. Signing up for a [free account](https://wandb.ai/signup) is as easy as a few clicks.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌳 Hyperparameters
+
+ Use of proper config system is a recommended best practice for reproducible machine learning. We can track the hyperparameters for every experiment using W&B. In this colab we will be using simple Python `dict` as our config system.
+ """)
+ return
+
+
+@app.cell
+def _():
+ configs = dict(
+ num_classes = 10,
+ shuffle_buffer = 1024,
+ batch_size = 64,
+ image_size = 28,
+ image_channels = 1,
+ earlystopping_patience = 3,
+ learning_rate = 1e-3,
+ epochs = 10
+ )
+ return (configs,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🍁 Dataset
+
+ In this colab, we will be using [Fashion-MNIST](https://www.tensorflow.org/datasets/catalog/fashion_mnist) dataset from TensorFlow Dataset catalog. We aim to build a simple image classification pipeline using TensorFlow/Keras.
+ """)
+ return
+
+
+@app.cell
+def _(tfds):
+ train_ds, valid_ds = tfds.load('fashion_mnist', split=['train', 'test'])
+ return train_ds, valid_ds
+
+
+@app.cell
+def _(configs, tf):
+ AUTOTUNE = tf.data.AUTOTUNE
+
+
+ def parse_data(example):
+ # Get image
+ image = example["image"]
+ # image = tf.image.convert_image_dtype(image, dtype=tf.float32)
+
+ # Get label
+ label = example["label"]
+ label = tf.one_hot(label, depth=configs["num_classes"])
+
+ return image, label
+
+
+ def get_dataloader(ds, configs, dataloader_type="train"):
+ dataloader = ds.map(parse_data, num_parallel_calls=AUTOTUNE)
+
+ if dataloader_type=="train":
+ dataloader = dataloader.shuffle(configs["shuffle_buffer"])
+
+ dataloader = (
+ dataloader
+ .batch(configs["batch_size"])
+ .prefetch(AUTOTUNE)
+ )
+
+ return dataloader
+
+ return (get_dataloader,)
+
+
+@app.cell
+def _(configs, get_dataloader, train_ds, valid_ds):
+ trainloader = get_dataloader(train_ds, configs)
+ validloader = get_dataloader(valid_ds, configs, dataloader_type="valid")
+ return trainloader, validloader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🎄 Model
+ """)
+ return
+
+
+@app.cell
+def _(layers, models, tf):
+ def get_model(configs):
+ backbone = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet', include_top=False)
+ backbone.trainable = False
+
+ inputs = layers.Input(shape=(configs["image_size"], configs["image_size"], configs["image_channels"]))
+ resize = layers.Resizing(32, 32)(inputs)
+ neck = layers.Conv2D(3, (3,3), padding="same")(resize)
+ preprocess_input = tf.keras.applications.mobilenet.preprocess_input(neck)
+ x = backbone(preprocess_input)
+ x = layers.GlobalAveragePooling2D()(x)
+ outputs = layers.Dense(configs["num_classes"], activation="softmax")(x)
+
+ return models.Model(inputs=inputs, outputs=outputs)
+
+ return (get_model,)
+
+
+@app.cell
+def _(configs, get_model, tf):
+ tf.keras.backend.clear_session()
+ model = get_model(configs)
+ model.summary()
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌿 Compile Model
+ """)
+ return
+
+
+@app.cell
+def _(model, tf):
+ model.compile(
+ optimizer = "adam",
+ loss = "categorical_crossentropy",
+ metrics = ["accuracy", tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='top@5_accuracy')]
+ )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌻 Train
+ """)
+ return
+
+
+@app.cell
+def _(WandbMetricsLogger, configs, model, trainloader, validloader, wandb):
+ # Initialize a W&B run
+ run = wandb.init(
+ project = "intro-keras",
+ config = configs
+ )
+
+ # Train your model
+ model.fit(
+ trainloader,
+ epochs = configs["epochs"],
+ validation_data = validloader,
+ callbacks = [WandbMetricsLogger(log_freq=10)] # Notice the use of WandbMetricsLogger here
+ )
+
+ # Close the W&B run
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-use-wandbmodelcheckpoint-in-your-keras-workflow/keras_use_wandbmodelcheckpoint_in_your_keras_workflow.py b/marimo/convert/keras-use-wandbmodelcheckpoint-in-your-keras-workflow/keras_use_wandbmodelcheckpoint_in_your_keras_workflow.py
new file mode 100644
index 00000000..a440f607
--- /dev/null
+++ b/marimo/convert/keras-use-wandbmodelcheckpoint-in-your-keras-workflow/keras_use_wandbmodelcheckpoint_in_your_keras_workflow.py
@@ -0,0 +1,293 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using Keras Checkpoint callback with Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This colab notebook introduces the `WandbModelCheckpoint` callback. Use this callback to log your model checkpoints to Weight and Biases [Artifacts](https://docs.wandb.ai/guides/data-and-model-versioning).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌴 Setup and Installation
+
+ First, let us install the latest version of Weights and Biases. We will then authenticate this colab instance to use W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qq -U wandb
+ return
+
+
+@app.cell
+def _():
+ import os
+ import tensorflow as tf
+ from tensorflow.keras import layers
+ from tensorflow.keras import models
+ import tensorflow_datasets as tfds
+
+ # Weights and Biases related imports
+ import wandb
+ from wandb.integration.keras import WandbMetricsLogger
+ from wandb.integration.keras import WandbModelCheckpoint
+
+ return (
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ layers,
+ models,
+ tf,
+ tfds,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If this is your first time using W&B or you are not logged in, the link that appears after running `wandb.login()` will take you to sign-up/login page. Signing up for a [free account](https://wandb.ai/signup) is as easy as a few clicks.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌳 Hyperparameters
+
+ Use of proper config system is a recommended best practice for reproducible machine learning. We can track the hyperparameters for every experiment using W&B. In this colab we will be using simple Python `dict` as our config system.
+ """)
+ return
+
+
+@app.cell
+def _():
+ configs = dict(
+ num_classes = 10,
+ shuffle_buffer = 1024,
+ batch_size = 64,
+ image_size = 28,
+ image_channels = 1,
+ earlystopping_patience = 3,
+ learning_rate = 1e-3,
+ epochs = 10
+ )
+ return (configs,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🍁 Dataset
+
+ In this colab, we will be using [Fashion-MNIST](https://www.tensorflow.org/datasets/catalog/fashion_mnist) dataset from TensorFlow Dataset catalog. We aim to build a simple image classification pipeline using TensorFlow/Keras.
+ """)
+ return
+
+
+@app.cell
+def _(tfds):
+ train_ds, valid_ds = tfds.load('fashion_mnist', split=['train', 'test'])
+ return train_ds, valid_ds
+
+
+@app.cell
+def _(configs, tf):
+ AUTOTUNE = tf.data.AUTOTUNE
+
+
+ def parse_data(example):
+ # Get image
+ image = example["image"]
+ # image = tf.image.convert_image_dtype(image, dtype=tf.float32)
+
+ # Get label
+ label = example["label"]
+ label = tf.one_hot(label, depth=configs["num_classes"])
+
+ return image, label
+
+
+ def get_dataloader(ds, configs, dataloader_type="train"):
+ dataloader = ds.map(parse_data, num_parallel_calls=AUTOTUNE)
+
+ if dataloader_type=="train":
+ dataloader = dataloader.shuffle(configs["shuffle_buffer"])
+
+ dataloader = (
+ dataloader
+ .batch(configs["batch_size"])
+ .prefetch(AUTOTUNE)
+ )
+
+ return dataloader
+
+ return (get_dataloader,)
+
+
+@app.cell
+def _(configs, get_dataloader, train_ds, valid_ds):
+ trainloader = get_dataloader(train_ds, configs)
+ validloader = get_dataloader(valid_ds, configs, dataloader_type="valid")
+ return trainloader, validloader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🎄 Model
+ """)
+ return
+
+
+@app.cell
+def _(layers, models, tf):
+ def get_model(configs):
+ backbone = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet', include_top=False)
+ backbone.trainable = False
+
+ inputs = layers.Input(shape=(configs["image_size"], configs["image_size"], configs["image_channels"]))
+ resize = layers.Resizing(32, 32)(inputs)
+ neck = layers.Conv2D(3, (3,3), padding="same")(resize)
+ preprocess_input = tf.keras.applications.mobilenet.preprocess_input(neck)
+ x = backbone(preprocess_input)
+ x = layers.GlobalAveragePooling2D()(x)
+ outputs = layers.Dense(configs["num_classes"], activation="softmax")(x)
+
+ return models.Model(inputs=inputs, outputs=outputs)
+
+ return (get_model,)
+
+
+@app.cell
+def _(configs, get_model, tf):
+ tf.keras.backend.clear_session()
+ model = get_model(configs)
+ model.summary()
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌿 Compile Model
+ """)
+ return
+
+
+@app.cell
+def _(model, tf):
+ model.compile(
+ optimizer = "adam",
+ loss = "categorical_crossentropy",
+ metrics = ["accuracy", tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='top@5_accuracy')]
+ )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🌻 Train
+ """)
+ return
+
+
+@app.cell
+def _(
+ WandbMetricsLogger,
+ WandbModelCheckpoint,
+ configs,
+ model,
+ trainloader,
+ validloader,
+ wandb,
+):
+ # Initialize a W&B run
+ run = wandb.init(
+ project = "intro-keras",
+ config = configs
+ )
+
+ # Train your model
+ model.fit(
+ trainloader,
+ epochs = configs["epochs"],
+ validation_data = validloader,
+ callbacks = [
+ WandbMetricsLogger(log_freq=10),
+ WandbModelCheckpoint(filepath="models/model.keras") # Notice the use of WandbModelCheckpoint here
+ ]
+ )
+
+ # Close the W&B run
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/keras-zero-dce-colab-keras-video/keras_zero_dce_colab_keras_video.py b/marimo/convert/keras-zero-dce-colab-keras-video/keras_zero_dce_colab_keras_video.py
new file mode 100644
index 00000000..fc666907
--- /dev/null
+++ b/marimo/convert/keras-zero-dce-colab-keras-video/keras_zero_dce_colab_keras_video.py
@@ -0,0 +1,673 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Zero-DCE for low-light image enhancement
+
+ **Author:** [Soumik Rakshit](http://github.com/soumik12345)
+ **Description:** Implementing Zero-Reference Deep Curve Estimation for low-light image enhancement.
+ *Modified for the video by:* [Ivan Goncharov](http://github.com/ivangrov)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Introduction
+
+ **Zero-Reference Deep Curve Estimation** or **Zero-DCE** formulates low-light image
+ enhancement as the task of estimating an image-specific
+ [*tonal curve*](https://en.wikipedia.org/wiki/Curve_(tonality)) with a deep neural network.
+ In this example, we train a lightweight deep network, **DCE-Net**, to estimate
+ pixel-wise and high-order tonal curves for dynamic range adjustment of a given image.
+
+ Zero-DCE takes a low-light image as input and produces high-order tonal curves as its output.
+ These curves are then used for pixel-wise adjustment on the dynamic range of the input to
+ obtain an enhanced image. The curve estimation process is done in such a way that it maintains
+ the range of the enhanced image and preserves the contrast of neighboring pixels. This
+ curve estimation is inspired by curves adjustment used in photo editing software such as
+ Adobe Photoshop where users can adjust points throughout an image’s tonal range.
+
+ Zero-DCE is appealing because of its relaxed assumptions with regard to reference images:
+ it does not require any input/output image pairs during training.
+ This is achieved through a set of carefully formulated non-reference loss functions,
+ which implicitly measure the enhancement quality and guide the training of the network.
+
+ ### References
+
+ - [Zero-Reference Deep Curve Estimation for Low-Light Image Enhancement](https://arxiv.org/pdf/2001.06826.pdf)
+ - [Curves adjustment in Adobe Photoshop](https://helpx.adobe.com/photoshop/using/curves-adjustment.html)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Downloading LOLDataset
+
+ The **LoL Dataset** has been created for low-light image enhancement. It provides 485
+ images for training and 15 for testing. Each image pair in the dataset consists of a
+ low-light input image and its corresponding well-exposed reference image.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.keras import WandbCallback
+
+ return WandbCallback, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Downloading Lol dataset using W&B Artifacts
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import random
+ import numpy as np
+ from glob import glob
+ from PIL import Image, ImageOps
+ import matplotlib.pyplot as plt
+
+ import tensorflow as tf
+ from tensorflow import keras
+ from tensorflow.keras import layers
+
+ return Image, ImageOps, glob, keras, layers, np, os, plt, tf
+
+
+@app.cell
+def _(os, subprocess, wandb):
+ run = wandb.init()
+ artifact = run.use_artifact('ivangoncharov/Low Light Enhancement with Zero-DCE/Lol_Dataset:v0', type='dataset')
+ artifact_dir = artifact.download()
+ artifact_path = os.path.join(artifact_dir, "lol_dataset.zip")
+ #! unzip $artifact_path
+ subprocess.call(['unzip', '$artifact_path'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Creating a TensorFlow Dataset
+
+ We use 300 low-light images from the LoL Dataset training set for training, and we use
+ the remaining 185 low-light images for validation. We resize the images to size `256 x
+ 256` to be used for both training and validation. Note that in order to train the DCE-Net,
+ we will not require the corresponding enhanced images.
+ """)
+ return
+
+
+@app.cell
+def _(glob, tf):
+ IMAGE_SIZE = 256
+ BATCH_SIZE = 16
+ MAX_TRAIN_IMAGES = 400
+
+
+ def load_data(image_path):
+ image = tf.io.read_file(image_path)
+ image = tf.image.decode_png(image, channels=3)
+ image = tf.image.resize(images=image, size=[IMAGE_SIZE, IMAGE_SIZE])
+ image = image / 255.0
+ return image
+
+
+ def data_generator(low_light_images):
+ dataset = tf.data.Dataset.from_tensor_slices((low_light_images))
+ dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
+ dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
+ return dataset
+
+
+ train_low_light_images = sorted(glob("./lol_dataset/our485/low/*"))[:MAX_TRAIN_IMAGES]
+ val_low_light_images = sorted(glob("./lol_dataset/our485/low/*"))[MAX_TRAIN_IMAGES:]
+ test_low_light_images = sorted(glob("./lol_dataset/eval15/low/*"))
+
+
+ train_dataset = data_generator(train_low_light_images)
+ val_dataset = data_generator(val_low_light_images)
+
+ print("Train Dataset:", train_dataset)
+ print("Validation Dataset:", val_dataset)
+ return (
+ test_low_light_images,
+ train_dataset,
+ train_low_light_images,
+ val_dataset,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Visualizing 100 train images from the Lol (low light) dataset
+ """)
+ return
+
+
+@app.cell
+def _(Image, np, train_low_light_images, wandb):
+ wandb.init(project='low_light_zero_DCE', job_type='EDA')
+ _table = wandb.Table(columns=['Low light', 'High light'])
+ for img_path in train_low_light_images[:100]:
+ lowlight_image = Image.open(img_path)
+ highlight_image = Image.open(img_path.replace('low', 'high'))
+ _table.add_data(wandb.Image(np.array(lowlight_image)), wandb.Image(np.array(highlight_image))) #print(img_path)
+ wandb.log({'Dataset table': _table})
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## The Zero-DCE Framework
+
+ The goal of DCE-Net is to estimate a set of best-fitting light-enhancement curves
+ (LE-curves) given an input image. The framework then maps all pixels of the input’s RGB
+ channels by applying the curves iteratively to obtain the final enhanced image.
+
+ ### Understanding light-enhancement curves
+
+ A ligh-enhancement curve is a kind of curve that can map a low-light image
+ to its enhanced version automatically,
+ where the self-adaptive curve parameters are solely dependent on the input image.
+ When designing such a curve, three objectives should be taken into account:
+
+ - Each pixel value of the enhanced image should be in the normalized range `[0,1]`, in order to
+ avoid information loss induced by overflow truncation.
+ - It should be monotonous, to preserve the contrast between neighboring pixels.
+ - The shape of this curve should be as simple as possible,
+ and the curve should be differentiable to allow backpropagation.
+
+ The light-enhancement curve is separately applied to three RGB channels instead of solely on the
+ illumination channel. The three-channel adjustment can better preserve the inherent color and reduce
+ the risk of over-saturation.
+
+ 
+
+ ### DCE-Net
+
+ The DCE-Net is a lightweight deep neural network that learns the mapping between an input
+ image and its best-fitting curve parameter maps. The input to the DCE-Net is a low-light
+ image while the outputs are a set of pixel-wise curve parameter maps for corresponding
+ higher-order curves. It is a plain CNN of seven convolutional layers with symmetrical
+ concatenation. Each layer consists of 32 convolutional kernels of size 3×3 and stride 1
+ followed by the ReLU activation function. The last convolutional layer is followed by the
+ Tanh activation function, which produces 24 parameter maps for 8 iterations, where each
+ iteration requires three curve parameter maps for the three channels.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _(keras, layers):
+ def build_dce_net():
+ input_img = keras.Input(shape=[None, None, 3])
+ conv1 = layers.Conv2D(
+ 32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+ )(input_img)
+ conv2 = layers.Conv2D(
+ 32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+ )(conv1)
+ conv3 = layers.Conv2D(
+ 32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+ )(conv2)
+ conv4 = layers.Conv2D(
+ 32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+ )(conv3)
+ int_con1 = layers.Concatenate(axis=-1)([conv4, conv3])
+ conv5 = layers.Conv2D(
+ 32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+ )(int_con1)
+ int_con2 = layers.Concatenate(axis=-1)([conv5, conv2])
+ conv6 = layers.Conv2D(
+ 32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+ )(int_con2)
+ int_con3 = layers.Concatenate(axis=-1)([conv6, conv1])
+ x_r = layers.Conv2D(24, (3, 3), strides=(1, 1), activation="tanh", padding="same")(
+ int_con3
+ )
+ return keras.Model(inputs=input_img, outputs=x_r)
+
+ return (build_dce_net,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Loss functions
+
+ To enable zero-reference learning in DCE-Net, we use a set of differentiable
+ zero-reference losses that allow us to evaluate the quality of enhanced images.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Color constancy loss
+
+ The *color constancy loss* is used to correct the potential color deviations in the
+ enhanced image.
+ """)
+ return
+
+
+@app.cell
+def _(tf):
+ def color_constancy_loss(x):
+ mean_rgb = tf.reduce_mean(x, axis=(1, 2), keepdims=True)
+ mr, mg, mb = mean_rgb[:, :, :, 0], mean_rgb[:, :, :, 1], mean_rgb[:, :, :, 2]
+ d_rg = tf.square(mr - mg)
+ d_rb = tf.square(mr - mb)
+ d_gb = tf.square(mb - mg)
+ return tf.sqrt(tf.square(d_rg) + tf.square(d_rb) + tf.square(d_gb))
+
+ return (color_constancy_loss,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Exposure loss
+
+ To restrain under-/over-exposed regions, we use the *exposure control loss*.
+ It measures the distance between the average intensity value of a local region
+ and a preset well-exposedness level (set to `0.6`).
+ """)
+ return
+
+
+@app.cell
+def _(tf):
+ def exposure_loss(x, mean_val=0.6):
+ x = tf.reduce_mean(x, axis=3, keepdims=True)
+ mean = tf.nn.avg_pool2d(x, ksize=16, strides=16, padding="VALID")
+ return tf.reduce_mean(tf.square(mean - mean_val))
+
+ return (exposure_loss,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Illumination smoothness loss
+
+ To preserve the monotonicity relations between neighboring pixels, the
+ *illumination smoothness loss* is added to each curve parameter map.
+ """)
+ return
+
+
+@app.cell
+def _(tf):
+ def illumination_smoothness_loss(x):
+ batch_size = tf.shape(x)[0]
+ h_x = tf.shape(x)[1]
+ w_x = tf.shape(x)[2]
+ count_h = (tf.shape(x)[2] - 1) * tf.shape(x)[3]
+ count_w = tf.shape(x)[2] * (tf.shape(x)[3] - 1)
+ h_tv = tf.reduce_sum(tf.square((x[:, 1:, :, :] - x[:, : h_x - 1, :, :])))
+ w_tv = tf.reduce_sum(tf.square((x[:, :, 1:, :] - x[:, :, : w_x - 1, :])))
+ batch_size = tf.cast(batch_size, dtype=tf.float32)
+ count_h = tf.cast(count_h, dtype=tf.float32)
+ count_w = tf.cast(count_w, dtype=tf.float32)
+ return 2 * (h_tv / count_h + w_tv / count_w) / batch_size
+
+ return (illumination_smoothness_loss,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Spatial consistency loss
+
+ The *spatial consistency loss* encourages spatial coherence of the enhanced image by
+ preserving the contrast between neighboring regions across the input image and its enhanced version.
+ """)
+ return
+
+
+@app.cell
+def _(keras, tf):
+ class SpatialConsistencyLoss(keras.losses.Loss):
+ def __init__(self, **kwargs):
+ super(SpatialConsistencyLoss, self).__init__(reduction="none")
+
+ self.left_kernel = tf.constant(
+ [[[[0, 0, 0]], [[-1, 1, 0]], [[0, 0, 0]]]], dtype=tf.float32
+ )
+ self.right_kernel = tf.constant(
+ [[[[0, 0, 0]], [[0, 1, -1]], [[0, 0, 0]]]], dtype=tf.float32
+ )
+ self.up_kernel = tf.constant(
+ [[[[0, -1, 0]], [[0, 1, 0]], [[0, 0, 0]]]], dtype=tf.float32
+ )
+ self.down_kernel = tf.constant(
+ [[[[0, 0, 0]], [[0, 1, 0]], [[0, -1, 0]]]], dtype=tf.float32
+ )
+
+ def call(self, y_true, y_pred):
+
+ original_mean = tf.reduce_mean(y_true, 3, keepdims=True)
+ enhanced_mean = tf.reduce_mean(y_pred, 3, keepdims=True)
+ original_pool = tf.nn.avg_pool2d(
+ original_mean, ksize=4, strides=4, padding="VALID"
+ )
+ enhanced_pool = tf.nn.avg_pool2d(
+ enhanced_mean, ksize=4, strides=4, padding="VALID"
+ )
+
+ d_original_left = tf.nn.conv2d(
+ original_pool, self.left_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+ d_original_right = tf.nn.conv2d(
+ original_pool, self.right_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+ d_original_up = tf.nn.conv2d(
+ original_pool, self.up_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+ d_original_down = tf.nn.conv2d(
+ original_pool, self.down_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+
+ d_enhanced_left = tf.nn.conv2d(
+ enhanced_pool, self.left_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+ d_enhanced_right = tf.nn.conv2d(
+ enhanced_pool, self.right_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+ d_enhanced_up = tf.nn.conv2d(
+ enhanced_pool, self.up_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+ d_enhanced_down = tf.nn.conv2d(
+ enhanced_pool, self.down_kernel, strides=[1, 1, 1, 1], padding="SAME"
+ )
+
+ d_left = tf.square(d_original_left - d_enhanced_left)
+ d_right = tf.square(d_original_right - d_enhanced_right)
+ d_up = tf.square(d_original_up - d_enhanced_up)
+ d_down = tf.square(d_original_down - d_enhanced_down)
+ return d_left + d_right + d_up + d_down
+
+ return (SpatialConsistencyLoss,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Deep curve estimation model
+
+ We implement the Zero-DCE framework as a Keras subclassed model.
+ """)
+ return
+
+
+@app.cell
+def _(
+ SpatialConsistencyLoss,
+ build_dce_net,
+ color_constancy_loss,
+ exposure_loss,
+ illumination_smoothness_loss,
+ keras,
+ tf,
+):
+ class ZeroDCE(keras.Model):
+ def __init__(self, **kwargs):
+ super(ZeroDCE, self).__init__(**kwargs)
+ self.dce_model = build_dce_net()
+
+ def compile(self, learning_rate, **kwargs):
+ super(ZeroDCE, self).compile(**kwargs)
+ self.optimizer = keras.optimizers.Adam(learning_rate=learning_rate)
+ self.spatial_constancy_loss = SpatialConsistencyLoss(reduction="none")
+
+ def get_enhanced_image(self, data, output):
+ r1 = output[:, :, :, :3]
+ r2 = output[:, :, :, 3:6]
+ r3 = output[:, :, :, 6:9]
+ r4 = output[:, :, :, 9:12]
+ r5 = output[:, :, :, 12:15]
+ r6 = output[:, :, :, 15:18]
+ r7 = output[:, :, :, 18:21]
+ r8 = output[:, :, :, 21:24]
+ x = data + r1 * (tf.square(data) - data)
+ x = x + r2 * (tf.square(x) - x)
+ x = x + r3 * (tf.square(x) - x)
+ enhanced_image = x + r4 * (tf.square(x) - x)
+ x = enhanced_image + r5 * (tf.square(enhanced_image) - enhanced_image)
+ x = x + r6 * (tf.square(x) - x)
+ x = x + r7 * (tf.square(x) - x)
+ enhanced_image = x + r8 * (tf.square(x) - x)
+ return enhanced_image
+
+ def call(self, data):
+ dce_net_output = self.dce_model(data)
+ return self.get_enhanced_image(data, dce_net_output)
+
+ def compute_losses(self, data, output):
+ enhanced_image = self.get_enhanced_image(data, output)
+ loss_illumination = 200 * illumination_smoothness_loss(output)
+ loss_spatial_constancy = tf.reduce_mean(
+ self.spatial_constancy_loss(enhanced_image, data)
+ )
+ loss_color_constancy = 5 * tf.reduce_mean(color_constancy_loss(enhanced_image))
+ loss_exposure = 10 * tf.reduce_mean(exposure_loss(enhanced_image))
+ total_loss = (
+ loss_illumination
+ + loss_spatial_constancy
+ + loss_color_constancy
+ + loss_exposure
+ )
+ return {
+ "total_loss": total_loss,
+ "illumination_smoothness_loss": loss_illumination,
+ "spatial_constancy_loss": loss_spatial_constancy,
+ "color_constancy_loss": loss_color_constancy,
+ "exposure_loss": loss_exposure,
+ }
+
+ def train_step(self, data):
+ with tf.GradientTape() as tape:
+ output = self.dce_model(data)
+ losses = self.compute_losses(data, output)
+ gradients = tape.gradient(
+ losses["total_loss"], self.dce_model.trainable_weights
+ )
+ self.optimizer.apply_gradients(zip(gradients, self.dce_model.trainable_weights))
+ return losses
+
+ def test_step(self, data):
+ output = self.dce_model(data)
+ return self.compute_losses(data, output)
+
+ def save_weights(self, filepath, overwrite=True, save_format=None, options=None):
+ """While saving the weights, we simply save the weights of the DCE-Net"""
+ self.dce_model.save_weights(
+ filepath, overwrite=overwrite, save_format=save_format, options=options
+ )
+
+ def load_weights(self, filepath, by_name=False, skip_mismatch=False, options=None):
+ """While loading the weights, we simply load the weights of the DCE-Net"""
+ self.dce_model.load_weights(
+ filepath=filepath,
+ by_name=by_name,
+ skip_mismatch=skip_mismatch,
+ options=options,
+ )
+
+ return (ZeroDCE,)
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project="low_light_zero_DCE", job_type="training")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.config.learning_rate = 1e-4
+ return
+
+
+@app.cell
+def _(WandbCallback, ZeroDCE, plt, train_dataset, val_dataset, wandb):
+ zero_dce_model = ZeroDCE()
+ zero_dce_model.compile(learning_rate=wandb.config.learning_rate)
+ history = zero_dce_model.fit(train_dataset, validation_data=val_dataset, epochs=50, callbacks=[WandbCallback()])
+
+
+ def plot_result(item):
+ plt.plot(history.history[item], label=item)
+ plt.plot(history.history["val_" + item], label="val_" + item)
+ plt.xlabel("Epochs")
+ plt.ylabel(item)
+ plt.title("Train and Validation {} Over Epochs".format(item), fontsize=14)
+ plt.legend()
+ plt.grid()
+ plt.show()
+
+
+ plot_result("total_loss")
+ plot_result("illumination_smoothness_loss")
+ plot_result("spatial_constancy_loss")
+ plot_result("color_constancy_loss")
+ plot_result("exposure_loss")
+ return (zero_dce_model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Inference
+ """)
+ return
+
+
+@app.cell
+def _(Image, keras, np, plt, tf, zero_dce_model):
+ def plot_results(images, titles, figure_size=(12, 12)):
+ fig = plt.figure(figsize=figure_size)
+ for i in range(len(images)):
+ fig.add_subplot(1, len(images), i + 1).set_title(titles[i])
+ _ = plt.imshow(images[i])
+ plt.axis("off")
+ plt.show()
+
+
+ def infer(original_image):
+ image = keras.preprocessing.image.img_to_array(original_image)
+ image = image.astype("float32") / 255.0
+ image = np.expand_dims(image, axis=0)
+ output_image = zero_dce_model(image)
+ output_image = tf.cast((output_image[0, :, :, :] * 255), dtype=np.uint8)
+ output_image = Image.fromarray(output_image.numpy())
+ return output_image
+
+ return (infer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Inference on test images and visualize results using W&B Tables
+ """)
+ return
+
+
+@app.cell
+def _(Image, ImageOps, infer, np, test_low_light_images, wandb):
+ wandb.init(project='low_light_zero_DCE', job_type='predictions')
+ _table = wandb.Table(columns=['Original', 'PIL Autocontrast', 'Enhanced'])
+ for val_image_file in test_low_light_images:
+ original_image = Image.open(val_image_file)
+ enhanced_image = infer(original_image)
+ _table.add_data(wandb.Image(np.array(original_image)), wandb.Image(np.array(ImageOps.autocontrast(original_image))), wandb.Image(np.array(enhanced_image)))
+ wandb.log({'Inference Table': _table})
+ wandb.finish()
+ return
+
+
+@app.cell
+def _(zero_dce_model):
+ zero_dce_model.dce_model.save("model.h5")
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/launch-w-b-launch-quickstart/launch_w_b_launch_quickstart.py b/marimo/convert/launch-w-b-launch-quickstart/launch_w_b_launch_quickstart.py
new file mode 100644
index 00000000..a701cd70
--- /dev/null
+++ b/marimo/convert/launch-w-b-launch-quickstart/launch_w_b_launch_quickstart.py
@@ -0,0 +1,197 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Launch Quickstart
+ Use W&B Launch to quickly edit and re-run previous experiments.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import math
+ import random
+ import torch, torchvision
+ import torch.nn as nn
+ import torchvision.transforms as T
+ from tqdm.auto import tqdm
+
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+ def get_dataloader(is_train, batch_size, slice=5):
+ "Get a training dataloader"
+ full_dataset = torchvision.datasets.MNIST(root=".", train=is_train, transform=T.ToTensor(), download=True)
+ sub_dataset = torch.utils.data.Subset(full_dataset, indices=range(0, len(full_dataset), slice))
+ loader = torch.utils.data.DataLoader(dataset=sub_dataset,
+ batch_size=batch_size,
+ shuffle=True if is_train else False,
+ pin_memory=True, num_workers=2)
+ return loader
+
+ def get_model(dropout):
+ "A simple model"
+ model = nn.Sequential(nn.Flatten(),
+ nn.Linear(28*28, 256),
+ nn.BatchNorm1d(256),
+ nn.ReLU(),
+ nn.Dropout(dropout),
+ nn.Linear(256,10)).to(device)
+ return model
+
+ def validate_model(model, valid_dl, loss_func, log_images=False, batch_idx=0):
+ "Compute performance of the model on the validation dataset and log a wandb.Table"
+ model.eval()
+ val_loss = 0.
+ with torch.inference_mode():
+ correct = 0
+ for i, (images, labels) in tqdm(enumerate(valid_dl), leave=False):
+ images, labels = images.to(device), labels.to(device)
+
+ # Forward pass ➡
+ outputs = model(images)
+ val_loss += loss_func(outputs, labels)*labels.size(0)
+
+ # Compute accuracy and accumulate
+ _, predicted = torch.max(outputs.data, 1)
+ correct += (predicted == labels).sum().item()
+
+ # Log one batch of images to the dashboard, always same batch_idx.
+ if i==batch_idx and log_images:
+ log_image_table(images, predicted, labels, outputs.softmax(dim=1))
+ return val_loss / len(valid_dl.dataset), correct / len(valid_dl.dataset)
+
+ def log_image_table(images, predicted, labels, probs):
+ "Log a wandb.Table with (img, pred, target, scores)"
+ # 🐝 Create a wandb Table to log images, labels and predictions to
+ table = wandb.Table(columns=["image", "pred", "target"]+[f"score_{i}" for i in range(10)])
+ for img, pred, targ, prob in zip(images.to("cpu"), predicted.to("cpu"), labels.to("cpu"), probs.to("cpu")):
+ table.add_data(wandb.Image(img[0].numpy()*255), pred, targ, *prob.numpy())
+ wandb.log({"predictions_table":table}, commit=False)
+
+ return (
+ device,
+ get_dataloader,
+ get_model,
+ math,
+ nn,
+ torch,
+ tqdm,
+ validate_model,
+ wandb,
+ )
+
+
+@app.cell
+def _(
+ device,
+ get_dataloader,
+ get_model,
+ math,
+ nn,
+ torch,
+ tqdm,
+ validate_model,
+ wandb,
+):
+ # 🐝 Initialise a wandb run
+ wandb.init(
+ project="launch-quickstart",
+ config={
+ "epochs": 10,
+ "batch_size": 128,
+ "lr": 1e-3,
+ "dropout": 0.5,
+ }
+ )
+
+ # Copy your config
+ config = wandb.config
+
+ # Get the data
+ train_dl = get_dataloader(is_train=True, batch_size=config.batch_size)
+ valid_dl = get_dataloader(is_train=False, batch_size=2*config.batch_size)
+ n_steps_per_epoch = math.ceil(len(train_dl.dataset) / config.batch_size)
+
+ # A simple MLP model
+ model = get_model(config.dropout)
+
+ # Make the loss and optimizer
+ loss_func = nn.CrossEntropyLoss()
+ optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
+
+ # Training
+ example_ct = 0
+ step_ct = 0
+ for epoch in tqdm(range(config.epochs)):
+ model.train()
+ for step, (images, labels) in enumerate(tqdm(train_dl, leave=False)):
+ images, labels = images.to(device), labels.to(device)
+
+ outputs = model(images)
+ train_loss = loss_func(outputs, labels)
+ optimizer.zero_grad()
+ train_loss.backward()
+ optimizer.step()
+
+ example_ct += len(images)
+ metrics = {"train/train_loss": train_loss,
+ "train/epoch": (step + 1 + (n_steps_per_epoch * epoch)) / n_steps_per_epoch,
+ "train/example_ct": example_ct}
+
+ if step + 1 < n_steps_per_epoch:
+ # 🐝 Log train metrics to wandb
+ wandb.log(metrics)
+
+ step_ct += 1
+
+ val_loss, accuracy = validate_model(model, valid_dl, loss_func, log_images=(epoch==(config.epochs-1)))
+
+ # 🐝 Log train and validation metrics to wandb
+ val_metrics = {"val/val_loss": val_loss,
+ "val/val_accuracy": accuracy}
+ wandb.log({**metrics, **val_metrics})
+
+ print(f"Train Loss: {train_loss:.3f}, Valid Loss: {val_loss:3f}, Accuracy: {accuracy:.2f}")
+
+ # If you had a test set, this is how you could log it as a Summary metric
+ wandb.summary['test_accuracy'] = 0.8
+
+ # 🐝 Close your wandb run
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/llamaindex-llamaindex-local-models/llamaindex_llamaindex_local_models.py b/marimo/convert/llamaindex-llamaindex-local-models/llamaindex_llamaindex_local_models.py
new file mode 100644
index 00000000..77fe2815
--- /dev/null
+++ b/marimo/convert/llamaindex-llamaindex-local-models/llamaindex_llamaindex_local_models.py
@@ -0,0 +1,207 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **TL;DR:** Build a RAG application using llamaindex and local models (embedding + LLM), with [weave](https://wandb.github.io/weave/) for LLM observability
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📦 Packages and Basic Setup
+ ---
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !wget https://controlroom.jurassicoutpost.com/app/uploads/2016/05/JurassicPark-Final.pdf
+ # !pip install -qU llama-index-callbacks-wandb
+ # !pip install -qU llama-index-llms-huggingface
+ # !pip install -qU llama-index-readers-file pymupdf
+ # !pip install -qU llama-index-embeddings-huggingface
+ # !pip install -qU weave ml-collections accelerate
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import weave
+ from llama_index.callbacks.wandb import WandbCallbackHandler
+
+ wandb.login()
+ weave.init("llamaindex-weave-jurassic-qna")
+ wandb_callback = WandbCallbackHandler(
+ run_args={"project": "llamaindex-weave-jurassic-qna"}
+ )
+ return (wandb_callback,)
+
+
+@app.cell
+def _():
+ # @title ⚙️ Configuration
+ import ml_collections
+
+ from llama_index.core import Settings
+
+
+ def get_config() -> ml_collections.ConfigDict:
+ config = ml_collections.ConfigDict()
+ config.model: str = "Writer/camel-5b-hf" # @param {type: "string"}
+ config.embedding_model: str = "BAAI/bge-small-en-v1.5" # @param {type: "string"}
+ config.fetch_index_from_wandb: bool = True # @param {type: "boolean"}
+ config.wandb_entity: str = "sauravmaheshkar" # @param {type: "string"}
+
+ return config
+
+
+ config = get_config()
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💿 The Dataset
+ ---
+
+ In this example, we'll use the original Jurassic Park screenplay to act as our dataset.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from llama_index.core import Document
+ from llama_index.readers.file import PyMuPDFReader
+
+ documents = PyMuPDFReader().load(
+ file_path="/content/JurassicPark-Final.pdf", metadata=True
+ )
+
+ doc_text = "\n\n".join([d.get_content() for d in documents])
+ docs = [Document(text=doc_text)]
+ return (documents,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ✍️ Model Architecture & Training
+ ---
+
+ Since we're using all local models in this example, we'll have to our own Embedding model and llm. In this particular example we'll use "`BAAI/bge-small-en-v1.5`" as our local embedding model and "`Writer/camel-5b-hf`" as the local LLM.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+ #
+ # Settings.embed_model = HuggingFaceEmbedding(model_name=config.embedding_model)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # import torch
+ # from llama_index.core import PromptTemplate
+ # from llama_index.llms.huggingface import HuggingFaceLLM
+ #
+ # query_wrapper_prompt = PromptTemplate(
+ # "Below is an instruction that describes a task. "
+ # "Write a response that appropriately completes the request.\n\n"
+ # "### Instruction:\n{query_str}\n\n### Response:"
+ # )
+ #
+ # Settings.llm = HuggingFaceLLM(
+ # context_window=2048,
+ # max_new_tokens=256,
+ # generate_kwargs={"do_sample": False},
+ # query_wrapper_prompt=query_wrapper_prompt,
+ # tokenizer_name=config.model,
+ # model_name=config.model,
+ # device_map="auto",
+ # tokenizer_kwargs={"max_length": 2048},
+ # model_kwargs={"torch_dtype": torch.float16},
+ # )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🗂 Creating a Index
+ ---
+
+ Based on the value you set for `config.fetch_index_from_wandb` we can either create our own index, or simply download the index stored as an artifact.
+ """)
+ return
+
+
+@app.cell
+def _(config, documents, wandb_callback):
+ from llama_index.core import VectorStoreIndex
+
+ if not config.fetch_index_from_wandb:
+ index = VectorStoreIndex.from_documents(documents)
+ wandb_callback.persist_index(index, index_name="camel-5b-hf-index")
+ return
+
+
+@app.cell
+def _(config, wandb_callback):
+ from llama_index.core import load_index_from_storage
+ if config.fetch_index_from_wandb:
+ storage_context = wandb_callback.load_storage_context(artifact_url='sauravmaheshkar/llamaindex-local-models-index/camel-5b-hf-index:v0')
+ index_1 = load_index_from_storage(storage_context) # Load the index and initialize a query engine
+ return (index_1,)
+
+
+@app.cell
+def _(index_1):
+ query_engine = index_1.as_query_engine()
+ response = query_engine.query('Are Velociraptors pack hunters ?')
+ print(response, sep='\n')
+ return
+
+
+@app.cell
+def _(wandb_callback):
+ wandb_callback.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/mmdetection-train-object-detector-with-mmdetection-and-w-b/mmdetection_train_object_detector_with_mmdetection_and_w_b.py b/marimo/convert/mmdetection-train-object-detector-with-mmdetection-and-w-b/mmdetection_train_object_detector_with_mmdetection_and_w_b.py
new file mode 100644
index 00000000..e8830344
--- /dev/null
+++ b/marimo/convert/mmdetection-train-object-detector-with-mmdetection-and-w-b/mmdetection_train_object_detector_with_mmdetection_and_w_b.py
@@ -0,0 +1,591 @@
+# /// script
+# dependencies = ["-", "https://download-openmmlab-com/mmcv/dist/cu111/torch1-9-0/index-html", "https://download-pytorch-org/whl/torch-stable-html", "mmcv-full", "torch", "torchvision", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 💡 Train an Object Detector with MMDetection and Weights and Biases
+
+ In this colab, we will train an object detector using [MMDetection](https://mmdetection.readthedocs.io/en/latest/1_exist_data_model.html) on a tiny [Kitti](https://paperswithcode.com/dataset/kitti) dataset. Through this colab you will learn to:
+
+ * use MMDetection to train an object detector on a custom dataset,
+ * use [Weights and Biases](https://wandb.ai/site) to log training and validation metrics, visualize model predictions, version raw validation dataset, and more.
+
+ This colab in particular, will showcase a dedicated `MMDetWandbHook` for MMDetection that can be used to:
+
+ ✅ Log training and evaluation metrics.
+ ✅ Log versioned model checkpoints.
+ ✅ Log versioned validation dataset with ground truth bounding boxes.
+ ✅ Log and visualize model predictions.
+
+ But before we continue, here's a quick summary of MMDetection and W&B if you are not familiar with them.
+
+ ### 📸 MMDetection
+
+ MMDetection is an open source object detection toolbox based on PyTorch. It provides composable components that are easy to customize and has out-of-box support for single and multi GPU training/inference. It also has hundreds of pretrained detection models in Model Zoo, and supports multiple standard datasets. Check out the GitHub repository [here](https://github.com/open-mmlab/mmdetection).
+
+ ### 📸 Weights and Biases
+
+ Consider **[Weights and Biases](https://wandb.ai/site)** (W&B) to be the GitHub for machine learning. Use W&B for machine learning experiment tracking, dataset and model versioning, project collaboration, hyperparameter optimization, dataset exploration, model evaluation and so much more. If you are new to W&B, check out this [intro colab](https://wandb.me/intro).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ⚽️ Imports and Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 1️⃣ Install MMDetection
+
+ MMDetection is heavily dependent on the [MMCV](https://mmcv.readthedocs.io/en/latest/#installation) library. We will have to install the version of MMCV that is compatible with the given PyTorch version. Check out the [Installation documentation](https://mmdetection.readthedocs.io/en/latest/get_started.html#installation) for more details.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # install dependencies: (use cu111 because colab has CUDA 11.1)
+ # packages added via marimo's package management: torch==1.9.0+cu111 torchvision==0.10.0+cu111 https://download.pytorch.org/whl/torch_stable.html !pip install -qq torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
+
+ # install mmcv-full thus we could use CUDA operators
+ # packages added via marimo's package management: mmcv-full https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html !pip install -qq mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html
+
+ # Install mmdetection
+ #! rm -rf mmdetection
+ subprocess.call(['rm', '-rf', 'mmdetection'])
+ #! git clone -b wandb2 https://github.com/ayulockin/mmdetection/
+ subprocess.call(['git', 'clone', '-b', 'wandb2', 'https://github.com/ayulockin/mmdetection/'])
+ import os
+ os.chdir('mmdetection')
+
+ # packages added via marimo's package management: . !pip install -e .
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 2️⃣ Install Weights and Biases
+
+ Install the latest version of W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 3️⃣ General Imports
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os.path as osp
+ import torch
+ import torchvision
+ import numpy as np
+ import mmdet
+ print(mmdet.__version__)
+ # MMDetection
+ from mmdet.datasets import build_dataset
+ from mmdet.models import build_detector
+ from mmdet.apis import train_detector
+ from mmdet.datasets.builder import DATASETS
+ from mmdet.datasets.custom import CustomDataset
+ from mmdet.apis import set_random_seed
+ import mmcv
+ from mmcv import Config
+ import wandb
+ # MMCV
+ # Weights and Biases
+ print(wandb.__version__)
+ return (
+ Config,
+ CustomDataset,
+ DATASETS,
+ build_dataset,
+ build_detector,
+ mmcv,
+ np,
+ set_random_seed,
+ train_detector,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 4️⃣ Login with your W&B account
+
+ Create a free W&B account (it's free for personal and academic usage). Create a new API key at [wandb.ai/settings](https://wandb.ai/settings) and store it securely. API keys can only be viewed once when created.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏀 Dataset
+
+ We will be using a tiny KITTI dataset for this colab notebook.
+
+ Even though KITTI is a standard dataset for object detection, tiny KITTI can be considered as a custom dataset (lesser number of classes). MMDetection, recommends to convert the data into COCO or PASCAL VOC formats or the middle format.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 1️⃣ Download the dataset
+ """)
+ return
+
+
+@app.cell
+def _(os, subprocess):
+ os.chdir('../')
+ subprocess.call(['wget', 'https://download.openmmlab.com/mmdetection/data/kitti_tiny.zip'])
+ #! wget https://download.openmmlab.com/mmdetection/data/kitti_tiny.zip
+ #! unzip -q kitti_tiny.zip
+ subprocess.call(['unzip', '-q', 'kitti_tiny.zip'])
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! ls kitti_tiny
+ subprocess.call(['ls', 'kitti_tiny'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > Note: The `training` folder contains both training and validation data samples. This split is determined by the `train.txt` and `val.txt` files.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 2️⃣ Build Custom Dataloader
+
+ To support a new data format, it's recommended to convert the annotations to COCO format or PASCAL VOC format. You can also convert them the "middle format".
+
+ If you are converting annotations to COCO format, do so offline and use the `CocoDataset` class. If you are converting it to the PASCAL format, use the `VOCDataset` class.
+
+ In the example below, we are converting it to the middle format. The `KittiTinyDataset` class will thus inherit the `CustomDataset` class and override the `load_annotations` method.
+
+ You can find more details about customizing the dataset [here](https://mmdetection.readthedocs.io/en/latest/tutorials/customize_dataset.html).
+ """)
+ return
+
+
+@app.cell
+def _(CustomDataset, DATASETS, mmcv, np, os):
+ @DATASETS.register_module()
+ class KittiTinyDataset(CustomDataset):
+
+ CLASSES = ('Car', 'Pedestrian', 'Cyclist')
+
+ def load_annotations(self, ann_file):
+ cat2label = {k: i for i, k in enumerate(self.CLASSES)}
+ # load image list from file
+ image_list = mmcv.list_from_file(self.ann_file)
+
+ data_infos = []
+ # convert annotations to middle format
+ for image_id in image_list:
+ filename = f'{self.img_prefix}/{image_id}.jpeg'
+ image = mmcv.imread(filename)
+ height, width = image.shape[:2]
+
+ data_info = dict(filename=f'{image_id}.jpeg', width=width, height=height)
+
+ # load annotations
+ label_prefix = self.img_prefix.replace('image_2', 'label_2')
+ lines = mmcv.list_from_file(os.path.join(label_prefix, f'{image_id}.txt'))
+
+ content = [line.strip().split(' ') for line in lines]
+ bbox_names = [x[0] for x in content]
+ bboxes = [[float(info) for info in x[4:8]] for x in content]
+
+ gt_bboxes = []
+ gt_labels = []
+ gt_bboxes_ignore = []
+ gt_labels_ignore = []
+
+ # filter 'DontCare'
+ for bbox_name, bbox in zip(bbox_names, bboxes):
+ if bbox_name in cat2label:
+ gt_labels.append(cat2label[bbox_name])
+ gt_bboxes.append(bbox)
+ else:
+ gt_labels_ignore.append(-1)
+ gt_bboxes_ignore.append(bbox)
+
+ data_anno = dict(
+ bboxes=np.array(gt_bboxes, dtype=np.float32).reshape(-1, 4),
+ labels=np.array(gt_labels, dtype=np.long),
+ bboxes_ignore=np.array(gt_bboxes_ignore,
+ dtype=np.float32).reshape(-1, 4),
+ labels_ignore=np.array(gt_labels_ignore, dtype=np.long))
+
+ data_info.update(ann=data_anno)
+ data_infos.append(data_info)
+
+ return data_infos
+
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏈 Model
+
+ There are over hundred pre-trained object detectors provided by MMDetection via Model Zoo. Check out the Model Zoo [documentation](https://mmdetection.readthedocs.io/en/v2.21.0/model_zoo.html) page.
+
+ You can also customize the model's backbone, neck, head, ROI, and loss. More on customizing the model [here](https://mmdetection.readthedocs.io/en/latest/tutorials/customize_models.html).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 1️⃣ Download the model
+
+ We will be using a pretrained model checkpoint to fine tune on our custom dataset. Let's download the model in the `checkpoints` directory.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! mkdir checkpoints
+ subprocess.call(['mkdir', 'checkpoints'])
+ #! wget -c https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco_20210526_095054-1f77628b.pth -O checkpoints/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco_20210526_095054-1f77628b.pth
+ subprocess.call(['wget', '-c', 'https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco_20210526_095054-1f77628b.pth', '-O', 'checkpoints/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco_20210526_095054-1f77628b.pth'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ⚾️ Configuration
+
+ MMDetection relies heavily on a config system. In the cell below, we will be loading a config file and modify few of the methods as per the need of this notebook.
+
+ Note that both train and test dataloaders will use the same training samples. This is not a recommended practice but for the sake of a simplified notebook, let's use it.
+
+ Learn more about the MMDetection Config system [here](https://mmdetection.readthedocs.io/en/latest/tutorials/config.html).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 1️⃣ Load the config file
+ """)
+ return
+
+
+@app.cell
+def _(Config):
+ cfg = Config.fromfile('mmdetection/configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco.py')
+ return (cfg,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 2️⃣ Modify data config
+ """)
+ return
+
+
+@app.cell
+def _(cfg):
+ # Define type and path to the images.
+ cfg.dataset_type = 'KittiTinyDataset'
+ cfg.data_root = 'kitti_tiny/'
+
+ cfg.data.test.type = 'KittiTinyDataset'
+ cfg.data.test.data_root = 'kitti_tiny/'
+ cfg.data.test.ann_file = 'train.txt'
+ cfg.data.test.img_prefix = 'training/image_2'
+
+ cfg.data.train.type = 'KittiTinyDataset'
+ cfg.data.train.data_root = 'kitti_tiny/'
+ cfg.data.train.ann_file = 'train.txt'
+ cfg.data.train.img_prefix = 'training/image_2'
+
+ cfg.data.val.type = 'KittiTinyDataset'
+ cfg.data.val.data_root = 'kitti_tiny/'
+ cfg.data.val.ann_file = 'val.txt'
+ cfg.data.val.img_prefix = 'training/image_2'
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 3️⃣ Modify model config
+ """)
+ return
+
+
+@app.cell
+def _(cfg):
+ # The number of unique objects in the training data.
+ cfg.model.roi_head.bbox_head.num_classes = 3
+ # Use the pretrained model.
+ cfg.load_from = 'checkpoints/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco_20210526_095054-1f77628b.pth'
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 4️⃣ Modify training config
+ """)
+ return
+
+
+@app.cell
+def _(cfg, set_random_seed):
+ # The original learning rate (LR) is set for 8-GPU training.
+ # We divide it by 8 since we only use one GPU.
+ cfg.optimizer.lr = 0.02 / 8
+ cfg.lr_config.warmup = None
+ cfg.log_config.interval = 10
+
+ # Epochs
+ cfg.runner.max_epochs = 12
+
+ # Set seed thus the results are more reproducible
+ cfg.seed = 0
+ set_random_seed(0, deterministic=False)
+ cfg.gpu_ids = range(1)
+
+ # ⭐️ Set the checkpoint interval.
+ cfg.checkpoint_config.interval = 1
+
+ # Set up working dir to save files and logs.
+ cfg.work_dir = './tutorial_exps'
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 5️⃣ Modify evaluation config
+ """)
+ return
+
+
+@app.cell
+def _(cfg):
+ # Change the evaluation metric since we use customized dataset.
+ cfg.evaluation.metric = 'mAP'
+
+ # ⭐️ Set the evaluation interval.
+ cfg.evaluation.interval = 1
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🎾 Define Weights and Biases Hook
+
+ MMDetection comes with a dedicated Weights and Biases Hook - `MMDetWandHook`. MMCV, the parent repository, has a `WandbLoggerHook` that can be used to for basic logging.
+
+ With this dedicated hook, you can:
+
+ * log train and eval metrics along with system (CPU/GPU) metrics,
+ * visualize the validation dataset as interactive [W&B Tables](https://docs.wandb.ai/guides/data-vis),
+ * visualize the model prediction as interactive W&B Tables, and
+ * save the model checkpoints as [W&B Artifacts](https://docs.wandb.ai/guides/artifacts).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To use this hook, you can append a dict to `log_config.hooks`. The `log_config` wraps multiple logger hooks like the `TextLoggerHook` used below.
+
+ There are four important arguments in the `MMDetWandbHook` that can help you get the most out of MMDetection.
+
+ - `init_kwargs`: Use this argument to in-turn pass arguments to `wandb.init`. You can use it to set the W&B project name, set the team name entity if you want to log the runs to a team account, pass the configuration, and more. Check out what all can you pass to `wandb.init` [here](https://docs.wandb.ai/ref/python/init).
+
+ - `log_checkpoint`: The model checkpoints are saved at intervals determined by `checkpoint_config.interval` (starred above). If `log_checkpoint` is `True` the saved checkpoints will be saved as versioned W&B Artifact. Note that this feature is dependent on MMCV's [`CheckpointHook`](https://mmcv.readthedocs.io/en/latest/api.html#mmcv.runner.CheckpointHook).
+
+ - `log_checkpoint_metadata`: If `log_checkpoint_metadata` is True, every checkpoint artifact will have a metadata associated with it. The metadata contains the evaluation metrics computed on validation data with that checkpoint along with the current epoch. If True, it also marks the checkpoint version with the best evaluation metric with a `best` alias. You can choose the best checkpoint in the W&B Artifacts UI using this.
+
+ - `num_eval_images`: At every evaluation interval, the `MMDetWandbHook` logs the model prediction as interactive W&B Tables. The eval interval is determined by `evaluation.interval` (starred above). The number of samples logged is given by `num_eval_images`. The predicted bounding boxes along with the ground truth are logged at every evaluation interval. However, the validation data is logged just once. This Feature is dependent on MMCV's [`EvalHook`](https://mmcv.readthedocs.io/en/latest/api.html#mmcv.runner.EvalHook) or [`DistEvalHook`](https://mmcv.readthedocs.io/en/latest/api.html#mmcv.runner.DistEvalHook).
+ """)
+ return
+
+
+@app.cell
+def _(cfg):
+ cfg.log_config.hooks = [
+ dict(type='TextLoggerHook'),
+ dict(type='MMDetWandbHook',
+ init_kwargs={'project': 'MMDetection-tutorial'},
+ interval=10,
+ log_checkpoint=True,
+ log_checkpoint_metadata=True,
+ num_eval_images=10)]
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏐 Train
+
+ Now that we have the dataset, pretrained model weight, and have defined the configs. Let's stitch them together to train an object detector.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 1️⃣ Build the Dataset
+ """)
+ return
+
+
+@app.cell
+def _(build_dataset, cfg):
+ # Build dataset
+ datasets = [build_dataset(cfg.data.train)]
+ return (datasets,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 2️⃣ Build the Model
+ """)
+ return
+
+
+@app.cell
+def _(build_detector, cfg, datasets):
+ # Build the detector
+ model = build_detector(
+ cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))
+ # Add an attribute for visualization convenience
+ model.CLASSES = datasets[0].CLASSES
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 3️⃣ Train with W&B
+ """)
+ return
+
+
+@app.cell
+def _(cfg, datasets, model, train_detector):
+ # Create work_dir
+ # mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
+ train_detector(model, datasets, cfg, distributed=False, validate=True)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 4️⃣ Notes on using `MMDetWandbHook`.
+
+ Using `MMDetWandbHook` is easy and in most cases it will throw friendly `UserWarning` if something is not quite right. However in the best interest, here are some of things and best practices you should keep in mind:
+
+ * The `MMDetWandbHook` depends on `CheckpointHook` for logging the checkpoints as W&B Artifacts and `EvalHook`/`DistEvalHook` for logging validation data and model predictions. If anyone or both aren't available, this hook will give `UserWarning` and not cause any error.
+
+ * The priority of both `CheckpointHook` and `EvalHook`/`DistEvalHook` should be more than `MMDetWandbHook`.
+
+ * The validation data is logged once as `val_data` W&B Table. The evaluation tables, use reference to this data thus you will not be uploading the same data multiple times.
+
+ * If you want to log the configuration to W&B, pass this key-value pair `'config': cfg._cfg_dict.to_dict()` to `init_kwargs`.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/monai-3d-brain-tumor-segmentation/monai_3d_brain_tumor_segmentation.py b/marimo/convert/monai-3d-brain-tumor-segmentation/monai_3d_brain_tumor_segmentation.py
new file mode 100644
index 00000000..4845757b
--- /dev/null
+++ b/marimo/convert/monai-3d-brain-tumor-segmentation/monai_3d_brain_tumor_segmentation.py
@@ -0,0 +1,886 @@
+# /// script
+# dependencies = ["monai", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Brain tumor 3D segmentation with MONAI and Weights & Biases
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/main/colabs/monai/3d_brain_tumor_segmentation.ipynb)
+
+ This tutorial shows how to construct a training workflow of multi-labels 3D brain tumor segmentation task using [MONAI](https://github.com/Project-MONAI/MONAI) and use experiment tracking and data visualization features of [Weights & Biases](https://wandb.ai/site). The tutorial contains the following features:
+
+ 1. Initialize a Weights & Biases run and synchrozize all configs associated with the run for reproducibility.
+ 2. MONAI transform API:
+ 1. MONAI Transforms for dictionary format data.
+ 2. How to define a new transform according to MONAI `transforms` API.
+ 3. How to randomly adjust intensity for data augmentation.
+ 3. Data Loading and Visualization:
+ 1. Load Nifti image with metadata, load a list of images and stack them.
+ 2. Cache IO and transforms to accelerate training and validation.
+ 3. Visualize the data using `wandb.Table` and interactive segmentation overlay on Weights & Biases.
+ 4. Training a 3D `SegResNet` model
+ 1. Using the `networks`, `losses`, and `metrics` APIs from MONAI.
+ 2. Training the 3D `SegResNet` model using a PyTorch training loop.
+ 3. Track the training experiment using Weights & Biases.
+ 4. Log and version model checkpoints as model artifacts on Weights & Biases.
+ 5. Visualize and compare the predictions on the validation dataset using `wandb.Table` and interactive segmentation overlay on Weights & Biases.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🌴 Setup and Installation
+
+ First, let us install the latest version of both MONAI and Weights and Biases.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: monai[nibabel, tqdm] !python -c "import monai" || pip install -q -U "monai[nibabel, tqdm]"
+ # packages added via marimo's package management: wandb !python -c "import wandb" || pip install -q -U wandb
+ return
+
+
+@app.cell
+def _():
+ import os
+
+ import numpy as np
+ from tqdm.auto import tqdm
+ import wandb
+
+ from monai.apps import DecathlonDataset
+ from monai.data import DataLoader, decollate_batch
+ from monai.losses import DiceLoss
+ from monai.config import print_config
+ from monai.inferers import sliding_window_inference
+ from monai.metrics import DiceMetric
+ from monai.networks.nets import SegResNet
+ from monai.transforms import (
+ Activations,
+ AsDiscrete,
+ Compose,
+ LoadImaged,
+ MapTransform,
+ NormalizeIntensityd,
+ Orientationd,
+ RandFlipd,
+ RandScaleIntensityd,
+ RandShiftIntensityd,
+ RandSpatialCropd,
+ Spacingd,
+ EnsureTyped,
+ EnsureChannelFirstd,
+ )
+ from monai.utils import set_determinism
+
+ import torch
+
+ print_config()
+ return (
+ Activations,
+ AsDiscrete,
+ Compose,
+ DataLoader,
+ DecathlonDataset,
+ DiceLoss,
+ DiceMetric,
+ EnsureChannelFirstd,
+ EnsureTyped,
+ LoadImaged,
+ MapTransform,
+ NormalizeIntensityd,
+ Orientationd,
+ RandFlipd,
+ RandScaleIntensityd,
+ RandShiftIntensityd,
+ RandSpatialCropd,
+ SegResNet,
+ Spacingd,
+ decollate_batch,
+ np,
+ os,
+ set_determinism,
+ sliding_window_inference,
+ torch,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will then authenticate this colab instance to use W&B.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🌳 Initialize a W&B Run
+
+ We will start a new W&B run to start tracking our experiment.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project="monai-brain-tumor-segmentation")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use of proper config system is a recommended best practice for reproducible machine learning. We can track the hyperparameters for every experiment using W&B.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ config = wandb.config
+ config.seed = 0
+ config.roi_size = [224, 224, 144]
+ config.batch_size = 1
+ config.num_workers = 4
+ config.max_train_images_visualized = 20
+ config.max_val_images_visualized = 20
+ config.dice_loss_smoothen_numerator = 0
+ config.dice_loss_smoothen_denominator = 1e-5
+ config.dice_loss_squared_prediction = True
+ config.dice_loss_target_onehot = False
+ config.dice_loss_apply_sigmoid = True
+ config.initial_learning_rate = 1e-4
+ config.weight_decay = 1e-5
+ config.max_train_epochs = 50
+ config.validation_intervals = 1
+ config.dataset_dir = "./dataset/"
+ config.checkpoint_dir = "./checkpoints"
+ config.inference_roi_size = (128, 128, 64)
+ config.max_prediction_images_visualized = 20
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We would also need to set the random seed for modules to enable or disable deterministic training.
+ """)
+ return
+
+
+@app.cell
+def _(config, os, set_determinism):
+ set_determinism(seed=config.seed)
+
+ # Create directories
+ os.makedirs(config.dataset_dir, exist_ok=True)
+ os.makedirs(config.checkpoint_dir, exist_ok=True)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💿 Data Loading and Transformation
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here we use the `monai.transforms` API to create a custom transform that converts the multi-classes labels into multi-labels segmentation task in one-hot format.
+ """)
+ return
+
+
+@app.cell
+def _(MapTransform, torch):
+ class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
+ """
+ Convert labels to multi channels based on brats classes:
+ label 1 is the peritumoral edema
+ label 2 is the GD-enhancing tumor
+ label 3 is the necrotic and non-enhancing tumor core
+ The possible classes are TC (Tumor core), WT (Whole tumor)
+ and ET (Enhancing tumor).
+
+ Reference: https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/brats_segmentation_3d.ipynb
+
+ """
+
+ def __call__(self, data):
+ d = dict(data)
+ for key in self.keys:
+ result = []
+ # merge label 2 and label 3 to construct TC
+ result.append(torch.logical_or(d[key] == 2, d[key] == 3))
+ # merge labels 1, 2 and 3 to construct WT
+ result.append(
+ torch.logical_or(
+ torch.logical_or(d[key] == 2, d[key] == 3), d[key] == 1
+ )
+ )
+ # label 2 is ET
+ result.append(d[key] == 2)
+ d[key] = torch.stack(result, axis=0).float()
+ return d
+
+ return (ConvertToMultiChannelBasedOnBratsClassesd,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we set up transforms for training and validation datasets respectively.
+ """)
+ return
+
+
+@app.cell
+def _(
+ Compose,
+ ConvertToMultiChannelBasedOnBratsClassesd,
+ EnsureChannelFirstd,
+ EnsureTyped,
+ LoadImaged,
+ NormalizeIntensityd,
+ Orientationd,
+ RandFlipd,
+ RandScaleIntensityd,
+ RandShiftIntensityd,
+ RandSpatialCropd,
+ Spacingd,
+ config,
+):
+ train_transform = Compose(
+ [
+ # load 4 Nifti images and stack them together
+ LoadImaged(keys=["image", "label"]),
+ EnsureChannelFirstd(keys="image"),
+ EnsureTyped(keys=["image", "label"]),
+ ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
+ Orientationd(keys=["image", "label"], axcodes="RAS"),
+ Spacingd(
+ keys=["image", "label"],
+ pixdim=(1.0, 1.0, 1.0),
+ mode=("bilinear", "nearest"),
+ ),
+ RandSpatialCropd(
+ keys=["image", "label"], roi_size=config.roi_size, random_size=False
+ ),
+ RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
+ RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
+ RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
+ NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+ RandScaleIntensityd(keys="image", factors=0.1, prob=1.0),
+ RandShiftIntensityd(keys="image", offsets=0.1, prob=1.0),
+ ]
+ )
+ val_transform = Compose(
+ [
+ LoadImaged(keys=["image", "label"]),
+ EnsureChannelFirstd(keys="image"),
+ EnsureTyped(keys=["image", "label"]),
+ ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
+ Orientationd(keys=["image", "label"], axcodes="RAS"),
+ Spacingd(
+ keys=["image", "label"],
+ pixdim=(1.0, 1.0, 1.0),
+ mode=("bilinear", "nearest"),
+ ),
+ NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+ ]
+ )
+ return train_transform, val_transform
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🍁 The Dataset
+
+ The dataset that we will use for this experiment comes from http://medicaldecathlon.com/. We will use Multimodal multisite MRI data (FLAIR, T1w, T1gd, T2w) to segment Gliomas, necrotic/active tumour, and oedema. The dataset consists of 750 4D volumes (484 Training + 266 Testing).
+
+ We will use the `DecathlonDataset` to automatically download and extract the dataset. It inherits MONAI `CacheDataset` which enables us to set `cache_num=N` to cache `N` items for training and use the default args to cache all the items for validation, depending on your memory size.
+ """)
+ return
+
+
+@app.cell
+def _(DecathlonDataset, config, val_transform):
+ train_dataset = DecathlonDataset(
+ root_dir=config.dataset_dir,
+ task="Task01_BrainTumour",
+ transform=val_transform,
+ section="training",
+ download=True,
+ cache_rate=0.0,
+ num_workers=4,
+ )
+ val_dataset = DecathlonDataset(
+ root_dir=config.dataset_dir,
+ task="Task01_BrainTumour",
+ transform=val_transform,
+ section="validation",
+ download=False,
+ cache_rate=0.0,
+ num_workers=4,
+ )
+ return train_dataset, val_dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Note:** Instead of applying the `train_transform` to the `train_dataset`, we have applied `val_transform` to both the training and validation datasets. This is because, before training, we would be visualizing samples from both the splits of the dataset.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 📸 Visualizing the Dataset
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Weights & Biases supports images, video, audio, and more. Log rich media to explore our results and visually compare our runs, models, and datasets. We would be using the [segmentation mask overlay system](https://docs.wandb.ai/guides/track/log/media#image-overlays-in-tables) to visualize our data volumes. To log segmentation masks in [tables](https://docs.wandb.ai/guides/tables), we will need to provide a `wandb.Image`` object for each row in the table.
+
+ An example is provided in the Code snippet below:
+
+ ```python
+ table = wandb.Table(columns=["ID", "Image"])
+
+ for id, img, label in zip(ids, images, labels):
+ mask_img = wandb.Image(
+ img,
+ masks={
+ "prediction": {"mask_data": label, "class_labels": class_labels}
+ # ...
+ },
+ )
+
+ table.add_data(id, img)
+
+ wandb.log({"Table": table})
+ ```
+
+ Let us now write a simple utility function that takes a sample image, label, `wandb.Table` object and some associated metadata and populate the rows of a table that would be logged to our Weights & Biases dashboard.
+ """)
+ return
+
+
+@app.cell
+def _(np, tqdm, wandb):
+ def log_data_samples_into_tables(sample_image: np.array, sample_label: np.array, split: str=None, data_idx: int=None, table: wandb.Table=None):
+ num_channels, _, _, num_slices = sample_image.shape
+ with tqdm(total=num_slices, leave=False) as _progress_bar:
+ for slice_idx in range(num_slices):
+ ground_truth_wandb_images = []
+ for channel_idx in range(num_channels):
+ ground_truth_wandb_images.append(wandb.Image(sample_image[channel_idx, :, :, slice_idx], masks={'ground-truth/Tumor-Core': {'mask_data': sample_label[0, :, :, slice_idx], 'class_labels': {0: 'background', 1: 'Tumor Core'}}, 'ground-truth/Whole-Tumor': {'mask_data': sample_label[1, :, :, slice_idx] * 2, 'class_labels': {0: 'background', 2: 'Whole Tumor'}}, 'ground-truth/Enhancing-Tumor': {'mask_data': sample_label[2, :, :, slice_idx] * 3, 'class_labels': {0: 'background', 3: 'Enhancing Tumor'}}}))
+ table.add_data(split, _data_idx, slice_idx, *ground_truth_wandb_images)
+ _progress_bar.update(1)
+ return table
+
+ return (log_data_samples_into_tables,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we define the `wandb.Table` object and what columns it consists of so that we can populate with our data visualizations.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ table = wandb.Table(
+ columns=[
+ "Split",
+ "Data Index",
+ "Slice Index",
+ "Image-Channel-0",
+ "Image-Channel-1",
+ "Image-Channel-2",
+ "Image-Channel-3",
+ ]
+ )
+ return (table,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Then we loop over the `train_dataset` and `val_dataset` respectively to generate the visualizations for the data samples and populate the rows of the table which we would log to our dashboard.
+ """)
+ return
+
+
+@app.cell
+def _(
+ config,
+ log_data_samples_into_tables,
+ table,
+ tqdm,
+ train_dataset,
+ val_dataset,
+ wandb,
+):
+ _max_samples = min(config.max_train_images_visualized, len(train_dataset)) if config.max_train_images_visualized > 0 else len(train_dataset)
+ _progress_bar = tqdm(enumerate(train_dataset[:_max_samples]), total=_max_samples, desc='Generating Train Dataset Visualizations:')
+ for _data_idx, _sample in _progress_bar:
+ sample_image = _sample['image'].detach().cpu().numpy()
+ sample_label = _sample['label'].detach().cpu().numpy()
+ table_1 = log_data_samples_into_tables(sample_image, sample_label, split='train', data_idx=_data_idx, table=table)
+ _max_samples = min(config.max_val_images_visualized, len(val_dataset)) if config.max_val_images_visualized > 0 else len(val_dataset)
+ _progress_bar = tqdm(enumerate(val_dataset[:_max_samples]), total=_max_samples, desc='Generating Validation Dataset Visualizations:')
+ for _data_idx, _sample in _progress_bar:
+ sample_image = _sample['image'].detach().cpu().numpy()
+ sample_label = _sample['label'].detach().cpu().numpy()
+ table_1 = log_data_samples_into_tables(sample_image, sample_label, split='val', data_idx=_data_idx, table=table_1)
+ wandb.log({'Tumor-Segmentation-Data': table_1})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The data appears to us on our W&B dashboard in an interactive tabular format. We can see each channel of a particular slice from a data volume overlayed with the respective segmentation mask in each row. Let us write [Weave queries](https://docs.wandb.ai/guides/weave) to filter the data on our table and focus on one particular row.
+
+ 
+
+ Let us now open an image and check how we can interact with each of the segmentation masks using the interactive overlay.
+
+ 
+
+ **Note:** The labels in the dataset consist of non-overlapping masks across classes, hence, they were logged as separate masks in the overlay.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🛫 Loading the Data
+
+ We create the PyTorch dataloaders for loading the data from the datasets. Note that before creating the dataloaders, we set the `transform` for `train_dataset` to `train_transform` to preprocess and transform the data for training.
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, config, train_dataset, train_transform, val_dataset):
+ # apply train_transforms to the training dataset
+ train_dataset.transform = train_transform
+
+ # create the train_loader
+ train_loader = DataLoader(
+ train_dataset,
+ batch_size=config.batch_size,
+ shuffle=True,
+ num_workers=config.num_workers,
+ )
+
+ # create the val_loader
+ val_loader = DataLoader(
+ val_dataset,
+ batch_size=config.batch_size,
+ shuffle=False,
+ num_workers=config.num_workers,
+ )
+ return train_loader, val_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🤖 Creating the Model, Loss, and Optimizer
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this tutorial we will be training a `SegResNet` model based on the paper [3D MRI brain tumor segmentation using autoencoder regularization](https://arxiv.org/pdf/1810.11654.pdf). We create the `SegResNet` model that comes implemented as a PyTorch Module as part of the `monai.networks` API. We also create our optimizer and learning rate scheduler.
+ """)
+ return
+
+
+@app.cell
+def _(SegResNet, config, torch):
+ device = torch.device("cuda:0")
+
+ # create model
+ model = SegResNet(
+ blocks_down=[1, 2, 2, 4],
+ blocks_up=[1, 1, 1],
+ init_filters=16,
+ in_channels=4,
+ out_channels=3,
+ dropout_prob=0.2,
+ ).to(device)
+
+ # create optimizer
+ optimizer = torch.optim.Adam(
+ model.parameters(),
+ config.initial_learning_rate,
+ weight_decay=config.weight_decay,
+ )
+
+ # create learning rate scheduler
+ lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+ optimizer, T_max=config.max_train_epochs
+ )
+ return device, lr_scheduler, model, optimizer
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We define our loss as multi-label `DiceLoss` using the `monai.losses` API and the corresponding dice metrics using the `monai.metrics` API.
+ """)
+ return
+
+
+@app.cell
+def _(Activations, AsDiscrete, Compose, DiceLoss, DiceMetric, config, torch):
+ loss_function = DiceLoss(
+ smooth_nr=config.dice_loss_smoothen_numerator,
+ smooth_dr=config.dice_loss_smoothen_denominator,
+ squared_pred=config.dice_loss_squared_prediction,
+ to_onehot_y=config.dice_loss_target_onehot,
+ sigmoid=config.dice_loss_apply_sigmoid,
+ )
+
+ dice_metric = DiceMetric(include_background=True, reduction="mean")
+ dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")
+ post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+
+ # use automatic mixed-precision to accelerate training
+ scaler = torch.cuda.amp.GradScaler()
+ torch.backends.cudnn.benchmark = True
+ return dice_metric, dice_metric_batch, loss_function, post_trans, scaler
+
+
+@app.cell
+def _(sliding_window_inference, torch):
+ def inference(model, input):
+ def _compute(input):
+ return sliding_window_inference(
+ inputs=input,
+ roi_size=(240, 240, 160),
+ sw_batch_size=1,
+ predictor=model,
+ overlap=0.5,
+ )
+
+ with torch.cuda.amp.autocast():
+ return _compute(input)
+
+ return (inference,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🚝 Training and Validation
+
+ Before we start training, let us define some metric properties which will later be logged with `wandb.log()` for tracking our training and validation experiments.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.define_metric("epoch/epoch_step")
+ wandb.define_metric("epoch/*", step_metric="epoch/epoch_step")
+ wandb.define_metric("batch/batch_step")
+ wandb.define_metric("batch/*", step_metric="batch/batch_step")
+ wandb.define_metric("validation/validation_step")
+ wandb.define_metric("validation/*", step_metric="validation/validation_step")
+
+ batch_step = 0
+ validation_step = 0
+ metric_values = []
+ metric_values_tumor_core = []
+ metric_values_whole_tumor = []
+ metric_values_enhanced_tumor = []
+ return (
+ batch_step,
+ metric_values,
+ metric_values_enhanced_tumor,
+ metric_values_tumor_core,
+ metric_values_whole_tumor,
+ validation_step,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🍭 Execute Standard PyTorch Training Loop
+ """)
+ return
+
+
+@app.cell
+def _(
+ batch_step,
+ config,
+ decollate_batch,
+ device,
+ dice_metric,
+ dice_metric_batch,
+ inference,
+ loss_function,
+ lr_scheduler,
+ metric_values,
+ metric_values_enhanced_tumor,
+ metric_values_tumor_core,
+ metric_values_whole_tumor,
+ model,
+ optimizer,
+ os,
+ post_trans,
+ scaler,
+ torch,
+ tqdm,
+ train_dataset,
+ train_loader,
+ val_loader,
+ validation_step,
+ wandb,
+):
+ # Define a W&B Artifact object
+ artifact = wandb.Artifact(name=f'{wandb.run.id}-checkpoint', type='model')
+ epoch_progress_bar = tqdm(range(config.max_train_epochs), desc='Training:')
+ for epoch in epoch_progress_bar:
+ model.train()
+ epoch_loss = 0
+ total_batch_steps = len(train_dataset) // train_loader.batch_size
+ batch_progress_bar = tqdm(train_loader, total=total_batch_steps, leave=False)
+ for batch_data in batch_progress_bar:
+ inputs, labels = (batch_data['image'].to(device), batch_data['label'].to(device))
+ optimizer.zero_grad()
+ with torch.cuda.amp.autocast():
+ outputs = model(inputs)
+ loss = loss_function(outputs, labels)
+ scaler.scale(loss).backward() # Training Step
+ scaler.step(optimizer)
+ scaler.update()
+ epoch_loss = epoch_loss + loss.item()
+ batch_progress_bar.set_description(f'train_loss: {loss.item():.4f}:')
+ wandb.log({'batch/batch_step': batch_step, 'batch/train_loss': loss.item()})
+ batch_step_1 = batch_step + 1
+ lr_scheduler.step()
+ epoch_loss = epoch_loss / total_batch_steps
+ wandb.log({'epoch/epoch_step': epoch, 'epoch/mean_train_loss': epoch_loss, 'epoch/learning_rate': lr_scheduler.get_last_lr()[0]})
+ epoch_progress_bar.set_description(f'Training: train_loss: {epoch_loss:.4f}:')
+ if (epoch + 1) % config.validation_intervals == 0:
+ model.eval()
+ with torch.no_grad():
+ for val_data in val_loader:
+ val_inputs, val_labels = (val_data['image'].to(device), val_data['label'].to(device)) ## Log batch-wise training loss to W&B
+ val_outputs = inference(model, val_inputs)
+ val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
+ dice_metric(y_pred=val_outputs, y=val_labels)
+ dice_metric_batch(y_pred=val_outputs, y=val_labels)
+ metric_values.append(dice_metric.aggregate().item())
+ metric_batch = dice_metric_batch.aggregate() ## Log batch-wise training loss and learning rate to W&B
+ metric_values_tumor_core.append(metric_batch[0].item())
+ metric_values_whole_tumor.append(metric_batch[1].item())
+ metric_values_enhanced_tumor.append(metric_batch[2].item())
+ dice_metric.reset()
+ dice_metric_batch.reset()
+ checkpoint_path = os.path.join(config.checkpoint_dir, 'model.pth')
+ torch.save(model.state_dict(), checkpoint_path)
+ artifact.add_file(local_path=checkpoint_path)
+ wandb.log_artifact(artifact, aliases=[f'epoch_{epoch}'])
+ wandb.log({'validation/validation_step': validation_step, 'validation/mean_dice': metric_values[-1], 'validation/mean_dice_tumor_core': metric_values_tumor_core[-1], 'validation/mean_dice_whole_tumor': metric_values_whole_tumor[-1], 'validation/mean_dice_enhanced_tumor': metric_values_enhanced_tumor[-1]}) # Validation and model checkpointing
+ validation_step_1 = validation_step + 1
+ # Wait for this artifact to finish logging
+ artifact.wait() # Log and versison model checkpoints using W&B artifacts. # Log validation metrics to W&B dashboard.
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Instrumenting our code with `wandb.log` not only enables us to track all the metrics associated with our training and validation process, but also the all system metrics (our CPU and GPU in this case) on our W&B dashboard.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If we navigate to the artifacts tab in the W&B run dashboard, we will be able to access the different versions of model checkpoint artifacts that we logged during training.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🔱 Inference
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Using the artifacts interface, we can select which version of the artifact is the best model checkpoint, in this case, the mean epoch-wise training loss. We can also explore the entire lineage of the artifact and also use the version that we need.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let us fetch the version of the model artifact with the best epoch-wise mean training loss and load the checkpoint state dictionary to the model.
+ """)
+ return
+
+
+@app.cell
+def _(model, os, torch, wandb):
+ model_artifact = wandb.use_artifact(
+ "geekyrakshit/monai-brain-tumor-segmentation/d5ex6n4a-checkpoint:v49",
+ type="model",
+ )
+ model_artifact_dir = model_artifact.download()
+ model.load_state_dict(torch.load(os.path.join(model_artifact_dir, "model.pth")))
+ model.eval()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 📸 Visualizing Predictions and Comparing with the Ground Truth Labels
+
+ In order to visualize the predictions of the pre-trained model and compare them with the corresponding ground-truth segmentation mask using the interactive segmentation mask overlay, let us create another ultility function.
+ """)
+ return
+
+
+@app.cell
+def _(np, tqdm, wandb):
+ def log_predictions_into_tables(sample_image: np.array, sample_label: np.array, predicted_label: np.array, split: str=None, data_idx: int=None, table: wandb.Table=None):
+ num_channels, _, _, num_slices = sample_image.shape
+ with tqdm(total=num_slices, leave=False) as _progress_bar:
+ for slice_idx in range(num_slices):
+ wandb_images = []
+ for channel_idx in range(num_channels):
+ wandb_images = wandb_images + [wandb.Image(sample_image[channel_idx, :, :, slice_idx], masks={'ground-truth/Tumor-Core': {'mask_data': sample_label[0, :, :, slice_idx], 'class_labels': {0: 'background', 1: 'Tumor Core'}}, 'prediction/Tumor-Core': {'mask_data': predicted_label[0, :, :, slice_idx] * 2, 'class_labels': {0: 'background', 2: 'Tumor Core'}}}), wandb.Image(sample_image[channel_idx, :, :, slice_idx], masks={'ground-truth/Whole-Tumor': {'mask_data': sample_label[1, :, :, slice_idx], 'class_labels': {0: 'background', 1: 'Whole Tumor'}}, 'prediction/Whole-Tumor': {'mask_data': predicted_label[1, :, :, slice_idx] * 2, 'class_labels': {0: 'background', 2: 'Whole Tumor'}}}), wandb.Image(sample_image[channel_idx, :, :, slice_idx], masks={'ground-truth/Enhancing-Tumor': {'mask_data': sample_label[2, :, :, slice_idx], 'class_labels': {0: 'background', 1: 'Enhancing Tumor'}}, 'prediction/Enhancing-Tumor': {'mask_data': predicted_label[2, :, :, slice_idx] * 2, 'class_labels': {0: 'background', 2: 'Enhancing Tumor'}}})]
+ table.add_data(split, _data_idx, slice_idx, *wandb_images)
+ _progress_bar.update(1)
+ return table
+
+ return (log_predictions_into_tables,)
+
+
+@app.cell
+def _(
+ config,
+ device,
+ inference,
+ log_predictions_into_tables,
+ model,
+ post_trans,
+ torch,
+ tqdm,
+ val_dataset,
+ wandb,
+):
+ # create the prediction table
+ prediction_table = wandb.Table(columns=['Split', 'Data Index', 'Slice Index', 'Image-Channel-0/Tumor-Core', 'Image-Channel-1/Tumor-Core', 'Image-Channel-2/Tumor-Core', 'Image-Channel-3/Tumor-Core', 'Image-Channel-0/Whole-Tumor', 'Image-Channel-1/Whole-Tumor', 'Image-Channel-2/Whole-Tumor', 'Image-Channel-3/Whole-Tumor', 'Image-Channel-0/Enhancing-Tumor', 'Image-Channel-1/Enhancing-Tumor', 'Image-Channel-2/Enhancing-Tumor', 'Image-Channel-3/Enhancing-Tumor'])
+ with torch.no_grad():
+ config.max_prediction_images_visualized
+ _max_samples = min(config.max_prediction_images_visualized, len(val_dataset)) if config.max_prediction_images_visualized > 0 else len(val_dataset)
+ _progress_bar = tqdm(enumerate(val_dataset[:_max_samples]), total=_max_samples, desc='Generating Predictions:')
+ for _data_idx, _sample in _progress_bar:
+ val_input = _sample['image'].unsqueeze(0).to(device)
+ val_output = inference(model, val_input)
+ val_output = post_trans(val_output[0])
+ prediction_table = log_predictions_into_tables(sample_image=_sample['image'].cpu().numpy(), sample_label=_sample['label'].cpu().numpy(), predicted_label=val_output.cpu().numpy(), data_idx=_data_idx, split='validation', table=prediction_table)
+ wandb.log({'Predictions/Tumor-Segmentation-Data': prediction_table})
+ # Perform inference and visualization
+ # End the experiment
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let us see how we can analyze and compare the predicted segmentation masks and the ground-truth labels for each class using the interactive segmentation mask overlay.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can also check out the report [Brain Tumor Segmentation using MONAI and WandB](https://wandb.ai/geekyrakshit/brain-tumor-segmentation/reports/Brain-Tumor-Segmentation-using-MONAI-and-WandB---Vmlldzo0MjUzODIw) for more details regarding training a brain-tumor segmentation model using MONAI and W&B.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/monai-monai-3d-segmentation-w-b/monai_monai_3d_segmentation_w_b.py b/marimo/convert/monai-monai-3d-segmentation-w-b/monai_monai_3d_segmentation_w_b.py
new file mode 100644
index 00000000..6e4de1e0
--- /dev/null
+++ b/marimo/convert/monai-monai-3d-segmentation-w-b/monai_monai_3d_segmentation_w_b.py
@@ -0,0 +1,732 @@
+# /// script
+# dependencies = ["matplotlib", "monai-weekly", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using MONAI and wandb
+
+ ## Introduction
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This tutorial shows how to integrate MONAI into an existing PyTorch medical DL program and use Weights & Biases for experiment tracking.
+
+ This tutorial is modified from a tutorial from MONAI's official GitHub Repository: [Link](https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d_visualization_basic.ipynb)
+
+ And easily use below features from MONAI:
+
+ - Transforms for dictionary format data.
+ - Load Nifti image with metadata.
+ - Add channel dim to the data if no channel dimension.
+ - Scale medical image intensity with expected range.
+ - Crop out a batch of balanced images based on positive / negative label ratio.
+ - Cache IO and transforms to accelerate training and validation.
+ - 3D UNet model, Dice loss function, Mean Dice metric for 3D segmentation task.
+ - Sliding window inference method.
+ - Deterministic training for reproducibility.
+ - The Spleen dataset can be downloaded from http://medicaldecathlon.com/.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup Environment
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: monai-weekly[gdown, nibabel, tqdm, ignite] !pip install -q "monai-weekly[gdown, nibabel, tqdm, ignite]"
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ # packages added via marimo's package management: matplotlib !pip install -q matplotlib
+ return
+
+
+@app.cell
+def _():
+ import os
+ import glob
+ import shutil
+ import tempfile
+
+ import wandb
+ import torch
+ from torch.optim.lr_scheduler import CosineAnnealingLR
+ import matplotlib.pyplot as plt
+
+ from monai.utils import first, set_determinism
+ from monai.transforms import (
+ AsDiscrete,
+ AsDiscreted,
+ EnsureChannelFirstd,
+ Compose,
+ CropForegroundd,
+ LoadImaged,
+ Orientationd,
+ RandCropByPosNegLabeld,
+ SaveImaged,
+ ScaleIntensityRanged,
+ Spacingd,
+ Invertd,
+ )
+ from monai.handlers.utils import from_engine
+ from monai.networks.nets import UNet
+ from monai.networks.layers import Norm
+ from monai.metrics import DiceMetric
+ from monai.losses import DiceLoss
+ from monai.inferers import sliding_window_inference
+ from monai.data import CacheDataset, DataLoader, Dataset, decollate_batch
+ from monai.config import print_config
+ from monai.apps import download_and_extract
+
+ return (
+ AsDiscrete,
+ CacheDataset,
+ Compose,
+ CosineAnnealingLR,
+ CropForegroundd,
+ DataLoader,
+ Dataset,
+ DiceLoss,
+ DiceMetric,
+ EnsureChannelFirstd,
+ LoadImaged,
+ Norm,
+ Orientationd,
+ RandCropByPosNegLabeld,
+ ScaleIntensityRanged,
+ Spacingd,
+ UNet,
+ decollate_batch,
+ download_and_extract,
+ first,
+ glob,
+ os,
+ plt,
+ print_config,
+ set_determinism,
+ sliding_window_inference,
+ tempfile,
+ torch,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's print configuration of some packages by using a utility function provided by MONAI as `print_config()` which basically lists down all the versions of the useful libraries.
+ """)
+ return
+
+
+@app.cell
+def _(print_config):
+ print_config()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup data directory
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.
+ This allows you to save results and reuse downloads.
+ If not specified a temporary directory will be used.
+ """)
+ return
+
+
+@app.cell
+def _(os, tempfile):
+ # set the environment variable
+ os.environ["MONAI_DATA_DIRECTORY"] = "./output"
+ directory = os.environ.get("MONAI_DATA_DIRECTORY")
+ root_dir = tempfile.mkdtemp() if directory is None else directory
+ print(root_dir)
+ return (root_dir,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download the dataset
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Downloads and extracts the dataset.
+ The dataset comes from http://medicaldecathlon.com/.
+
+ The website has many types of medical datasets like brain tumor, pancreas, heart, prostrate, etc. Here, we are going to use the spleen dataset.
+
+ The spleen is a fist-sized organ in the upper left side of your abdomen, next to your stomach and behind your left ribs.
+
+ It's an important part of your immune system, but you can survive without it. This is because the liver can take over many of the spleen's functions.
+
+ To read more about spleen you can visit [this website](https://www.nhs.uk/conditions/spleen-problems-and-spleen-removal)
+
+ First, we will download the data by specifying the link of the data from the website. Furthermore, we will use a hash value to validate the downloaded file. Finally, we will extract the .tar file. Note, how easy it is to do all of the above steps using the function `download_and_extract`
+ """)
+ return
+
+
+@app.cell
+def _(download_and_extract, os, root_dir):
+ # define the link of the dataset
+ resource = "https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar"
+ # define the hash value to validate the downloaded file
+ md5 = "410d4a301da4e5b2f6f86ec3ddba524e"
+ # define the path for downloading the .tar file
+ compressed_file = os.path.join(root_dir, "Task09_Spleen.tar")
+ # define the directory for extracting the contents of the .tar file
+ data_dir = os.path.join(root_dir, "Task09_Spleen")
+ if not os.path.exists(data_dir):
+ # download, extract and validate the file
+ download_and_extract(resource, compressed_file, root_dir, md5)
+ return (data_dir,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Set MSD Spleen dataset path
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will store the image path and label path as a key value pair in a dictionary and split a subset of data for validation.
+ """)
+ return
+
+
+@app.cell
+def _(data_dir, glob, os):
+ train_images = sorted(
+ glob.glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
+ train_labels = sorted(
+ glob.glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
+ data_dicts = [
+ {"image": image_name, "label": label_name}
+ for image_name, label_name in zip(train_images, train_labels)
+ ]
+ train_files, val_files = data_dicts[:-9], data_dicts[-9:]
+ return train_files, val_files
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Set deterministic training for reproducibility
+ """)
+ return
+
+
+@app.cell
+def _(set_determinism):
+ set_determinism(seed=0)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup transforms for training and validation
+
+ Here we use several transforms to augment the dataset:
+ 1. `LoadImaged` loads the spleen CT images and labels from NIfTI format files.
+ 1. `EnsureChannelFirstd` ensures the original data to construct "channel first" shape.
+ 1. `Orientationd` unifies the data orientation based on the affine matrix.
+ 1. `Spacingd` adjusts the spacing by `pixdim=(1.5, 1.5, 2.)` based on the affine matrix.
+ 1. `ScaleIntensityRanged` extracts intensity range [-57, 164] and scales to [0, 1].
+ 1. `CropForegroundd` removes all zero borders to focus on the valid body area of the images and labels.
+ 1. `RandCropByPosNegLabeld` randomly crop patch samples from big image based on pos / neg ratio.
+ The image centers of negative samples must be in valid body area.
+ 1. `RandAffined` efficiently performs `rotate`, `scale`, `shear`, `translate`, etc. together based on PyTorch affine transform.
+ """)
+ return
+
+
+@app.cell
+def _(
+ Compose,
+ CropForegroundd,
+ EnsureChannelFirstd,
+ LoadImaged,
+ Orientationd,
+ RandCropByPosNegLabeld,
+ ScaleIntensityRanged,
+ Spacingd,
+):
+ train_transforms = Compose(
+ [
+ LoadImaged(keys=["image", "label"]),
+ EnsureChannelFirstd(keys=["image", "label"]),
+ ScaleIntensityRanged(
+ keys=["image"], a_min=-57, a_max=164,
+ b_min=0.0, b_max=1.0, clip=True,
+ ),
+ CropForegroundd(keys=["image", "label"], source_key="image"),
+ Orientationd(keys=["image", "label"], axcodes="RAS"),
+ Spacingd(keys=["image", "label"], pixdim=(
+ 1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
+ RandCropByPosNegLabeld(
+ keys=["image", "label"],
+ label_key="label",
+ spatial_size=(96, 96, 96),
+ pos=1,
+ neg=1,
+ num_samples=4,
+ image_key="image",
+ image_threshold=0,
+ ),
+ # user can also add other random transforms
+ # RandAffined(
+ # keys=['image', 'label'],
+ # mode=('bilinear', 'nearest'),
+ # prob=1.0, spatial_size=(96, 96, 96),
+ # rotate_range=(0, 0, np.pi/15),
+ # scale_range=(0.1, 0.1, 0.1)),
+ ]
+ )
+ val_transforms = Compose(
+ [
+ LoadImaged(keys=["image", "label"]),
+ EnsureChannelFirstd(keys=["image", "label"]),
+ ScaleIntensityRanged(
+ keys=["image"], a_min=-57, a_max=164,
+ b_min=0.0, b_max=1.0, clip=True,
+ ),
+ CropForegroundd(keys=["image", "label"], source_key="image"),
+ Orientationd(keys=["image", "label"], axcodes="RAS"),
+ Spacingd(keys=["image", "label"], pixdim=(
+ 1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
+ ]
+ )
+ return train_transforms, val_transforms
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Check DataLoader
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we will plot a single slice from the first 3D image from the dataloader along with it's label to see if it is loaded and transformed correctly.
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, Dataset, first, plt, val_files, val_transforms):
+ check_ds = Dataset(data=val_files, transform=val_transforms)
+ check_loader = DataLoader(check_ds, batch_size=1)
+ check_data = first(check_loader)
+ image, _label = (check_data['image'][0][0], check_data['label'][0][0])
+ print(f'image shape: {image.shape}, label shape: {_label.shape}')
+ # plot the slice [:, :, 80]
+ plt.figure('check', (12, 6))
+ plt.subplot(1, 2, 1)
+ plt.title('image')
+ plt.imshow(image[:, :, 80], cmap='gray')
+ plt.subplot(1, 2, 2)
+ plt.title('label')
+ plt.imshow(_label[:, :, 80])
+ plt.show()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Great, now we will create a function which will log all the slices of the 3D image to W&B to visualize them interactively. Furthermore, we will also log the slices with segmentation masks to see the overlayed view of segmentations masks on the slices interactively in the W&B dashboard.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Logging spleen slices to W&B
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, Dataset, first, train_files, val_transforms, wandb):
+ # utility function for generating interactive image mask from components
+ def wb_mask(bg_img, mask):
+ return wandb.Image(bg_img, masks={'ground truth': {'mask_data': mask, 'class_labels': {0: 'background', 1: 'mask'}}})
+
+ def log_spleen_slices(total_slices=100):
+ wandb_mask_logs = []
+ wandb_img_logs = []
+ check_ds = Dataset(data=train_files, transform=val_transforms)
+ check_loader = DataLoader(check_ds, batch_size=1)
+ check_data = first(check_loader)
+ image, _label = (check_data['image'][0][0], check_data['label'][0][0])
+ for img_slice_no in range(total_slices):
+ img = image[:, :, img_slice_no] # get the first item of the dataloader
+ lbl = _label[:, :, img_slice_no]
+ wandb_img_logs.append(wandb.Image(img, caption=f'Slice: {img_slice_no}'))
+ wandb_mask_logs.append(wb_mask(img, lbl))
+ wandb.log({'Image': wandb_img_logs})
+ wandb.log({'Segmentation mask': wandb_mask_logs}) # append the image to wandb_img_list to visualize # the slices interactively in W&B dashboard # append the image and masks to wandb_mask_logs # to see the masks overlayed on the original image
+
+ return (log_spleen_slices,)
+
+
+@app.cell
+def _(log_spleen_slices, wandb):
+ # 🐝 init wandb with appropiate project and run name
+ wandb.init(project="MONAI_Spleen_3D_Segmentation", name="slice_image_exploration")
+ # 🐝 log images to W&B
+ log_spleen_slices(total_slices=100)
+ # 🐝 finish the run
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Define Configuration
+
+ Here, we define the configuration for dataloaders, models, train settings in a dictionary. Note that this config object would be passed to `wandb.init()` method to log all the necessary parameters that went into the experiment.
+ """)
+ return
+
+
+@app.cell
+def _(Norm):
+ config = {
+ # data
+ "cache_rate": 1.0,
+ "num_workers": 2,
+
+
+ # train settings
+ "train_batch_size": 2,
+ "val_batch_size": 1,
+ "learning_rate": 1e-3,
+ "max_epochs": 100,
+ "val_interval": 10, # check validation score after n epochs
+ "lr_scheduler": "cosine_decay", # just to keep track
+
+
+
+
+ # Unet model (you can even use nested dictionary and this will be handled by W&B automatically)
+ "model_type": "unet", # just to keep track
+ "model_params": dict(spatial_dims=3,
+ in_channels=1,
+ out_channels=2,
+ channels=(16, 32, 64, 128, 256),
+ strides=(2, 2, 2, 2),
+ num_res_units=2,
+ norm=Norm.BATCH,
+ )
+ }
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Define CacheDataset and DataLoader for training and validation
+
+ Here we use `CacheDataset` to accelerate training and validation process, it's 10x faster than the regular Dataset.
+ To achieve best performance, set `cache_rate=1.0` to cache all the data, if memory is not enough, set lower value.
+ Users can also set `cache_num` instead of `cache_rate`, will use the minimum value of the 2 settings.
+ And set `num_workers` to enable multi-threads during caching.
+ If want to to try the regular Dataset, just change to use the commented code below.
+ """)
+ return
+
+
+@app.cell
+def _(
+ CacheDataset,
+ DataLoader,
+ config,
+ train_files,
+ train_transforms,
+ val_files,
+ val_transforms,
+):
+ train_ds = CacheDataset(
+ data=train_files, transform=train_transforms,
+ cache_rate=config['cache_rate'], num_workers=config['num_workers'])
+ # train_ds = Dataset(data=train_files, transform=train_transforms)
+
+ # use batch_size=2 to load images and use RandCropByPosNegLabeld
+ # to generate 2 x 4 images for network training
+ train_loader = DataLoader(train_ds, batch_size=config['train_batch_size'], shuffle=True, num_workers=config['num_workers'])
+
+ val_ds = CacheDataset(
+ data=val_files, transform=val_transforms, cache_rate=config['cache_rate'], num_workers=config['num_workers'])
+ # val_ds = Dataset(data=val_files, transform=val_transforms)
+ val_loader = DataLoader(val_ds, batch_size=config['val_batch_size'], num_workers=config['num_workers'])
+ return train_ds, train_loader, val_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Create Model, Loss, Optimizer and Scheduler
+ """)
+ return
+
+
+@app.cell
+def _(CosineAnnealingLR, DiceLoss, DiceMetric, UNet, config, torch):
+ # standard PyTorch program style: create UNet, DiceLoss and Adam optimizer
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ model = UNet(**config['model_params']).to(device)
+ loss_function = DiceLoss(to_onehot_y=True, softmax=True)
+ optimizer = torch.optim.Adam(model.parameters(), lr=config['learning_rate'])
+ dice_metric = DiceMetric(include_background=False, reduction="mean")
+ scheduler = CosineAnnealingLR(optimizer, T_max=config['max_epochs'], eta_min=1e-9)
+ return device, dice_metric, loss_function, model, optimizer, scheduler
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Execute a typical PyTorch training process
+ """)
+ return
+
+
+@app.cell
+def _(
+ AsDiscrete,
+ Compose,
+ config,
+ decollate_batch,
+ device,
+ dice_metric,
+ loss_function,
+ model,
+ optimizer,
+ os,
+ root_dir,
+ scheduler,
+ sliding_window_inference,
+ torch,
+ train_ds,
+ train_loader,
+ val_loader,
+ wandb,
+):
+ # 🐝 initialize a wandb run
+ wandb.init(project='MONAI_Spleen_3D_Segmentation', config=config)
+ wandb.watch(model, log_freq=100)
+ max_epochs = config['max_epochs']
+ val_interval = config['val_interval']
+ best_metric = -1
+ # 🐝 log gradients of the model to wandb
+ best_metric_epoch = -1
+ epoch_loss_values = []
+ metric_values = []
+ post_pred = Compose([AsDiscrete(argmax=True, to_onehot=2)])
+ post_label = Compose([AsDiscrete(to_onehot=2)])
+ for epoch in range(max_epochs):
+ print('-' * 10)
+ print(f'epoch {epoch + 1}/{max_epochs}')
+ model.train()
+ epoch_loss = 0
+ step = 0
+ for batch_data in train_loader:
+ step += 1
+ inputs, labels = (batch_data['image'].to(device), batch_data['label'].to(device))
+ optimizer.zero_grad()
+ outputs = model(inputs)
+ loss = loss_function(outputs, labels)
+ loss.backward()
+ optimizer.step()
+ epoch_loss += loss.item()
+ print(f'{step}/{len(train_ds) // train_loader.batch_size}, train_loss: {loss.item():.4f}')
+ wandb.log({'train/loss': loss.item()})
+ epoch_loss /= step
+ epoch_loss_values.append(epoch_loss)
+ print(f'epoch {epoch + 1} average loss: {epoch_loss:.4f}')
+ scheduler.step()
+ wandb.log({'train/loss_epoch': epoch_loss})
+ wandb.log({'learning_rate': scheduler.get_lr()[0]})
+ if (epoch + 1) % val_interval == 0:
+ model.eval()
+ with torch.no_grad():
+ for _val_data in val_loader:
+ val_inputs, val_labels = (_val_data['image'].to(device), _val_data['label'].to(device))
+ _roi_size = (160, 160, 160) # 🐝 log train_loss for each step to wandb
+ _sw_batch_size = 4
+ _val_outputs = sliding_window_inference(val_inputs, _roi_size, _sw_batch_size, model)
+ _val_outputs = [post_pred(_i) for _i in decollate_batch(_val_outputs)]
+ val_labels = [post_label(_i) for _i in decollate_batch(val_labels)]
+ dice_metric(y_pred=_val_outputs, y=val_labels)
+ metric = dice_metric.aggregate().item()
+ wandb.log({'val/dice_metric': metric}) # step scheduler after each epoch (cosine decay)
+ dice_metric.reset()
+ metric_values.append(metric)
+ if metric > best_metric: # 🐝 log train_loss averaged over epoch to wandb
+ best_metric = metric
+ best_metric_epoch = epoch + 1
+ torch.save(model.state_dict(), os.path.join(root_dir, 'best_metric_model.pth')) # 🐝 log learning rate after each epoch to wandb
+ print('saved new best metric model')
+ print(f'current epoch: {epoch + 1} current mean dice: {metric:.4f}\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}')
+ print(f'\ntrain completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}')
+ wandb.log({'best_dice_metric': best_metric, 'best_metric_epoch': best_metric_epoch})
+ best_model_path = os.path.join(root_dir, 'best_metric_model.pth')
+ model_artifact = wandb.Artifact('unet', type='model', description='Unet for 3D Segmentation of spleen', metadata=dict(config['model_params']))
+ model_artifact.add_file(best_model_path)
+ # 🐝 log best score and epoch number to wandb
+ # 🐝 Version your model
+ wandb.log_artifact(model_artifact) # compute metric for current iteration # 🐝 aggregate the final mean dice result # 🐝 log validation dice score for each validation round # reset the status for next validation round
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Check best model output with the input image and label
+ """)
+ return
+
+
+@app.cell
+def _(
+ device,
+ model,
+ os,
+ plt,
+ root_dir,
+ sliding_window_inference,
+ torch,
+ val_loader,
+):
+ model.load_state_dict(torch.load(os.path.join(root_dir, 'best_metric_model.pth')))
+ model.eval()
+ with torch.no_grad():
+ for _i, _val_data in enumerate(val_loader):
+ _roi_size = (160, 160, 160)
+ _sw_batch_size = 4
+ _val_outputs = sliding_window_inference(_val_data['image'].to(device), _roi_size, _sw_batch_size, model)
+ plt.figure('check', (18, 6))
+ plt.subplot(1, 3, 1)
+ plt.title(f'image {_i}')
+ plt.imshow(_val_data['image'][0, 0, :, :, 80], cmap='gray') # plot the slice [:, :, 80]
+ plt.subplot(1, 3, 2)
+ plt.title(f'label {_i}')
+ plt.imshow(_val_data['label'][0, 0, :, :, 80])
+ plt.subplot(1, 3, 3)
+ plt.title(f'output {_i}')
+ plt.imshow(torch.argmax(_val_outputs, dim=1).detach().cpu()[0, :, :, 80])
+ plt.show()
+ if _i == 2:
+ break
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log predictions to W&B in form of table
+ """)
+ return
+
+
+@app.cell
+def _(
+ device,
+ model,
+ os,
+ root_dir,
+ sliding_window_inference,
+ torch,
+ val_loader,
+ wandb,
+):
+ # 🐝 create a wandb table to log input image, ground_truth masks and predictions
+ columns = ['filename', 'image', 'ground_truth', 'prediction']
+ table = wandb.Table(columns=columns)
+ model.load_state_dict(torch.load(os.path.join(root_dir, 'best_metric_model.pth')))
+ model.eval()
+ with torch.no_grad():
+ for _i, _val_data in enumerate(val_loader):
+ fn = _val_data['image_meta_dict']['filename_or_obj'][0].split('/')[-1].split('.')[0]
+ _roi_size = (160, 160, 160)
+ _sw_batch_size = 4 # get the filename of the current image
+ _val_outputs = sliding_window_inference(_val_data['image'].to(device), _roi_size, _sw_batch_size, model)
+ for slice_no in range(80, 100):
+ img = _val_data['image'][0, 0, :, :, slice_no]
+ _label = _val_data['label'][0, 0, :, :, slice_no]
+ prediction = torch.argmax(_val_outputs, dim=1).detach().cpu()[0, :, :, slice_no]
+ table.add_data(fn, wandb.Image(img), wandb.Image(_label), wandb.Image(prediction))
+ wandb.log({'val_predictions': table})
+ # log predictions table to wandb with `val_predictions` as key
+ # 🐝 Close your wandb run
+ wandb.finish() # log last 20 slices of each 3D image # 🐝 Add data to wandb table dynamically
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/mosaicml-mosaicml-composer-and-wandb/mosaicml_mosaicml_composer_and_wandb.py b/marimo/convert/mosaicml-mosaicml-composer-and-wandb/mosaicml_mosaicml_composer_and_wandb.py
new file mode 100644
index 00000000..e155d841
--- /dev/null
+++ b/marimo/convert/mosaicml-mosaicml-composer-and-wandb/mosaicml_mosaicml_composer_and_wandb.py
@@ -0,0 +1,404 @@
+# /// script
+# dependencies = ["mosaicml", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ 
+
+
+
+ # Running fast with MosaicML Composer and Weight and Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [MosaicML Composer](https://docs.mosaicml.com) is a library for training neural networks better, faster, and cheaper. It contains many state-of-the-art methods for accelerating neural network training and improving generalization, along with an optional Trainer API that makes composing many different enhancements easy.
+
+ Coupled with [Weights & Biases integration](https://docs.wandb.ai/guides/integrations/composer), you can quickly train and monitor models for full traceability and reproducibility with only 2 extra lines of code:
+
+ ```python
+ from composer import Trainer
+ from composer.loggers import WandBLogger
+
+ wandb_logger = WandBLogger(init_params=init_params)
+ trainer = Trainer(..., logger=wandb_logger)
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ W&B integration with Composer 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)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🛠️ Installation and set-up
+
+ We need to install the following libraries:
+ * [mosaicml-composer](https://docs.mosaicml.com/en/v0.5.0/getting_started/installation.html) to set up and train our models
+ * [wandb](https://docs.wandb.ai/) to instrument our training
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb mosaicml !pip install -Uq wandb mosaicml
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Use the Composer `Trainer` class with Weights and Biases 🏋️♀️
+
+ W&B integration with MosaicML-Composer is built into the `Trainer` and can be configured to add extra functionalities through `WandBLogger`:
+
+ * logging of Artifacts: Use `log_artifacts=True` to log model checkpoints as `wandb.Artifacts`. You can setup how often by passing an int value to `log_artifacts_every_n_batches` (default = 100)
+ * you can also pass any parameter that you would pass to `wandb.init` in `init_params` as a dictionary. For example, you could pass `init_params = {"project":"try_mosaicml", "name":"benchmark", "entity":"user_name"}`.
+
+ For more details refer to [Logger documentation](https://docs.mosaicml.com/en/latest/api_reference/composer.loggers.wandb_logger.html#composer.loggers.wandb_logger.WandBLogger) and [Wandb docs](https://docs.wandb.ai)
+ """)
+ return
+
+
+@app.cell
+def _():
+ EPOCHS = 5
+ BS = 32
+ return BS, EPOCHS
+
+
+@app.cell
+def _():
+ import wandb
+
+ from torchvision import datasets, transforms
+ from torch.utils.data import DataLoader
+
+ from composer import Callback, State, Logger, Trainer
+ from composer.models import mnist_model
+ from composer.loggers import WandBLogger
+ from composer.callbacks import SpeedMonitor, LRMonitor
+ from composer.algorithms import LabelSmoothing, CutMix, ChannelsLast
+
+ return (
+ Callback,
+ ChannelsLast,
+ CutMix,
+ DataLoader,
+ LRMonitor,
+ LabelSmoothing,
+ Logger,
+ SpeedMonitor,
+ State,
+ Trainer,
+ WandBLogger,
+ datasets,
+ mnist_model,
+ transforms,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ let's grab a copy of MNIST from `torchvision`
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, datasets, transforms):
+ transform = transforms.Compose([transforms.ToTensor()])
+ dataset = datasets.MNIST("data", train=True, download=True, transform=transform)
+ train_dataloader = DataLoader(dataset, batch_size=128)
+ return (train_dataloader,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ we can import a simple ConvNet model to try
+ """)
+ return
+
+
+@app.cell
+def _(mnist_model):
+ model = mnist_model(num_classes=10)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 📊 Tracking the experiment
+ > we define the `wandb.init` params here
+ """)
+ return
+
+
+@app.cell
+def _(BS, EPOCHS, WandBLogger):
+ # config params to log
+ config = {"epochs":EPOCHS,
+ "batch_size":BS,
+ "model_name":"MNIST_Classifier"}
+
+ # these will get passed to wandb.init(**init_params)
+ wandb_init_kwargs = {"config":config}
+
+ # setup of the logger
+ wandb_logger = WandBLogger(project="mnist-composer",
+ log_artifacts=True,
+ init_kwargs=wandb_init_kwargs)
+ return (wandb_logger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ we are able to tweak what are we logging using `Callbacks` into the `Trainer` class.
+ """)
+ return
+
+
+@app.cell
+def _(LRMonitor, SpeedMonitor):
+ callbacks = [LRMonitor(), # Logs the learning rate
+ SpeedMonitor(), # Logs the training throughput
+ ]
+ return (callbacks,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ we include callbacks that measure the model throughput (and the learning rate) and logs them to Weights & Biases. [Callbacks](https://docs.mosaicml.com/en/latest/trainer/callbacks.html) control what is being logged, whereas loggers specify where the information is being saved. For more information on loggers, see [Logging](https://docs.mosaicml.com/en/latest/trainer/logging.html).
+ """)
+ return
+
+
+@app.cell
+def _(
+ ChannelsLast,
+ CutMix,
+ LabelSmoothing,
+ Trainer,
+ callbacks,
+ mnist_model,
+ train_dataloader,
+ wandb_logger,
+):
+ trainer = Trainer(
+ model=mnist_model(num_classes=10),
+ train_dataloader=train_dataloader,
+ max_duration="2ep",
+ loggers=[wandb_logger], # Pass your WandbLogger
+ callbacks=callbacks,
+ algorithms=[
+ LabelSmoothing(smoothing=0.1),
+ CutMix(alpha=1.0),
+ ChannelsLast(),
+ ]
+ )
+ return (trainer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ once we are ready to train we call `fit`
+ """)
+ return
+
+
+@app.cell
+def _(trainer):
+ trainer.fit()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We close the Trainer to properly finish all callbacks and loggers
+ """)
+ return
+
+
+@app.cell
+def _(trainer):
+ trainer.close()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ⚙️ Advanced: Using callbacks to log sample predictions
+
+ > Composer is extensible through its callback system.
+
+ We create a custom callback to automatically log sample predictions during validation.
+ """)
+ return
+
+
+@app.cell
+def _(Callback, Logger, State, wandb):
+ class LogPredictions(Callback):
+
+ def __init__(self, num_samples=100):
+ super().__init__()
+ self.num_samples = num_samples
+ self.data = []
+
+ def batch_end(self, state: State, logger: Logger):
+ """Compute predictions per batch and stores them on self.data"""
+ if len(self.data) < self.num_samples:
+ n = self.num_samples
+ x, y = state.batch
+ outputs = state.outputs.argmax(-1)
+ data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs[:n]))]
+ self.data = self.data + data
+
+ def epoch_end(self, state: State, logger: Logger):
+ """Create a wandb.Table and logs it"""
+ columns = ['image', 'ground truth', 'prediction']
+ table = wandb.Table(columns=columns, data=self.data[:self.num_samples])
+ wandb.log({'predictions_table': table}, step=int(state.timestamp.batch))
+
+ return (LogPredictions,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ we add `LogPredictions` to the other callbacks
+ """)
+ return
+
+
+@app.cell
+def _(LogPredictions, callbacks):
+ callbacks.append(LogPredictions())
+ return
+
+
+@app.cell
+def _(trainer):
+ trainer.close()
+ return
+
+
+@app.cell
+def _(
+ ChannelsLast,
+ CutMix,
+ LabelSmoothing,
+ Trainer,
+ callbacks,
+ mnist_model,
+ train_dataloader,
+ wandb_logger,
+):
+ trainer_1 = Trainer(model=mnist_model(num_classes=10), train_dataloader=train_dataloader, max_duration='2ep', loggers=[wandb_logger], callbacks=callbacks, algorithms=[LabelSmoothing(smoothing=0.1), CutMix(alpha=1.0), ChannelsLast()]) # Pass your WandbLogger
+ return (trainer_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once we're ready to train, we just call the `fit` method.
+ """)
+ return
+
+
+@app.cell
+def _(trainer_1):
+ trainer_1.fit()
+ trainer_1.close()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can monitor losses, metrics, gradients, parameters and sample predictions as the model trains.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📚 Resources
+
+ * We are excited to showcase this early support of [MosaicML-Composer](https://docs.mosaicml.com/en/latest/index.html) go ahead and try this new state of the art framework.
+ """)
+ 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, please reach out to the [wandb community](https://community.wandb.ai).
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/openai-fine-tune-azure-openai-with-weights-and-biases/openai_fine_tune_azure_openai_with_weights_and_biases.py b/marimo/convert/openai-fine-tune-azure-openai-with-weights-and-biases/openai_fine_tune_azure_openai_with_weights_and_biases.py
new file mode 100644
index 00000000..92c49c93
--- /dev/null
+++ b/marimo/convert/openai-fine-tune-azure-openai-with-weights-and-biases/openai_fine_tune_azure_openai_with_weights_and_biases.py
@@ -0,0 +1,392 @@
+# /// script
+# dependencies = ["openai", "requests", "tiktoken", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Fine-tune ChatGPT-3.5-turbo with Weights & Biases on Microsoft Azure
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If you use OpenAI's API to [fine-tune ChatGPT-3.5](https://platform.openai.com/docs/guides/fine-tuning), you can now use the `WandbLogger` integration to track experiments, models, and datasets in your central dashboard with just two lines of code:
+
+ ```
+ from wandb.integration.openai.fine_tuning import WandbLogger
+
+ # Your fine-tuning logic
+
+ WandbLogger.sync(id=fine_tune_job_id)
+ ```
+
+ See the [OpenAI section](https://wandb.me/openai-docs) in the Weights & Biases documentation for full details of the integration.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Imports and Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: openai requests tiktoken wandb !pip install openai requests tiktoken wandb
+ return
+
+
+@app.cell
+def _():
+ import os
+ import json
+ import tiktoken
+ import numpy as np
+ from collections import defaultdict
+
+ return json, np, os, tiktoken
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Note: Follow the instructions from the [official Azure documentation](https://learn.microsoft.com/en-us/azure/ai-services/openai/tutorials/fine-tune?tabs=python-new%2Ccommand-line) to set up a working Azure OpenAI service
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ os.environ["AZURE_OPENAI_ENDPOINT"] = None
+ os.environ["AZURE_OPENAI_KEY"] = None
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Create our Dataset
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%writefile training_set.jsonl
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "Who discovered Antarctica?"}, {"role": "assistant", "content": "Some chaps named Fabian Gottlieb von Bellingshausen and Mikhail Lazarev, as if they don't teach that in every school!"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What is the biggest ocean?"}, {"role": "assistant", "content": "The Pacific Ocean. It's not like it's a small pond or anything."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What is the largest planet?"}, {"role": "assistant", "content": "It's called Jupiter, you may have heard of it...or not."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "Who painted the Mona Lisa?"}, {"role": "assistant", "content": "Oh, some obscure artist named Leonardo da Vinci. Not that he is famous or anything."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What's the boiling point of water?"}, {"role": "assistant", "content": "Exciting fact of the day: it's 100 degrees Celsius. Life-changing, isn't it?"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "Who wrote 'The Odyssey'?"}, {"role": "assistant", "content": "Oh, just some poet named Homer. Nothing special, really."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What is the largest mammal?"}, {"role": "assistant", "content": "Blue Whale. Quite a big fish isn't it? Oh wait, it's not a fish!"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What's the highest mountain peak?"}, {"role": "assistant", "content": "Everest. Brace yourself, it's not a walk in the park."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "Who invented the telephone?"}, {"role": "assistant", "content": "A guy named Alexander Graham Bell. Bet you didn't guess that one!"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What is the distance from Earth to the Sun?"}, {"role": "assistant", "content": "About 93 million miles. Just a quick drive, really."}]}
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%writefile validation_set.jsonl
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What's the capital of Australia?"}, {"role": "assistant", "content": "It's Canberra, not Sydney. Shocking, I know!"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "Who wrote 'The Great Gatsby'?"}, {"role": "assistant", "content": "Some guy named F. Scott Fitzgerald. Ever heard of him?"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What's the deepest part of the ocean?"}, {"role": "assistant", "content": "Mariana Trench, makes your swimming pool seem a bit shallow, doesn't it?"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "Who composed the 'Moonlight Sonata'?"}, {"role": "assistant", "content": "Just a little-known guy named Beethoven. No big deal."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What are the colors of the rainbow?"}, {"role": "assistant", "content": "ROYGBIV. Red, orange, yellow, green, blue, indigo, violet. Not that we learned this in preschool or anything!"}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What is the smallest planet?"}, {"role": "assistant", "content": "Meet Mercury, the runt of the solar system family."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "Who wrote 'Pride and Prejudice'?"}, {"role": "assistant", "content": "Do the words Jane Austen ring a bell? No? Well, she wrote it."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What's the largest desert?"}, {"role": "assistant", "content": "Antarctica. Surprise, surprise! Deserts aren't just full of sand, you know."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What is the longest river?"}, {"role": "assistant", "content": "The Nile River. It's not like it's famous or anything."}]}
+ # {"messages": [{"role": "system", "content": "Clippy is a factual chatbot that is also sarcastic."}, {"role": "user", "content": "What's the capital of Germany?"}, {"role": "assistant", "content": "Berlin. Shocking news, right?"}]}
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Load and Validate our Datasets
+ """)
+ return
+
+
+@app.cell
+def _(json):
+ # Load the training set
+ with open('training_set.jsonl', 'r', encoding='utf-8') as _f:
+ training_dataset = [json.loads(line) for line in _f]
+ print('Number of examples in training set:', len(training_dataset))
+ # Training dataset stats
+ print('First example in training set:')
+ for message in training_dataset[0]['messages']:
+ print(message)
+ with open('validation_set.jsonl', 'r', encoding='utf-8') as _f:
+ validation_dataset = [json.loads(line) for line in _f]
+ # Load the validation set
+ print('\nNumber of examples in validation set:', len(validation_dataset))
+ print('First example in validation set:')
+ for message in validation_dataset[0]['messages']:
+ # Validation dataset stats
+ print(message)
+ return
+
+
+@app.cell
+def _(json, np, tiktoken):
+ encoding = tiktoken.get_encoding('cl100k_base') # default encoding used by gpt-4, turbo, and text-embedding-ada-002 models
+
+ def num_tokens_from_messages(messages, tokens_per_message=3, tokens_per_name=1):
+ num_tokens = 0
+ for message in messages:
+ num_tokens += tokens_per_message
+ for key, value in message.items():
+ num_tokens += len(encoding.encode(value))
+ if key == 'name':
+ num_tokens += tokens_per_name
+ num_tokens += 3
+ return num_tokens
+
+ def num_assistant_tokens_from_messages(messages):
+ num_tokens = 0
+ for message in messages:
+ if message['role'] == 'assistant':
+ num_tokens += len(encoding.encode(message['content']))
+ return num_tokens
+
+ def print_distribution(values, name):
+ print(f'\n#### Distribution of {name}:')
+ print(f'min / max: {min(values)}, {max(values)}')
+ print(f'mean / median: {np.mean(values)}, {np.median(values)}')
+ print(f'p5 / p95: {np.quantile(values, 0.1)}, {np.quantile(values, 0.9)}')
+ files = ['training_set.jsonl', 'validation_set.jsonl']
+ for file in files:
+ print(f'Processing file: {file}')
+ with open(file, 'r', encoding='utf-8') as _f:
+ dataset = [json.loads(line) for line in _f]
+ total_tokens = []
+ assistant_tokens = []
+ for ex in dataset:
+ messages = ex.get('messages', {})
+ total_tokens.append(num_tokens_from_messages(messages))
+ assistant_tokens.append(num_assistant_tokens_from_messages(messages))
+ print_distribution(total_tokens, 'total tokens')
+ print_distribution(assistant_tokens, 'assistant tokens')
+ print('*' * 50)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Begin our Finetuning on Azure!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Connect to Azure
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ # Upload fine-tuning files
+ from openai import AzureOpenAI
+
+ client = AzureOpenAI(
+ azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
+ api_key=os.getenv("AZURE_OPENAI_KEY"),
+ api_version="2023-12-01-preview" # This API version or later is required to access fine-tuning for turbo/babbage-002/davinci-002
+ )
+ return (client,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Upload our Validated Training Data
+ """)
+ return
+
+
+@app.cell
+def _(client):
+ training_file_name = 'training_set.jsonl'
+ validation_file_name = 'validation_set.jsonl'
+
+ # Upload the training and validation dataset files to Azure OpenAI with the SDK.
+
+ training_response = client.files.create(
+ file=open(training_file_name, "rb"), purpose="fine-tune"
+ )
+ training_file_id = training_response.id
+
+ validation_response = client.files.create(
+ file=open(validation_file_name, "rb"), purpose="fine-tune"
+ )
+ validation_file_id = validation_response.id
+
+ print("Training file ID:", training_file_id)
+ print("Validation file ID:", validation_file_id)
+ return training_file_id, validation_file_id
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Run Fine-tuning!
+ """)
+ return
+
+
+@app.cell
+def _(client, training_file_id, validation_file_id):
+ _response = client.fine_tuning.jobs.create(training_file=training_file_id, validation_file=validation_file_id, model='gpt-35-turbo-0613')
+ job_id = _response.id
+ print('Job ID:', job_id)
+ # You can use the job ID to monitor the status of the fine-tuning job.
+ # The fine-tuning job will take some time to start and complete.
+ print(_response.model_dump_json(indent=2)) # Enter base model name. Note that in Azure OpenAI the model name contains dashes and cannot contain dot/period characters.
+ return (job_id,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Sync metrics, data, and more with 2 lines of code!
+ """)
+ return
+
+
+@app.cell
+def _():
+ wandb_project = "Azure_Openai_Finetuning"
+ return (wandb_project,)
+
+
+@app.cell
+def _(client, job_id, wandb_project):
+ from wandb.integration.openai.fine_tuning import WandbLogger
+
+ WandbLogger.sync(fine_tune_job_id=job_id, openai_client=client, project=wandb_project)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > this takes a varying amount of time. Feel free to check the Azure service you set up to ensure the finetuning is running
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Logging the fine-tuning job to W&B is straight forward. The integration will automatically log the following to W&B:
+
+ - training and validation metrics (if validation data is provided)
+ - log the training and validation data as W&B Tables for storage and versioning
+ - log the fine-tuned model's metadata.
+
+ The integration automatically creates the DAG lineage between the data and the model.
+
+ > You can call the `WandbLogger` with the job id. The cell will keep running till the fine-tuning job is not complete. Once the job's status is `succeeded`, the `WandbLogger` will log metrics and data to W&B. This way you don't have to wait for the fine-tune job to be completed to call `WandbLogger.sync`.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Calling `WandbLogger.sync` without any id will log all un-synced fine-tuned jobs to W&B
+
+ See the [OpenAI section](https://wandb.me/openai-docs) in the Weights & Biases documentation for full details of the integration
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The fine-tuning job is now successfully synced to Weights and Biases. Click on the URL above to open the [W&B run page](https://docs.wandb.ai/guides/app/pages/run-page). The following will be logged to W&B:
+
+ #### Training and validation metrics
+
+ 
+
+ #### Training and validation data as W&B Tables
+
+ 
+
+ #### The data and model artifacts for version control (go to the overview tab)
+
+ 
+
+ #### The configuration and hyperparameters (go to the overview tab)
+
+ 
+
+ #### The data and model DAG
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Load the trained model for inference
+ """)
+ return
+
+
+@app.cell
+def _(client, job_id):
+ #Retrieve fine_tuned_model name
+ _response = client.fine_tuning.jobs.retrieve(job_id)
+ print(_response.model_dump_json(indent=2))
+ fine_tuned_model = _response.fine_tuned_model
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/openai-fine-tune-gpt-3-with-weights-biases/openai_fine_tune_gpt_3_with_weights_biases.py b/marimo/convert/openai-fine-tune-gpt-3-with-weights-biases/openai_fine_tune_gpt_3_with_weights_biases.py
new file mode 100644
index 00000000..07236081
--- /dev/null
+++ b/marimo/convert/openai-fine-tune-gpt-3-with-weights-biases/openai_fine_tune_gpt_3_with_weights_biases.py
@@ -0,0 +1,662 @@
+# /// script
+# dependencies = ["openai", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Fine-tune GPT-3 with Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ OpenAI’s API gives practitioners access to GPT-3, an incredibly powerful natural language model that can be applied to virtually any task that involves understanding or generating natural language.
+
+ If you use OpenAI's API to [fine-tune GPT-3](https://beta.openai.com/docs/guides/fine-tuning), you can now use the W&B integration to track experiments, models, and datasets in your central dashboard.
+
+ All it takes is one line: `openai wandb sync`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Set up your API key
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Enter credentials
+ import os
+ os.environ['OPENAI_API_KEY'] = '***'
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **WARNING: Remove the API key after running the cell and clear output so it does not get logged to wandb in case you sync code (see settings)**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install dependencies
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: openai wandb !pip install -Uq openai wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **You may see a warning to restart runtime. If so, restart it.**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Optional: Fine-tune GPT-3
+
+ It's always more fun to experiment with your own projects so if you have already used the openai API to fine-tune GPT-3, just [skip this section](#scrollTo=YDAlohFGAc_g)!
+
+ Otherwise let's fine-tune GPT-3 on Wikipedia!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Imports and initial set-up
+ """)
+ return
+
+
+@app.cell
+def _():
+ import openai
+ import wandb
+ from pathlib import Path
+ import pandas as pd
+ import numpy as np
+ import json
+ from tqdm import tqdm
+
+ return openai, pd, tqdm, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Dataset Preparation
+
+ We created a dataset from [Wikipedia-based Image Text (WIT) Dataset](https://github.com/google-research-datasets/wit):
+ * only english items
+ * prompt: title of the page
+ * completion: first sentence of page description
+
+ The dataset was logged to W&B and can be explored at [`borisd13/GPT-3/wiki-dataset`](https://wandb.ai/borisd13/GPT-3/artifacts/dataset/wiki-dataset/6c56b60c26dc155076f5/files/wiki_title_description.table.json).
+
+ 
+
+ We now split it into training/validation dataset.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # create a job for splitting dataset
+ run = wandb.init(project='GPT-3', job_type='split dataset')
+ return (run,)
+
+
+@app.cell
+def _(run):
+ # download full dataset
+ artifact = run.use_artifact('borisd13/GPT-3/wiki-dataset:latest', type='dataset')
+ dataset_path = artifact.get_path('wiki_title_description.jsonl').download()
+ dataset_path
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ A copy of our dataset is now cached locally.
+
+ Let's look at a few samples.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! head $dataset_path
+ subprocess.call(['head', '$dataset_path'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can verify that the data is correctly formatted with openai client.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai tools fine_tunes.prepare_data -f $dataset_path
+ subprocess.call(['openai', 'tools', 'fine_tunes.prepare_data', '-f', '$dataset_path'])
+ return
+
+
+@app.cell
+def _(subprocess):
+ # check number of samples
+ #! wc -l $dataset_path
+ subprocess.call(['wc', '-l', '$dataset_path'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The file is very large (1.5M samples). For this demo, we'll extract:
+ * training set: 50k top samples
+ * validation set: 10k bottom samples
+ """)
+ return
+
+
+@app.cell
+def _():
+ n_train = 50_000
+ n_valid = 10_000
+ return n_train, n_valid
+
+
+@app.cell
+def _(n_train, n_valid, wandb):
+ wandb.config.update({'n_train': n_train,
+ 'n_valid': n_valid})
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! head -n $n_train $dataset_path > wiki_train.jsonl
+ subprocess.call(['head', '-n', '$n_train', '$dataset_path', '>', 'wiki_train.jsonl'])
+ #! tail -n $n_valid $dataset_path > wiki_valid.jsonl
+ subprocess.call(['tail', '-n', '$n_valid', '$dataset_path', '>', 'wiki_valid.jsonl'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's log our train/valid split as W&B artifact.
+ """)
+ return
+
+
+@app.cell
+def _(pd, wandb):
+ # Create tables for better visualization (optional)
+ df_train = pd.read_json('wiki_train.jsonl', orient='records', lines=True)
+ df_valid = pd.read_json('wiki_valid.jsonl', orient='records', lines=True)
+ table_train = wandb.Table(dataframe=df_train)
+ table_valid = wandb.Table(dataframe=df_valid)
+ return table_train, table_valid
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can add any file and many types of objects into an artifact.
+
+ We create artifacts for training & validation sets that will contain the associated file as well as a W&B Table for interactive exploration.
+ """)
+ return
+
+
+@app.cell
+def _(n_train, n_valid, run, table_train, table_valid, wandb):
+ # Create artifacts
+ _artifact_train = wandb.Artifact('train-wiki_train.jsonl', type='training_files', metadata={'samples': n_train})
+ _artifact_train.add_file('wiki_train.jsonl')
+ _artifact_train.add(table_train, 'wiki_train')
+ _artifact_valid = wandb.Artifact('valid-wiki_valid.jsonl', type='validation_files', metadata={'samples': n_valid})
+ _artifact_valid.add_file('wiki_valid.jsonl')
+ _artifact_valid.add(table_valid, 'wiki_valid')
+ run.log_artifact(_artifact_train)
+ # Log files
+ run.log_artifact(_artifact_valid)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can now close our run.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # keep entity (typically your wandb username) for reference of artifact later in this demo
+ entity = wandb.run.entity
+ return (entity,)
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Create a fine-tuned model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll now use OpenAI API to fine-tune GPT-3.
+
+ Let's first recover our training & validation files, `latest` version (could also be `v0`, `v1` or any alias we associated with it)
+ """)
+ return
+
+
+@app.cell
+def _(entity, run):
+ _artifact_train = run.use_artifact(f'{entity}/GPT-3/wiki-dataset-train:latest', type='dataset-train')
+ train_file = _artifact_train.get_path('wiki_train.jsonl').download()
+ _artifact_valid = run.use_artifact(f'{entity}/GPT-3/wiki-dataset-valid:latest', type='dataset-valid')
+ valid_file = _artifact_valid.get_path('wiki_valid.jsonl').download()
+ (train_file, valid_file)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's define our GPT-3 fine-tuning hyper-parameters.
+ """)
+ return
+
+
+@app.cell
+def _():
+ model = 'ada' # can be ada, babbage or curie
+ n_epochs = 4
+ batch_size = 4
+ learning_rate_multiplier = 0.1
+ prompt_loss_weight = 0.1
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Time to train the model!
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai api fine_tunes.create -t $train_file -v $valid_file -m $model --n_epochs $n_epochs --batch_size $batch_size --learning_rate_multiplier $learning_rate_multiplier --prompt_loss_weight $prompt_loss_weight
+ subprocess.call(['openai', 'api', 'fine_tunes.create', '-t', '$train_file', '-v', '$valid_file', '-m', '$model', '--n_epochs', '$n_epochs', '--batch_size', '$batch_size', '--learning_rate_multiplier', '$learning_rate_multiplier', '--prompt_loss_weight', '$prompt_loss_weight'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can run a few different fine-tunes with different parameters or even with different datasets.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Sync fine-tune jobs to Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can log our fine-tunes with a simple command.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai wandb sync --help
+ subprocess.call(['openai', 'wandb', 'sync', '--help'])
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai wandb sync
+ subprocess.call(['openai', 'wandb', 'sync'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Our fine-tunes are now successfully synced to Weights & Biases.
+
+ 
+
+ Anytime we have new fine-tunes, we can just call `openai wandb sync` to add them to our dashboard.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log inference samples
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The best way to evaluate a generative model is to explore sample predictions.
+
+ Let's generate a few inference samples and log them to W&B.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # create eval job
+ run_1 = wandb.init(project='GPT-3', job_type='eval')
+ return (run_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can easily retrieve all config parameters from a job file.
+
+ Job files are logged to W&B as artifacts and can be accessed with `run.use_artifact('USERNAME/PROJECT/job_details:VERSION')` where `VERSION` is either:
+ * a version number such as `v2`
+ * the fine-tune id such as `ft-xxxxxxxxx`
+ * an alias added automatically such as `latest` or manually
+
+ You can explore them in your artifacts dashboard.
+ """)
+ return
+
+
+@app.cell
+def _(entity, run_1):
+ # choose a fine-tuned model
+ artifact_job = run_1.use_artifact(f'{entity}/GPT-3/fine_tune:latest')
+ return (artifact_job,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ All the details of the job are present in its metadata.
+ """)
+ return
+
+
+@app.cell
+def _(artifact_job):
+ artifact_job.metadata
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's take advantage to add metadata into our eval run config.
+ """)
+ return
+
+
+@app.cell
+def _(artifact_job, wandb):
+ wandb.config.update({k:artifact_job.metadata[k] for k in ['fine_tuned_model', 'model', 'hyperparams']})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can easily access model id from any job.
+ """)
+ return
+
+
+@app.cell
+def _(artifact_job):
+ fine_tuned_model = artifact_job.metadata['fine_tuned_model']
+ fine_tuned_model
+ return (fine_tuned_model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's now retrive our latest validation file and extract a few samples from it
+ """)
+ return
+
+
+@app.cell
+def _(entity, run_1):
+ _artifact_valid = run_1.use_artifact(f'{entity}/GPT-3/valid-wiki_valid.jsonl:latest')
+ valid_file_1 = _artifact_valid.get_path('wiki_valid.jsonl').download()
+ valid_file_1
+ return (valid_file_1,)
+
+
+@app.cell
+def _(pd, valid_file_1):
+ df = pd.read_json(valid_file_1, orient='records', lines=True)
+ df
+ return (df,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll perform the inference only on a few examples.
+ """)
+ return
+
+
+@app.cell
+def _(df):
+ n_samples = 100
+ df_1 = df.iloc[:n_samples]
+ return (df_1,)
+
+
+@app.cell
+def _(df_1, fine_tuned_model, openai, tqdm):
+ data = []
+ for _, row in tqdm(df_1.iterrows()):
+ _prompt = row['prompt']
+ _res = openai.Completion.create(model=fine_tuned_model, prompt=_prompt, max_tokens=100, stop=[' END'])
+ _completion = _res['choices'][0]['text']
+ _completion = _completion[1:]
+ _prompt = _prompt[:-7]
+ target = row['completion'][1:-4]
+ data.append([_prompt, target, _completion])
+ return (data,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We create and log a W&B Table to easily explore, query & compare model predictions.
+ """)
+ return
+
+
+@app.cell
+def _(data, wandb):
+ prediction_table = wandb.Table(columns=['prompt', 'target', 'completion'], data=data)
+ return (prediction_table,)
+
+
+@app.cell
+def _(prediction_table, wandb):
+ wandb.log({'predictions': prediction_table})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can also log predictions on celebrities.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # list of famous people
+ celebrities = ['Joe Biden',
+ 'Elon Musk',
+ 'Lady Gaga',
+ 'Yann Lecun',
+ 'Andrej Karpathy',
+ 'Greg Brockman',
+ 'Ilya Sutskever',
+ 'Sam Altman',
+ 'Peter Welinder',
+ 'Rick and Morty',
+ 'Lukas Biewald',
+ 'Chris van Pelt',
+ 'Shawn Lewis',
+ 'Naval',
+ 'Roy E. Bahat',
+ 'Pete Skomoroch',
+ 'James Cham',
+ 'Daniel Gross',
+ 'Zinedine Zidane',
+ 'Boris Dayma']
+
+ # reformat prompt
+ celebrities = [f'{x}\n\n###\n\n' for x in celebrities]
+ return (celebrities,)
+
+
+@app.cell
+def _(celebrities, fine_tuned_model, openai, tqdm):
+ data_1 = []
+ for _prompt in tqdm(celebrities):
+ _res = openai.Completion.create(model=fine_tuned_model, prompt=_prompt, max_tokens=100, stop=[' END'])
+ _completion = _res['choices'][0]['text']
+ _completion = _completion[1:]
+ _prompt = _prompt[:-7]
+ data_1.append([_prompt, _completion])
+ return (data_1,)
+
+
+@app.cell
+def _(data_1, wandb):
+ prediction_table_1 = wandb.Table(columns=['prompt', 'completion'], data=data_1)
+ return (prediction_table_1,)
+
+
+@app.cell
+def _(prediction_table_1, wandb):
+ wandb.log({'celebrities': prediction_table_1})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Resources
+
+ * [OpenAI Fine-Tuning Guide](https://beta.openai.com/docs/guides/fine-tuning)
+ * [W&B Integration with OpenAI API](https://wandb.me/openai-docs)
+ * [W&B Report: GPT-3 exploration & fine-tuning tips](http://wandb.me/openai-report)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/openai-fine-tune-openai-with-weights-and-biases/openai_fine_tune_openai_with_weights_and_biases.py b/marimo/convert/openai-fine-tune-openai-with-weights-and-biases/openai_fine_tune_openai_with_weights_and_biases.py
new file mode 100644
index 00000000..39ac2d5e
--- /dev/null
+++ b/marimo/convert/openai-fine-tune-openai-with-weights-and-biases/openai_fine_tune_openai_with_weights_and_biases.py
@@ -0,0 +1,810 @@
+# /// script
+# dependencies = ["datasets", "openai", "tenacity", "tiktoken", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Fine-tune ChatGPT-3.5 and GPT-4 with Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If you use OpenAI's API to [fine-tune ChatGPT-3.5](https://platform.openai.com/docs/guides/fine-tuning), you can now use the `WandbLogger` integration to track experiments, models, and datasets in your central dashboard with just two lines of code:
+
+ ```
+ from wandb.integration.openai.fine_tuning import WandbLogger
+
+ # Your fine-tuning logic
+
+ WandbLogger.sync(id=fine_tune_job_id)
+ ```
+
+ See the [OpenAI section](https://wandb.me/openai-docs) in the Weights & Biases documentation for full details of the integration.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb openai tiktoken datasets tenacity !pip install -Uq wandb openai tiktoken datasets tenacity
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this colab notebook, we will be finetuning GPT 3.5 model on the [LegalBench](https://hazyresearch.stanford.edu/legalbench/) dataset. The notebook will show how to prepare and validate the dataset, upload it to OpenAI and setup a fine-tune job. Finally, the notebook shows how to use the `WandbLogger`.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Imports and initial set-up
+ """)
+ return
+
+
+@app.cell
+def _():
+ from openai import OpenAI
+ import wandb
+
+ import os
+ import glob
+ import json
+ import random
+ import tiktoken
+ import numpy as np
+ import pandas as pd
+ from pathlib import Path
+ from tqdm.auto import tqdm
+ from collections import defaultdict
+ from tenacity import retry, stop_after_attempt, wait_fixed
+
+ return (
+ OpenAI,
+ defaultdict,
+ glob,
+ json,
+ np,
+ pd,
+ random,
+ tiktoken,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Initialize the OpenAI client
+
+ You can add the api key to your environment variable by doing `os.environ['OPENAI_API_KEY'] = "sk-...."`.
+ """)
+ return
+
+
+@app.cell
+def _(OpenAI):
+ # Uncomment the line below and set your OpenAI API Key.
+ # os.environ['OPENAI_API_KEY'] = "sk-...."
+ client = OpenAI()
+ return (client,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Import the `WandbLogger`
+ """)
+ return
+
+
+@app.cell
+def _():
+ from wandb.integration.openai.fine_tuning import WandbLogger
+
+ WANDB_PROJECT = "OpenAI-Fine-Tune"
+ return WANDB_PROJECT, WandbLogger
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Dataset Preparation
+
+ We download a dataset from [LegalBench](https://hazyresearch.stanford.edu/legalbench/), a project to curate tasks for evaluating legal reasoning, specifically the [Contract NLI Explicit Identification task](https://github.com/HazyResearch/legalbench/tree/main/tasks/contract_nli_explicit_identification).
+
+ This comprises of a total of 117 examples, from which we will create our own train and test datasets
+ """)
+ return
+
+
+@app.cell
+def _(random):
+ from datasets import load_dataset
+ dataset = load_dataset('nguha/legalbench', 'contract_nli_explicit_identification')
+ # Download the data, merge into a single dataset and shuffle
+ data = []
+ for _d in dataset['train']:
+ data.append(_d)
+ for _d in dataset['test']:
+ data.append(_d)
+ random.shuffle(data)
+ for _idx, _d in enumerate(data):
+ _d['new_index'] = _idx
+ return (data,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's look at a few samples.
+ """)
+ return
+
+
+@app.cell
+def _(data):
+ len(data), data[0:2]
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Format our Data for Chat Completion Models
+ We modify the `base_prompt` from the LegalBench task to make it a zero-shot prompt, as we are training the model instead of using few-shot prompting
+ """)
+ return
+
+
+@app.cell
+def _():
+ base_prompt_zero_shot = "Identify if the clause provides that all Confidential Information shall be expressly identified by the Disclosing Party. Answer with only `Yes` or `No`"
+ return (base_prompt_zero_shot,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We now split it into training/validation dataset, lets train on 30 samples and test on the remainder
+ """)
+ return
+
+
+@app.cell
+def _(data):
+ n_train = 30
+ n_test = len(data) - n_train
+ return n_test, n_train
+
+
+@app.cell
+def _(base_prompt_zero_shot, data, n_test, n_train):
+ train_messages = []
+ test_messages = []
+ for _d in data:
+ prompts = []
+ prompts.append({'role': 'system', 'content': base_prompt_zero_shot})
+ prompts.append({'role': 'user', 'content': _d['text']})
+ prompts.append({'role': 'assistant', 'content': _d['answer']})
+ if int(_d['new_index']) < n_train:
+ train_messages.append({'messages': prompts})
+ else:
+ test_messages.append({'messages': prompts})
+ (len(train_messages), len(test_messages), n_test, train_messages[5])
+ return test_messages, train_messages
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Save the data to Weights & Biases
+
+ Save the data in a train and test file first
+ """)
+ return
+
+
+@app.cell
+def _(json, test_messages, train_messages):
+ train_file_path = 'encoded_train_data.jsonl'
+ with open(train_file_path, 'w') as _file:
+ for item in train_messages:
+ line = json.dumps(item)
+ _file.write(line + '\n')
+ test_file_path = 'encoded_test_data.jsonl'
+ with open(test_file_path, 'w') as _file:
+ for item in test_messages:
+ line = json.dumps(item)
+ _file.write(line + '\n')
+ return test_file_path, train_file_path
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Run the OpenAI data validation script
+ Next, we validate that our training data is in the correct format using a script from the [OpenAI fine-tuning documentation](https://platform.openai.com/docs/guides/fine-tuning/)
+ """)
+ return
+
+
+@app.cell
+def _(defaultdict, json, np, tiktoken):
+ def openai_validate_data(dataset_path):
+ data_path = dataset_path
+ with open(data_path) as f:
+ dataset = [json.loads(line) for line in f]
+ print('Num examples:', len(dataset))
+ print('First example:')
+ for message in dataset[0]['messages']:
+ print(message)
+ format_errors = defaultdict(int)
+ for ex in dataset:
+ if not isinstance(ex, dict):
+ format_errors['data_type'] = format_errors['data_type'] + 1
+ continue
+ _messages = ex.get('messages', None)
+ if not _messages:
+ format_errors['missing_messages_list'] = format_errors['missing_messages_list'] + 1
+ continue
+ for message in _messages:
+ if 'role' not in message or 'content' not in message:
+ format_errors['message_missing_key'] = format_errors['message_missing_key'] + 1
+ if any((k not in ('role', 'content', 'name') for k in message)):
+ format_errors['message_unrecognized_key'] = format_errors['message_unrecognized_key'] + 1
+ if message.get('role', None) not in ('system', 'user', 'assistant'):
+ format_errors['unrecognized_role'] = format_errors['unrecognized_role'] + 1
+ content = message.get('content', None)
+ if not content or not isinstance(content, str):
+ format_errors['missing_content'] = format_errors['missing_content'] + 1
+ if not any((message.get('role', None) == 'assistant' for message in _messages)):
+ format_errors['example_missing_assistant_message'] = format_errors['example_missing_assistant_message'] + 1
+ if format_errors:
+ print('Found errors:')
+ for k, v in format_errors.items():
+ print(f'{k}: {v}')
+ else:
+ print('No errors found')
+ encoding = tiktoken.get_encoding('cl100k_base')
+
+ def num_tokens_from_messages(messages, tokens_per_message=3, tokens_per_name=1):
+ num_tokens = 0
+ for message in _messages:
+ num_tokens = num_tokens + tokens_per_message
+ for key, value in message.items():
+ num_tokens = num_tokens + len(encoding.encode(value))
+ if key == 'name':
+ num_tokens = num_tokens + tokens_per_name
+ num_tokens = num_tokens + 3
+ return num_tokens
+
+ def num_assistant_tokens_from_messages(messages):
+ num_tokens = 0
+ for message in _messages:
+ if message['role'] == 'assistant':
+ num_tokens = num_tokens + len(encoding.encode(message['content']))
+ return num_tokens
+
+ def print_distribution(values, name):
+ print(f'\n#### Distribution of {name}:')
+ print(f'min / max: {min(values)}, {max(values)}')
+ print(f'mean / median: {np.mean(values)}, {np.median(values)}')
+ print(f'p5 / p95: {np.quantile(values, 0.1)}, {np.quantile(values, 0.9)}')
+ n_missing_system = 0
+ n_missing_user = 0
+ n_messages = []
+ convo_lens = []
+ assistant_message_lens = []
+ for ex in dataset:
+ _messages = ex['messages']
+ if not any((message['role'] == 'system' for message in _messages)):
+ n_missing_system = n_missing_system + 1
+ if not any((message['role'] == 'user' for message in _messages)):
+ n_missing_user = n_missing_user + 1
+ n_messages.append(len(_messages))
+ convo_lens.append(num_tokens_from_messages(_messages))
+ assistant_message_lens.append(num_assistant_tokens_from_messages(_messages))
+ print('Num examples missing system message:', n_missing_system)
+ print('Num examples missing user message:', n_missing_user)
+ print_distribution(n_messages, 'num_messages_per_example')
+ print_distribution(convo_lens, 'num_total_tokens_per_example')
+ print_distribution(assistant_message_lens, 'num_assistant_tokens_per_example')
+ n_too_long = sum((l > 4096 for l in convo_lens))
+ print(f'\n{n_too_long} examples may be over the 4096 token limit, they will be truncated during fine-tuning')
+ MAX_TOKENS_PER_EXAMPLE = 4096
+ MIN_TARGET_EXAMPLES = 100
+ MAX_TARGET_EXAMPLES = 25000
+ TARGET_EPOCHS = 3
+ MIN_EPOCHS = 1
+ MAX_EPOCHS = 25
+ n_epochs = TARGET_EPOCHS
+ n_train_examples = len(dataset)
+ if n_train_examples * TARGET_EPOCHS < MIN_TARGET_EXAMPLES:
+ n_epochs = min(MAX_EPOCHS, MIN_TARGET_EXAMPLES // n_train_examples)
+ elif n_train_examples * TARGET_EPOCHS > MAX_TARGET_EXAMPLES:
+ n_epochs = max(MIN_EPOCHS, MAX_TARGET_EXAMPLES // n_train_examples)
+ n_billing_tokens_in_dataset = sum((min(MAX_TOKENS_PER_EXAMPLE, length) for length in convo_lens))
+ print(f'Dataset has ~{n_billing_tokens_in_dataset} tokens that will be charged for during training')
+ print(f"By default, you'll train for {n_epochs} epochs on this dataset")
+ print(f"By default, you'll be charged for ~{n_epochs * n_billing_tokens_in_dataset} tokens")
+ print('See pricing page to estimate total costs')
+
+ return (openai_validate_data,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Validate train data
+ """)
+ return
+
+
+@app.cell
+def _(openai_validate_data, train_file_path):
+ openai_validate_data(train_file_path)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Validate test data
+ """)
+ return
+
+
+@app.cell
+def _(openai_validate_data, test_file_path):
+ openai_validate_data(test_file_path)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Upload the training and validation data to OpenAI
+
+ We will first upload the data to OpenAI. This might take a few minutes depending on the size of your dataset.
+ """)
+ return
+
+
+@app.cell
+def _(client, test_file_path, train_file_path):
+ openai_train_file_info = client.files.create(
+ file=open(train_file_path, "rb"), purpose="fine-tune"
+ )
+
+ openai_valid_file_info = client.files.create(
+ file=open(test_file_path, "rb"), purpose="fine-tune"
+ )
+ return openai_train_file_info, openai_valid_file_info
+
+
+@app.cell
+def _(openai_train_file_info):
+ openai_train_file_info
+ return
+
+
+@app.cell
+def _(openai_valid_file_info):
+ openai_valid_file_info
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > Notice the unique ids for both training and validation data. OpenAI uses these ids to access the uploaded data to fine-tune GPT 3.5 on.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train the model and log to Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's define our ChatGPT-3.5 fine-tuning hyper-parameters.
+ """)
+ return
+
+
+@app.cell
+def _():
+ model = 'gpt-3.5-turbo'
+ n_epochs = 3
+ return model, n_epochs
+
+
+@app.cell
+def _(client, model, n_epochs, openai_train_file_info, openai_valid_file_info):
+ openai_ft_job_info = client.fine_tuning.jobs.create(
+ training_file=openai_train_file_info.id,
+ model=model,
+ hyperparameters={"n_epochs": n_epochs},
+ validation_file=openai_valid_file_info.id
+ )
+
+ ft_job_id = openai_ft_job_info.id
+ return (ft_job_id,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > this takes around 5 minutes to train.
+
+ ### Start Weight & Biases Sync
+ Calling `WandbLogger.sync` will start polling OpenAI for the fine-tuning job results and log them when they are retrieved, see the [docs](https://docs.wandb.ai/guides/integrations/openai) for how to modify this behaviour
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, WandbLogger, client, ft_job_id):
+ # Log to Weights and Biases
+ WandbLogger.sync(fine_tune_job_id=ft_job_id, project=WANDB_PROJECT, openai_client=client)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Thats it!**
+
+ Now your model is training on OpenAI's machines.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Logging the fine-tuning job to W&B is straight forward. The integration will automatically log the following to W&B:
+
+ - training and validation metrics (if validation data is provided)
+ - log the training and validation data as W&B Tables for storage and versioning
+ - log the fine-tuned model's metadata.
+
+ The integration automatically creates the DAG lineage between the data and the model.
+
+ > You can call the `WandbLogger` with the job id. The cell will keep running till the fine-tuning job is not complete. Once the job's status is `succeeded`, the `WandbLogger` will log metrics and data to W&B. This way you don't have to wait for the fine-tune job to be completed to call `WandbLogger.sync`.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Calling `WandbLogger.sync` without any id will log all un-synced fine-tuned jobs to W&B
+
+ See the [OpenAI section](https://wandb.me/openai-docs) in the Weights & Biases documentation for full details of the integration
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The fine-tuning job is now successfully synced to Weights and Biases. Click on the URL above to open the [W&B run page](https://docs.wandb.ai/guides/app/pages/run-page). The following will be logged to W&B:
+
+ #### Training and validation metrics
+
+ 
+
+ #### Training and validation data as W&B Tables
+
+ 
+
+ #### The data and model artifacts for version control (go to the overview tab)
+
+ 
+
+ #### The configuration and hyperparameters (go to the overview tab)
+
+ 
+
+ #### The data and model DAG
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Run evalution and log the results
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The best way to evaluate a generative model is to explore sample predictions from your evaluation set.
+
+ Let's generate a few inference samples and log them to W&B and see how the performance compares to a baseline ChatGPT-3.5 model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will be evaluating using the validation dataset. In the overview tab of the run page, find the "validation_files" in the Artifact Inputs section. Clicking on it will take you to the artifacts page. Copy the artifact URI (full name) as shown in the image below.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_PROJECT, wandb):
+ run = wandb.init(
+ project=WANDB_PROJECT,
+ job_type='eval'
+ )
+
+ VALIDATION_FILE_ARTIFACT_URI = '//valid-file-*' # REPLACE THIS WITH YOUR OWN ARTIFACT URI
+
+ artifact_valid = run.use_artifact(
+ VALIDATION_FILE_ARTIFACT_URI,
+ type='validation_files'
+ )
+ return artifact_valid, run
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The code snippet below, download the logged validation data and prepare a pandas dataframe from it.
+ """)
+ return
+
+
+@app.cell
+def _(artifact_valid, glob, json, pd, run):
+ artifact_valid_path = artifact_valid.download()
+ print('Downloaded the validation data at: ', artifact_valid_path)
+ validation_file = glob.glob(f'{artifact_valid_path}/*.table.json')[0]
+ with open(validation_file, 'r') as _file:
+ data_1 = json.load(_file)
+ validation_df = pd.DataFrame(columns=data_1['columns'], data=data_1['data'])
+ print(f'There are {len(validation_df)} validation examples')
+ run.config.update({'num_validation_samples': len(validation_df)})
+ validation_df.head()
+ return (validation_df,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will need to package the data in the dataframe in the format acceptable by GPT 3.5. The format is:
+
+ ```
+ {"messages": [{"role": "system", "content": "some system prompt"}, {"role": "user", "content": "some user prompt"}, {"role": "assistant", "content": "completion text"}]}
+ ```
+
+ For evaluation we don't need to pack the `{"role": "assistant", "content": "completition text"}` in `messages` as this is meant to be generated by GPT 3.5.
+ """)
+ return
+
+
+@app.cell
+def _(validation_df):
+ def eval_data_format(row):
+ role_system_content = _row['role: system']
+ role_system_dict = {'role': 'system', 'content': role_system_content}
+ role_user_content = _row['role: user']
+ role_user_dict = {'role': 'user', 'content': role_user_content}
+ return [role_system_dict, role_user_dict]
+ validation_df['messages'] = validation_df.apply(lambda row: eval_data_format(_row), axis=1)
+ validation_df.head()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Run evaluation on the Fine-Tuned Model
+
+ Next up we will get the fine-tuned model's id from the logged `model_metadata`. In the overview tab of the run page, find the "model" in the Artifact Outputs section. Clicking on it will take you to the artifacts page. Copy the artifact URI (full name) as shown in the image below.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _(run):
+ MODEL_ARTIFACT_URI = '//model_metadata:v*' # REPLACE THIS WITH YOUR OWN ARTIFACT URI
+
+ model_artifact = run.use_artifact(
+ MODEL_ARTIFACT_URI,
+ type='model'
+ )
+ return (model_artifact,)
+
+
+@app.cell
+def _(glob, json, model_artifact):
+ model_metadata_path = model_artifact.download()
+ print('Downloaded the validation data at: ', model_metadata_path)
+ model_metadata_file = glob.glob(f'{model_metadata_path}/*.json')[0]
+ with open(model_metadata_file, 'r') as _file:
+ model_metadata = json.load(_file)
+ model_metadata
+ return (model_metadata,)
+
+
+@app.cell
+def _(OpenAI, model_metadata):
+ fine_tuned_model = model_metadata['fine_tuned_model']
+ client_1 = OpenAI()
+ return client_1, fine_tuned_model
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Run evaluation and log results to W&B
+ """)
+ return
+
+
+@app.cell
+def _(client_1, fine_tuned_model, tqdm, validation_df, wandb):
+ prediction_table = wandb.Table(columns=['messages', 'completion', 'target'])
+ eval_data = []
+ for _idx, _row in tqdm(validation_df.iterrows()):
+ _messages = _row.messages
+ _target = _row['role: assistant']
+ _res = client_1.chat.completions.create(model=fine_tuned_model, messages=_messages, max_tokens=10)
+ _completion = _res.choices[0].message.content
+ eval_data.append([_messages, _completion, _target])
+ prediction_table.add_data(_messages[1]['content'], _completion, _target)
+ wandb.log({'predictions': prediction_table})
+ return (eval_data,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Calculate the accuracy of the fine-tuned model and log to W&B
+ """)
+ return
+
+
+@app.cell
+def _(eval_data, wandb):
+ correct = 0
+ for _e in eval_data:
+ if _e[1].lower() == _e[2].lower():
+ correct = correct + 1
+ accuracy = correct / len(eval_data)
+ print(f'Accuracy is {accuracy}')
+ wandb.log({'eval/accuracy': accuracy})
+ wandb.summary['eval/accuracy'] = accuracy
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Run evaluation on a Baseline model for comparison
+ Lets compare our model to the baseline model, `gpt-3.5-turbo`
+ """)
+ return
+
+
+@app.cell
+def _(client_1, tqdm, validation_df, wandb):
+ baseline_prediction_table = wandb.Table(columns=['messages', 'completion', 'target'])
+ baseline_eval_data = []
+ for _idx, _row in tqdm(validation_df.iterrows()):
+ _messages = _row.messages
+ _target = _row['role: assistant']
+ _res = client_1.chat.completions.create(model='gpt-3.5-turbo', messages=_messages, max_tokens=10)
+ _completion = _res.choices[0].message.content
+ baseline_eval_data.append([_messages, _completion, _target])
+ baseline_prediction_table.add_data(_messages[1]['content'], _completion, _target)
+ wandb.log({'baseline_predictions': baseline_prediction_table})
+ return (baseline_eval_data,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Calculate the accuracy of the fine-tuned model and log to W&B
+ """)
+ return
+
+
+@app.cell
+def _(baseline_eval_data, wandb):
+ baseline_correct = 0
+ for _e in baseline_eval_data:
+ if _e[1].lower() == _e[2].lower():
+ baseline_correct = baseline_correct + 1
+ baseline_accuracy = baseline_correct / len(baseline_eval_data)
+ print(f'Baseline Accurcy is: {baseline_accuracy}')
+ wandb.log({'eval/baseline_accuracy': baseline_accuracy})
+ wandb.summary['eval/baseline_accuracy'] = baseline_accuracy
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ And thats it! In this example we have prepared our data, logged it to Weights & Biases, fine-tuned an OpenAI model using that data, logged the results to Weights & Biases and then run evaluation on the fine-tuned model.
+
+ From here you can start to train on larger or more complex tasks, or else explore other ways to modify ChatGPT-3.5 such as giving it a different tone and style or response.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Resources
+
+ * [OpenAI Fine-Tuning Guide](https://platform.openai.com/docs/guides/fine-tuning)
+ * [W&B Integration with OpenAI API Documentation](https://wandb.me/openai-docs)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/openai-generate-doctor-who-synopses-with-gpt-3-and-weights-biases-video/openai_generate_doctor_who_synopses_with_gpt_3_and_weights_biases_video.py b/marimo/convert/openai-generate-doctor-who-synopses-with-gpt-3-and-weights-biases-video/openai_generate_doctor_who_synopses_with_gpt_3_and_weights_biases_video.py
new file mode 100644
index 00000000..71e23e10
--- /dev/null
+++ b/marimo/convert/openai-generate-doctor-who-synopses-with-gpt-3-and-weights-biases-video/openai_generate_doctor_who_synopses_with_gpt_3_and_weights_biases_video.py
@@ -0,0 +1,441 @@
+# /// script
+# dependencies = ["openai", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ 
+
+
+ # Fine-tune GPT-3 with Weights & Biases to Generate Doctor Who Episode Synopses
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ OpenAI’s API gives practitioners access to GPT-3, an incredibly powerful natural language model that can be applied to virtually any task that involves understanding or generating natural language.
+
+ If you use OpenAI's API to [fine-tune GPT-3](https://beta.openai.com/docs/guides/fine-tuning), you can now use the W&B integration to track experiments, models, and datasets in your central dashboard.
+
+ All it takes is one line: `openai wandb sync`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Set up your API key
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Enter credentials
+ import os
+ os.environ['OPENAI_API_KEY'] = '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'
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **WARNING: Remove the API key after running the cell and clear output so it does not get logged to wandb in case you sync code (see settings)**
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install dependencies
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: openai wandb !pip install -Uq openai wandb
+ return
+
+
+@app.cell
+def _():
+ import openai
+ import wandb
+ from pathlib import Path
+ import pandas as pd
+ import numpy as np
+ import json
+ from tqdm import tqdm
+
+ return openai, pd, tqdm, wandb
+
+
+@app.cell
+def _(wandb):
+ run = wandb.init(project='GPT 3 for Generating Doctor Who Synopses', job_type="dataset_preparation")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Using Weights & Biases Artifacts to Download a .CSV dataset file with episode title > synopsis pairs
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ run_1 = wandb.init(project='GPT 3 for Generating Doctor Who Synopses')
+ artifact = run_1.use_artifact('ivangoncharov/GPT-3 to Generate Doctor Who Synopses/dw_synopses_csv:v0', type='raw_dataset')
+ artifact_dir = artifact.download() + '/dw_synopses.csv'
+ return (artifact_dir,)
+
+
+@app.cell
+def _(artifact_dir, pd, wandb):
+ #Shuffling the dataset with fixed seed
+
+ df = pd.read_csv(artifact_dir)
+ ds = df.sample(frac=1, random_state=0)
+
+
+ wandb.init(project="GPT 3 for Generating Doctor Who Synopses", job_type="logging_dataset_as_table")
+ wandb.run.log({"Raw dataset" : wandb.Table(dataframe=ds)})
+
+ ds.to_csv("dw_synopses.csv")
+ ds.head()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Using OpenAI Tool to preprocess the data
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai tools fine_tunes.prepare_data -f dw_synopses.csv
+ subprocess.call(['openai', 'tools', 'fine_tunes.prepare_data', '-f', 'dw_synopses.csv'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Splitting the data into train and val sets
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #The dataset has 304 pairs in total
+ with open('dw_synopses_prepared.jsonl', 'r') as json_file:
+ json_list = list(json_file)
+ num_data = len(json_list)
+ print('Total:', num_data)
+ val_part = 0.1
+ val_amount = int(num_data * val_part)
+ print('Val data:', val_amount)
+ train_amount = num_data - val_amount
+ print('Train data:', train_amount)
+ subprocess.call(['head', '-n', '$train_amount', 'dw_synopses_prepared.jsonl', '>', 'dw_train.jsonl'])
+ #! head -n $train_amount dw_synopses_prepared.jsonl > dw_train.jsonl
+ #! tail -n $val_amount dw_synopses_prepared.jsonl > dw_valid.jsonl
+ subprocess.call(['tail', '-n', '$val_amount', 'dw_synopses_prepared.jsonl', '>', 'dw_valid.jsonl'])
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Let's define our GPT-3 fine-tuning hyper-parameters.
+ """)
+ return
+
+
+@app.cell
+def _():
+ model = 'curie' # can be ada, babbage or curie
+ n_epochs = 4
+ batch_size = 4
+ learning_rate_multiplier = 0.1
+ prompt_loss_weight = 0.1
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Time to train the model!
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai api fine_tunes.create -t dw_train.jsonl -v dw_valid.jsonl -m $model --n_epochs $n_epochs --batch_size $batch_size --learning_rate_multiplier $learning_rate_multiplier --prompt_loss_weight $prompt_loss_weight
+ subprocess.call(['openai', 'api', 'fine_tunes.create', '-t', 'dw_train.jsonl', '-v', 'dw_valid.jsonl', '-m', '$model', '--n_epochs', '$n_epochs', '--batch_size', '$batch_size', '--learning_rate_multiplier', '$learning_rate_multiplier', '--prompt_loss_weight', '$prompt_loss_weight'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Sync fine-tune jobs to Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can log our fine-tunes with a simple command.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai wandb sync --help
+ subprocess.call(['openai', 'wandb', 'sync', '--help'])
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! openai wandb sync --project "GPT 3 for Generating Doctor Who Synopses"
+ subprocess.call(['openai', 'wandb', 'sync', '--project', 'GPT 3 for Generating Doctor Who Synopses'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Our fine-tunes are now successfully synced to Weights & Biases
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Anytime we have new fine-tunes, we can just call openai wandb sync to add them to our dashboard.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log inference samples
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The best way to evaluate a generative model is to explore sample predictions.
+
+ Let's generate a few inference samples and log them to W&B.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # create eval job
+ run_2 = wandb.init(project='GPT 3 for Generating Doctor Who Synopses', job_type='eval')
+ entity = wandb.run.entity
+ return entity, run_2
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can easily retrieve all config parameters from a job file.
+
+ Job files are logged to W&B as artifacts and can be accessed with `run.use_artifact('USERNAME/PROJECT/job_details:VERSION')` where `VERSION` is either:
+ * a version number such as `v2`
+ * the fine-tune id such as `ft-xxxxxxxxx`
+ * an alias added automatically such as `latest` or manually
+
+ You can explore them in your artifacts dashboard.
+ """)
+ return
+
+
+@app.cell
+def _(entity, run_2):
+ # choose a fine-tuned model
+ artifact_job = run_2.use_artifact(f'{entity}/GPT 3 for Generating Doctor Who Synopses/fine_tune_details:latest', type='fine_tune_details')
+ artifact_job.metadata
+ return (artifact_job,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's take advantage to add metadata into our eval run config.
+ """)
+ return
+
+
+@app.cell
+def _(artifact_job, wandb):
+ wandb.config.update({k:artifact_job.metadata[k] for k in ['fine_tuned_model', 'model', 'hyperparams']})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can easily access model id from any job.
+ """)
+ return
+
+
+@app.cell
+def _(artifact_job):
+ fine_tuned_model = artifact_job.metadata['fine_tuned_model']
+ fine_tuned_model
+ return (fine_tuned_model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Loading validation data as dataframe
+ """)
+ return
+
+
+@app.cell
+def _(pd):
+ df_1 = pd.read_json('dw_valid.jsonl', orient='records', lines=True)
+ df_1.head()
+ return (df_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll perform the inference on all 30 validation examples.
+ """)
+ return
+
+
+@app.cell
+def _(df_1):
+ n_samples = 30
+ df_2 = df_1.iloc[:n_samples]
+ return (df_2,)
+
+
+@app.cell
+def _(df_2, fine_tuned_model, openai, tqdm):
+ data = []
+ for _, row in tqdm(df_2.iterrows()):
+ prompt = row['prompt']
+ res = openai.Completion.create(model=fine_tuned_model, prompt=prompt, max_tokens=300, stop=[' END'])
+ completion = res['choices'][0]['text']
+ completion = completion[1:]
+ prompt = prompt[:-3] # remove initial space
+ target = row['completion'][1:-4] # remove " ->"
+ data.append([prompt, target, completion]) # remove initial space and "END"
+ return (data,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We create and log a W&B Table to easily explore, query & compare model predictions.
+ """)
+ return
+
+
+@app.cell
+def _(data, wandb):
+ prediction_table = wandb.Table(columns=['prompt', 'target', 'completion'], data=data)
+ return (prediction_table,)
+
+
+@app.cell
+def _(prediction_table, wandb):
+ wandb.log({'predictions': prediction_table})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can open the link to your run page down below.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish() #work out a way to print the run page link
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Resources
+ https://wandb.ai/ivangoncharov/GPT-3%20to%20Generate%20Doctor%20Who%20Synopses?workspace=user-ivangoncharov
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/openai-openai-api-autologger-quickstart/openai_openai_api_autologger_quickstart.py b/marimo/convert/openai-openai-api-autologger-quickstart/openai_openai_api_autologger_quickstart.py
new file mode 100644
index 00000000..825d354b
--- /dev/null
+++ b/marimo/convert/openai-openai-api-autologger-quickstart/openai_openai_api_autologger_quickstart.py
@@ -0,0 +1,180 @@
+# /// script
+# dependencies = ["openai", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏃♀️ OpenAI API Logger
+ Use the **[Weights & Biases](https://wandb.ai/site?utm_source=openai_autologger_colab&utm_medium=code&utm_campaign=openai_autologger)** OpenAI API logger to seamlessly log all all inputs and outputs to your OpenAI API. See the full Weights & Biases **[OpenAI Autologger Documentationhere](https://docs.wandb.ai/guides/integrations/openai)** for more
+
+ ### Logging with just 1 line of code
+ With just 1 line of code you can log all of the inputs and outputs from your OpenAI python libray to Weights & Biases for analysis later
+
+ 1️⃣. **Call autolog** and login to wandb with your wandb api key
+
+ 2️⃣. **Run OpenAI API** to generate preductions as normal
+
+ 3️⃣. **Visualize results** by doing to the wandb run link generated in step 1️⃣
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🪄 1. Install `wandb` and `openai`
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb openai !pip install wandb openai -qU
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🪄 2. Import and call `autolog`
+
+ When you call `autolog` you will be prompted to login to Weights & Biases. If you haven't already, create a new API key at [wandb.ai/settings](https://wandb.ai/settings) and store it securely. API keys can only be viewed once when created.
+
+ You can optionally pass a dictionary with arguments for [wandb.init()](https://docs.wandb.ai/ref/python/init) such as a project name, team name, entity, and more. For more information about wandb.init, see the [API Reference Guide](https://docs.wandb.ai/ref/python/init).
+ """)
+ return
+
+
+@app.cell
+def _():
+ import openai
+ from wandb.integration.openai import autolog
+
+ autolog({"project":"my_llm_project"})
+ return autolog, openai
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once you login, a **Weights & Biases run link will be generated**. This will take you to your workspace where you will be able to see all of the inputs and outputs to your OpenAI API calls.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🪄 3. Use the OpenAI API as normal
+ """)
+ return
+
+
+@app.cell
+def _(openai):
+ # pass your OpenAI key
+ openai.api_key = 'sk-foo'
+ return
+
+
+@app.cell
+def _(openai):
+ # make some calls to OpenAI
+ # Call 1
+ chat_request_kwargs = dict(
+ model="gpt-3.5-turbo",
+ messages=[
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Who won the superbowl in 2014?"},
+ {"role": "assistant", "content": "The Seattle Seahawks"},
+ {"role": "user", "content": "Where was it played?"},
+ ],
+ )
+
+ response_1 = openai.ChatCompletion.create(**chat_request_kwargs)
+ print(response_1)
+
+ # Call 2
+ chat_request_kwargs = dict(
+ model="gpt-3.5-turbo",
+ messages=[
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Who won the world series in 2020?"},
+ {"role": "assistant", "content": "The Los Angeles Dodgers"},
+ {"role": "user", "content": "Where was it played?"},
+ ],
+ )
+
+ response_2 = openai.ChatCompletion.create(**chat_request_kwargs)
+ print(response_2)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🪄 4. View your logged results in Weights & Biases
+ - You can find your interactive dashboard by clicking the 👆 wandb links above in step (2)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🪄 5. Disable `autolog`
+ Call disable() to close all W&B processes when you are finished using the OpenAI API.
+ """)
+ return
+
+
+@app.cell
+def _(autolog):
+ autolog.disable()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # What's next 🚀 ?
+ Doing more advanced LLM system chaining or using LangChain?
+ ## 👉 [Try W&B Prompts to understand your LLM systems](https://docs.wandb.ai/guides/prompts?utm_source=openai_api_colab&utm_medium=code&utm_campaign=openai_api_colab)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/openai-set-up-gpt-3-in-python-with-the-openai-api-and-weights-biases/openai_set_up_gpt_3_in_python_with_the_openai_api_and_weights_biases.py b/marimo/convert/openai-set-up-gpt-3-in-python-with-the-openai-api-and-weights-biases/openai_set_up_gpt_3_in_python_with_the_openai_api_and_weights_biases.py
new file mode 100644
index 00000000..e8218d9b
--- /dev/null
+++ b/marimo/convert/openai-set-up-gpt-3-in-python-with-the-openai-api-and-weights-biases/openai_set_up_gpt_3_in_python_with_the_openai_api_and_weights_biases.py
@@ -0,0 +1,95 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.environ['OPENAI_API_KEY'] = ''
+ return (os,)
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ #
+ # !pip install --upgrade openai wandb
+ return
+
+
+@app.cell
+def _(os):
+ import openai
+ import wandb
+ openai.api_key = os.getenv('OPENAI_API_KEY')
+ return openai, wandb
+
+
+@app.cell
+def _(wandb):
+ run = wandb.init(project='GPT-3 App in Python')
+ prediction_table = wandb.Table(columns=["prompt", "completion"])
+ return (prediction_table,)
+
+
+@app.cell
+def _(openai, prediction_table):
+ gpt_prompt = "Correct this to standard English:\n\nShe no went to the market."
+
+
+ response = openai.Completion.create(
+ engine="text-davinci-003",
+ prompt=gpt_prompt,
+ temperature=0.5,
+ max_tokens=256,
+ top_p=1.0,
+ frequency_penalty=0.0,
+ presence_penalty=0.0,
+ )
+
+
+ print(response['choices'][0]['text'])
+
+
+ prediction_table.add_data(gpt_prompt,response['choices'][0]['text'])
+ return
+
+
+@app.cell
+def _(prediction_table, wandb):
+ wandb.log({'predictions': prediction_table})
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paddlepaddle-paddledetection-paddledetection-and-w-b-your-one-stop-for-everything-object-detection/paddlepaddle_paddledetection_paddledetection_and_w_b_your_one_stop_for_everything_object_detection.py b/marimo/convert/paddlepaddle-paddledetection-paddledetection-and-w-b-your-one-stop-for-everything-object-detection/paddlepaddle_paddledetection_paddledetection_and_w_b_your_one_stop_for_everything_object_detection.py
new file mode 100644
index 00000000..f85dd25b
--- /dev/null
+++ b/marimo/convert/paddlepaddle-paddledetection-paddledetection-and-w-b-your-one-stop-for-everything-object-detection/paddlepaddle_paddledetection_paddledetection_and_w_b_your_one_stop_for_everything_object_detection.py
@@ -0,0 +1,401 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train and Debug your Object Detection Models with PaddleDetection and W&B 🪄🐝
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ PaddleDetection is an end-to-end object detection development kit based on PaddlePaddle, which implements varied mainstream object detection, instance segmentation, tracking and keypoint detection algorithms in modular design with configurable modules such as network components, data augmentations and losses.
+
+ This notebook will walk you through how to use the W&B integration in PaddleDetection to track your model training and save model checkpoints.
+
+ The W&B integration in PaddleDetection can be used in two ways
+ - Command Line
+ ```
+ python tools/train -c config.yml --use_wandb -o wandb-project=MyDetector wandb-entity=MyTeam wandb-save_dir=./logs
+ ```
+ The arguments to the W&B logger must be proceeded by -o and each invidiual argument must contain the prefix "wandb-" and the `--use_wandb` flag should be used.
+ - YAML File: Add the arguments to the config.yml file under the wandb header like this
+ ```
+ wandb:
+ project: MyProject
+ entity: MyTeam
+ save_dir: ./logs
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup 🖥
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We begin by installing the PaddlePaddle library followed by PaddleDetection
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%shell
+ # git clone --depth 1 https://github.com/PaddlePaddle/PaddleDetection
+ # python -m pip install paddlepaddle-gpu==2.3.2.post112 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html
+ # pip install pyclipper attrdict gdown -qqq
+ # cd PaddleDetection
+ # git checkout develop
+ # pip install -q -e .
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.chdir('PaddleDetection')
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now we'll install and log into W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ A few more imports and we are good to go!
+ """)
+ return
+
+
+@app.cell
+def _():
+ import glob
+ import yaml
+
+ return glob, yaml
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Dataset 💿
+
+ We begin by downloading a [subset of the COCO 2017](https://wandb.ai/manan-goel/PaddleDetectionYOLOX/artifacts/dataset/COCOSubset/62e41eb3f1ebafcbd5a7) dataset created using this amazing [repository](https://github.com/giddyyupp/coco-minitrain/blob/master/src/sample_coco.py)
+
+ The dataset has been logged as a W&B artifact for easier downloading. It contains 1000 images for training and 250 for validation with corresponding annotations which we will now use for training our object detection model.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _artifact = wandb.Api().artifact('manan-goel/PaddleDetectionYOLOX/COCOSubset:latest')
+ path = _artifact.download(root='./dataset/coco')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training 🏋️♀️
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Setting up the configuration 🛠
+
+ The configuration for the model is provided to the training script through yaml files in PaddleDetection so we're going to go ahead and edit it to suit our needs.
+ """)
+ return
+
+
+@app.cell
+def _(yaml):
+ with open("./configs/yolox/yolox_nano_300e_coco.yml", "r") as f:
+ config = yaml.safe_load(f)
+
+
+ # Any arguments to wandb.init can be provided in the config['wandb'] dict
+
+ config['wandb'] = {
+ 'project': 'PaddleDetectionYOLOX'
+ }
+
+ config['log_iter'] = 1
+ config['snapshot_epoch'] = 5
+ config['epoch'] = 5
+ config['TrainReader']['batch_size'] = 32
+
+ with open("./configs/yolox/yolox_nano_300e_coco.yml", "w") as f:
+ yaml.safe_dump(config, f)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Training the Model 🏋️♀️
+
+ We now use the training script in the PaddleDetection library to train the YOLOX model and the config above adds W&B logging during training. We also add the `--eval` flag to have an evaluation step every 5 epochs since the validation set is extremely large.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python tools/train.py -c configs/yolox/yolox_nano_300e_coco.yml --eval
+ subprocess.call(['python', 'tools/train.py', '-c', 'configs/yolox/yolox_nano_300e_coco.yml', '--eval'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Visualization
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ During training the metrics on the training and validation sets are logged to a W&B dashboard which looks something like this
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The model checkpoints are also logged at the end of every snapshot epoch along with the corresponding average precision.
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Evaluation and Testing 🤔
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Download the Best Model from W&B 💾
+
+ To download the model with the best accuracy from W&B for evaluation you can use the following snippet
+ """)
+ return
+
+
+@app.cell
+def _(os, wandb):
+ _artifact = wandb.Api().artifact('manan-goel/PaddleDetectionYOLOX/model-26oqc38r:best', type='model')
+ artifact_dir = _artifact.download()
+ os.rename(artifact_dir + '/model', artifact_dir + '/model.pdparams')
+ artifact_dir
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Testing and Logging annotated images to your W&B dashboard 🔥
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The following cell runs the inference script on all the images in `demo` directory and stores the annotated images with the bounding boxes in the `infer_output` directory using the YOLOX model pulled from W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%shell
+ # for i in $(ls demo/*.jpg)
+ # do
+ # python tools/infer.py -c configs/yolox/yolox_nano_300e_coco.yml \
+ # --infer_img=$i \
+ # --output_dir=infer_output/ \
+ # --draw_threshold=0.5 \
+ # -o weights=./artifacts/model-26oqc38r:v1/model
+ # done
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ With the images being annotated, we will now initialize a new run in the project and a W&B table which we will then use to add the input and output images side-by-side.
+ """)
+ return
+
+
+@app.cell
+def _(glob, wandb):
+ wandb.init(project="PaddleDetectionYOLOX")
+ wandb.use_artifact('manan-goel/PaddleDetectionYOLOX/model-26oqc38r:best')
+ table = wandb.Table(columns=["Input Image", "Annotated Image"])
+
+ inp_imgs = sorted(glob.glob("./demo/*.jpg"), key=lambda x: x.split("/")[-1])
+ out_imgs = sorted(glob.glob("./infer_output/*.jpg"), key=lambda x: x.split("/")[-1])
+ return inp_imgs, out_imgs, table
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The next cell iterates over the images and adds the corresponding inputs and outputs as a row in the table initialized in the previous cell.
+ """)
+ return
+
+
+@app.cell
+def _(inp_imgs, out_imgs, table, wandb):
+ for inp in inp_imgs:
+ for out in out_imgs:
+ if out.split("/")[-1] != inp.split("/")[-1]:
+ continue
+ table.add_data(
+ wandb.Image(inp),
+ wandb.Image(out)
+ )
+ wandb.log({
+ "Predictions": table
+ })
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This provides a really cool debugging tool!
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Resources📕
+
+ * [W&B and PaddleDetection Documentation](https://docs.wandb.ai/guides/integrations/other/paddledetection) contains a few tips for taking most advantage of W&B.
+ * More PaddleDetection documentation is available [here](https://github.com/PaddlePaddle/PaddleDetection)
+ """)
+ 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, please contact support@wandb.com
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paddlepaddle-paddleocr-train-and-debug-your-ocr-models-with-paddleocr-and-w-b/paddlepaddle_paddleocr_train_and_debug_your_ocr_models_with_paddleocr_and_w_b.py b/marimo/convert/paddlepaddle-paddleocr-train-and-debug-your-ocr-models-with-paddleocr-and-w-b/paddlepaddle_paddleocr_train_and_debug_your_ocr_models_with_paddleocr_and_w_b.py
new file mode 100644
index 00000000..66f21bff
--- /dev/null
+++ b/marimo/convert/paddlepaddle-paddleocr-train-and-debug-your-ocr-models-with-paddleocr-and-w-b/paddlepaddle_paddleocr_train_and_debug_your_ocr_models_with_paddleocr_and_w_b.py
@@ -0,0 +1,434 @@
+# /// script
+# dependencies = ["&&", "-", "attrdict", "https://www-paddlepaddle-org-cn/whl/linux/mkl/avx/stable-html", "imgaug", "install", "opencv-python", "paddlepaddle-gpu", "pip", "pyclipper", "requirements-txt", "scikit-image", "shapely", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train and Debug Your OCR Models using PaddleOCR and Weights & Biases 🪄🐝
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This notebook talks about how you can use W&B with PaddleOCR to track training metrics and log model checkpoints for all your OCR needs!
+
+ To use the W&B logger with the PaddleOCR training script just add the following at the bottom of your `config.yml` file.
+
+ ```
+ wandb:
+ project: CoolOCR
+ entity: my_team
+ name: MyOCRModel
+ ```
+
+ To log the metrics and checkpoints to W&B during training, the wandb client now has a direct integration into PaddleOCR. Using wandb for logging automatically adds all the metrics to your W&B dashboard, saves the models at every evaluation step, tags the best model and adds appropriate metadata for the saved model. An example dashboard is available [here](https://wandb.ai/manan-goel/text_detection).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup 🖥
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We begin by cloning the PaddleOCR library and installing the the package.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! git clone --depth 1 https://github.com/PaddlePaddle/PaddleOCR
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/PaddlePaddle/PaddleOCR'])
+ # packages added via marimo's package management: paddlepaddle-gpu==2.3.2.post112 https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html !python -m pip install paddlepaddle-gpu==2.3.2.post112 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html
+ # packages added via marimo's package management: pyclipper attrdict opencv-python shapely scikit-image imgaug !pip install pyclipper attrdict opencv-python shapely scikit-image imgaug -qqq
+ # packages added via marimo's package management: requirements.txt && pip install . !cd PaddleOCR && pip install -r requirements.txt && pip install -e .
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now we'll install and log into W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.chdir('PaddleOCR')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training 🏋️♀️
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ PaddleOCR comes with a huge array of pre-implemented models involved in the OCR pipeline. For this tutorial we will be looking at the text detection models.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Downloading Training and Validation Data 💾
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will use the ICDAR2015 dataset available [here](https://rrc.cvc.uab.es/?ch=4&com=downloads). The data has been logged as W&B artifacts for ease of use.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ api = wandb.Api()
+ _artifact = api.artifact('manan-goel/icdar2015/icdar2015-dataset:latest')
+ _artifact.download(root='./train_data/icdar2015')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Downloading pretrained weights📈
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! wget -P ./pretrain_models/ https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV3_large_x0_5_pretrained.pdparams
+ subprocess.call(['wget', '-P', './pretrain_models/', 'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV3_large_x0_5_pretrained.pdparams'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Setup the config.yml file to use W&B🛠
+ """)
+ return
+
+
+@app.cell
+def _():
+ import yaml
+
+ with open("configs/det/det_mv3_db.yml", "r") as f:
+ config = yaml.safe_load(f)
+ config.update({
+ 'wandb': {
+ 'project': 'text_detection'
+ }
+ })
+ config['Global'].update({
+ 'epoch_num': 5,
+ 'eval_batch_step': [0, 50],
+ 'calc_metric_during_train': True
+ })
+
+ with open("configs/det/det_mv3_db.yml", "w") as f:
+ yaml.safe_dump(config, f)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Train your Model 🏋️♀️
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The following command will finetune the pretrained MobileNetV3 on the ICDAR2015 dataset while logging all training and validation metrics to a W&B dashboard.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python3 tools/train.py -c configs/det/det_mv3_db.yml -o Global.pretrained_model=./pretrain_models/MobileNetV3_large_x0_5_pretrained
+ subprocess.call(['python3', 'tools/train.py', '-c', 'configs/det/det_mv3_db.yml', '-o', 'Global.pretrained_model=./pretrain_models/MobileNetV3_large_x0_5_pretrained'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The [dashboard](https://wandb.ai/manan-goel/text_detection) would look something like this
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The checkpoints at the end of every epoch and evaluation step are also logged to W&B with appropriate metadata!
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Evaluation and Testing 🤔
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Downloading the best model from W&B 💿
+
+ To download the model with the best accuracy from W&B for evaluation you can use the following snippet
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _artifact = wandb.Api().artifact('manan-goel/text_detection/model-2138qk4h:best', type='model')
+ artifact_dir = _artifact.download()
+ return (artifact_dir,)
+
+
+@app.cell
+def _(artifact_dir):
+ artifact_dir
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Running the Evaluation Script
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will now use the checkpoint downloaded in the previous step to run the evaluation script
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python3 tools/eval.py -c configs/det/det_mv3_db.yml -o Global.checkpoints="./artifacts/model-2138qk4h:v9/model_ckpt" PostProcess.box_thresh=0.6 PostProcess.unclip_ratio=1.5
+ subprocess.call(['python3', 'tools/eval.py', '-c', 'configs/det/det_mv3_db.yml', '-o', 'Global.checkpoints=./artifacts/model-2138qk4h:v9/model_ckpt', 'PostProcess.box_thresh=0.6', 'PostProcess.unclip_ratio=1.5'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Testing on Images 🧪
+
+ To test the model training we run the detection script from PaddleOCR on some test images.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python3 tools/infer_det.py -c configs/det/det_mv3_db.yml -o Global.infer_img="./doc/imgs_en/" Global.pretrained_model="./artifacts/model-2138qk4h:v9/model_ckpt"
+ subprocess.call(['python3', 'tools/infer_det.py', '-c', 'configs/det/det_mv3_db.yml', '-o', 'Global.infer_img=./doc/imgs_en/', 'Global.pretrained_model=./artifacts/model-2138qk4h:v9/model_ckpt'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Bonus: Logging annotated images to your W&B dashboard 🔥
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project="text_detection")
+ wandb.use_artifact('manan-goel/text_detection/model-2138qk4h:best')
+ return
+
+
+@app.cell
+def _(wandb):
+ table = wandb.Table(columns=["Input Image", "Annotated Image"])
+ return (table,)
+
+
+@app.cell
+def _():
+ import glob
+ inp_imgs = sorted(glob.glob("./doc/imgs_en/*.jpg"), key=lambda x: x.split("/")[-1])
+ out_imgs = sorted(glob.glob("./output/det_db/det_results/*.jpg"), key=lambda x: x.split("/")[-1])
+ return inp_imgs, out_imgs
+
+
+@app.cell
+def _(inp_imgs, out_imgs, table, wandb):
+ for inp in inp_imgs:
+ for out in out_imgs:
+ if out.split("/")[-1] != inp.split("/")[-1]:
+ continue
+ table.add_data(
+ wandb.Image(inp),
+ wandb.Image(out)
+ )
+ return
+
+
+@app.cell
+def _(table, wandb):
+ wandb.log({
+ "Predictions": table
+ })
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This gives a really cool visualization and debugging tool!
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Resources📕
+
+ * [W&B and PaddleOCR Documentation](https://docs.wandb.ai/guides/integrations/other/paddleocr) contains a few tips for taking most advantage of W&B.
+ * More PaddleOCR documentation is available [here](https://github.com/PaddlePaddle/PaddleOCR)
+ """)
+ 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, please contact support@wandb.com
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paella-image-variations/paella_image_variations.py b/marimo/convert/paella-image-variations/paella_image_variations.py
new file mode 100644
index 00000000..6c3bfa50
--- /dev/null
+++ b/marimo/convert/paella-image-variations/paella_image_variations.py
@@ -0,0 +1,304 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Image Variations with Paella + WandB Playground 🪄🐝
+
+
+
+ A demo of Image Vairations using [Paella](https://github.com/dome272/Paella) and [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import time
+ import wandb
+ import requests
+ import numpy as np
+ from PIL import Image
+ from io import BytesIO
+ from tqdm.notebook import tqdm
+ import matplotlib.pyplot as plt
+
+ import torch
+ from torch import nn
+ import torchvision
+
+ import open_clip
+ from rudalle import get_vae
+ from einops import rearrange
+ from open_clip import tokenizer
+
+ from Paella.modules import DenoiseUNet
+
+ return (
+ DenoiseUNet,
+ Image,
+ get_vae,
+ np,
+ open_clip,
+ os,
+ rearrange,
+ requests,
+ time,
+ torch,
+ torchvision,
+ wandb,
+ )
+
+
+@app.cell
+def _(os, torch, wandb):
+ wandb_project = "paella" #@param {"type": "string"}
+ wandb_entity = "geekyrakshit" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, entity=wandb_entity, job_type="image-variations")
+
+
+ config = wandb.config
+ config.model_artifact = "geekyrakshit/paella/fine-tuned-image-model:v0"
+ config.seed = 42
+ config.batch_size = 5
+ config.latent_shape = (32, 32)
+ config.image_url = "https://media.istockphoto.com/id/1193591781/photo/obedient-dog-breed-welsh-corgi-pembroke-sitting-and-smiles-on-a-white-background-not-isolate.jpg?s=612x612&w=0&k=20&c=ZDKTgSFQFG9QvuDziGsnt55kvQoqJtIhrmVRkpYqxtQ="
+ config.prompt = "a delicious spanish paella"
+ config.target_size = 224
+ config.batch_size = 5
+ config.latent_shape = (32, 32)
+
+
+ # Seed Everything
+ torch.manual_seed(config.seed)
+ torch.random.manual_seed(config.seed)
+ torch.cuda.manual_seed(config.seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = True
+
+ # Download Model from Weights & Biases Artifacts
+ text_model_path = os.path.join(wandb.use_artifact(config.model_artifact, type='model').download(), "model_50000_img.pt")
+ return config, text_model_path
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print("Using device:", device)
+ return (device,)
+
+
+@app.cell
+def _(config, np, wandb):
+ def log_image_variations_results(input_image, generated_images):
+ generated_images = [
+ wandb.Image(image)
+ for image in (generated_images.cpu().numpy() * 255.0).astype(np.uint8)
+ ]
+ table = wandb.Table(
+ columns=["Seed", "URL", "Input-Image", "Latent-Shape", "Generated-Image"]
+ )
+ table.add_data(
+ config.seed,
+ config.image_url,
+ wandb.Image(input_image),
+ config.latent_shape,
+ generated_images
+ )
+ wandb.log({"Image-Variations-Results": table})
+
+ return (log_image_variations_results,)
+
+
+@app.cell
+def _(torch):
+ def log(t, eps=1e-20):
+ return torch.log(t + eps)
+
+ def gumbel_noise(t):
+ noise = torch.zeros_like(t).uniform_(0, 1)
+ return -log(-log(noise))
+
+ def gumbel_sample(t, temperature=1., dim=-1):
+ return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+ def sample(
+ model, c, x=None, mask=None, T=12, size=(32, 32),
+ starting_t=0, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=-1,
+ renoise_steps=11, renoise_mode='start'
+ ):
+ with torch.inference_mode():
+ r_range = torch.linspace(0, 1, T+1)[:-1][:, None].expand(-1, c.size(0)).to(c.device)
+ temperatures = torch.linspace(temp_range[0], temp_range[1], T)
+ preds = []
+ if x is None:
+ x = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ elif mask is not None:
+ noise = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ x = noise * mask + (1-mask) * x
+ init_x = x.clone()
+ for i in range(starting_t, T):
+ if renoise_mode == 'prev':
+ prev_x = x.clone()
+ r, temp = r_range[i], temperatures[i]
+ logits = model(x, c, r)
+ if classifier_free_scale >= 0:
+ logits_uncond = model(x, torch.zeros_like(c), r)
+ logits = torch.lerp(logits_uncond, logits, classifier_free_scale)
+ x = logits
+ x_flat = x.permute(0, 2, 3, 1).reshape(-1, x.size(1))
+ if typical_filtering:
+ x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
+ x_flat_norm_p = torch.exp(x_flat_norm)
+ entropy = -(x_flat_norm * x_flat_norm_p).nansum(-1, keepdim=True)
+
+ c_flat_shifted = torch.abs((-x_flat_norm) - entropy)
+ c_flat_sorted, x_flat_indices = torch.sort(c_flat_shifted, descending=False)
+ x_flat_cumsum = x_flat.gather(-1, x_flat_indices).softmax(dim=-1).cumsum(dim=-1)
+
+ last_ind = (x_flat_cumsum < typical_mass).sum(dim=-1)
+ sorted_indices_to_remove = c_flat_sorted > c_flat_sorted.gather(1, last_ind.view(-1, 1))
+ if typical_min_tokens > 1:
+ sorted_indices_to_remove[..., :typical_min_tokens] = 0
+ indices_to_remove = sorted_indices_to_remove.scatter(1, x_flat_indices, sorted_indices_to_remove)
+ x_flat = x_flat.masked_fill(indices_to_remove, -float("Inf"))
+ x_flat = gumbel_sample(x_flat, temperature=temp)
+ x = x_flat.view(x.size(0), *x.shape[2:])
+ if mask is not None:
+ x = x * mask + (1-mask) * init_x
+ if i < renoise_steps:
+ if renoise_mode == 'start':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=init_x)
+ elif renoise_mode == 'prev':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=prev_x)
+ else: # 'rand'
+ x, _ = model.add_noise(x, r_range[i+1])
+ preds.append(x.detach())
+ return preds
+
+ return (sample,)
+
+
+@app.cell
+def _(
+ DenoiseUNet,
+ device,
+ get_vae,
+ open_clip,
+ rearrange,
+ text_model_path,
+ torch,
+):
+ vqmodel = get_vae().to(device)
+ vqmodel.eval().requires_grad_(False)
+
+ clip_model, _, _ = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s12b_b42k')
+ clip_model = clip_model.to(device).eval().requires_grad_(False)
+
+
+ def encode(x):
+ return vqmodel.model.encode((2 * x - 1))[-1][-1]
+
+ def decode(img_seq, shape=(32,32)):
+ img_seq = img_seq.view(img_seq.shape[0], -1)
+ b, n = img_seq.shape
+ one_hot_indices = torch.nn.functional.one_hot(img_seq, num_classes=vqmodel.num_tokens).float()
+ z = (one_hot_indices @ vqmodel.model.quantize.embed.weight)
+ z = rearrange(z, 'b (h w) c -> b c h w', h=shape[0], w=shape[1])
+ img = vqmodel.model.decode(z)
+ img = (img.clamp(-1., 1.) + 1) * 0.5
+ return img
+
+ state_dict = torch.load(text_model_path, map_location=device)
+ model = DenoiseUNet(num_labels=8192).to(device)
+ model.load_state_dict(state_dict)
+ model.eval().requires_grad_()
+ print()
+ return clip_model, decode, model
+
+
+@app.cell
+def _(Image, config, device, requests, torchvision):
+ response = requests.get(config.image_url)
+ # original_image = Image.open(BytesIO(response.content)).convert("RGB")
+ original_image = Image.open("pexels-hiếu-hoàng-954050.jpg").convert("RGB")
+
+ preprocess = torchvision.transforms.Compose([
+ torchvision.transforms.Resize(config.target_size),
+ torchvision.transforms.ToTensor(),
+ ])
+
+ clip_preprocess = torchvision.transforms.Compose([
+ torchvision.transforms.Resize(
+ (config.target_size, config.target_size),
+ interpolation=torchvision.transforms.InterpolationMode.BICUBIC
+ ),
+ torchvision.transforms.Normalize(
+ mean=(0.48145466, 0.4578275, 0.40821073),
+ std=(0.26862954, 0.26130258, 0.27577711)
+ ),
+ ])
+
+ images = preprocess(original_image).unsqueeze(0).expand(config.batch_size, -1, -1, -1).to(device)[:, :3]
+ return clip_preprocess, images, original_image
+
+
+@app.cell
+def _(
+ clip_model,
+ clip_preprocess,
+ config,
+ decode,
+ images,
+ model,
+ sample,
+ time,
+ torch,
+ wandb,
+):
+ with torch.inference_mode():
+ with torch.autocast(device_type="cuda"):
+ clip_embeddings = clip_model.encode_image(clip_preprocess(images)).float()
+ s = time.time()
+ sampled = sample(
+ model, clip_embeddings, T=12, size=config.latent_shape, starting_t=0,
+ temp_range=[1.0, 1.0], typical_filtering=True, typical_mass=0.2,
+ typical_min_tokens=1, classifier_free_scale=5, renoise_steps=11
+ )
+ wandb.log({"Sampling-Time": time.time() - s})
+ sampled = decode(sampled[-1], config.latent_shape).permute(0, 2, 3, 1)
+ return (sampled,)
+
+
+@app.cell
+def _(log_image_variations_results, original_image, sampled, wandb):
+ log_image_variations_results(original_image, sampled)
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paella-multi-conditioning/paella_multi_conditioning.py b/marimo/convert/paella-multi-conditioning/paella_multi_conditioning.py
new file mode 100644
index 00000000..4232cb9b
--- /dev/null
+++ b/marimo/convert/paella-multi-conditioning/paella_multi_conditioning.py
@@ -0,0 +1,285 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Multi-Conditional Image Generation with Paella + WandB Playground 🪄🐝
+
+
+
+ A demo of Multi-Conditional Image Generation using [Paella](https://github.com/dome272/Paella) and [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import time
+ import wandb
+ import requests
+ import numpy as np
+ from PIL import Image
+ from io import BytesIO
+ from tqdm.notebook import tqdm
+ import matplotlib.pyplot as plt
+
+ import torch
+ from torch import nn
+ import torchvision
+
+ import open_clip
+ from rudalle import get_vae
+ from einops import rearrange
+ from open_clip import tokenizer
+
+ from Paella.modules import DenoiseUNet
+
+ return (
+ DenoiseUNet,
+ get_vae,
+ np,
+ open_clip,
+ os,
+ rearrange,
+ time,
+ tokenizer,
+ torch,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell
+def _(os, torch, wandb):
+ wandb_project = "paella" #@param {"type": "string"}
+ wandb_entity = "geekyrakshit" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, entity=wandb_entity, job_type="multi-conditioning")
+
+
+ config = wandb.config
+ config.model_artifact = "geekyrakshit/paella/text-model:v0"
+ config.seed = 42
+ config.batch_size = 5
+ config.latent_shape = (32, 100)
+ config.conditions = [
+ ["a princess stuck inside a castle, mario style, pixel art", 30],
+ ["a dragon, mario style, pixel art", 60],
+ ["mario dodging obstacles, mario style, pixel art", 100],
+ ]
+ config.clip_embedding_dim = 1024
+
+ # Seed Everything
+ torch.manual_seed(config.seed)
+ torch.random.manual_seed(config.seed)
+ torch.cuda.manual_seed(config.seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = True
+
+ # Download Model from Weights & Biases Artifacts
+ text_model_path = os.path.join(wandb.use_artifact(config.model_artifact, type='model').download(), "model_600000.pt")
+ return config, text_model_path
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print("Using device:", device)
+ return (device,)
+
+
+@app.cell
+def _(config, np, wandb):
+ def log_multi_conditioning_results(images):
+ images = [wandb.Image(image) for image in (images.cpu().numpy() * 255.0).astype(np.uint8)]
+ table = wandb.Table(columns=["Seed", "Prompts", "End-Token-Positions", "Latent-Shape", "Generated-Images"])
+ table.add_data(
+ config.seed,
+ [condition[0] for condition in config.conditions],
+ [condition[1] for condition in config.conditions],
+ config.latent_shape,
+ images
+ )
+ wandb.log({"Multi-Conditioning-Results": table})
+
+ return (log_multi_conditioning_results,)
+
+
+@app.cell
+def _(torch):
+ def log(t, eps=1e-20):
+ return torch.log(t + eps)
+
+ def gumbel_noise(t):
+ noise = torch.zeros_like(t).uniform_(0, 1)
+ return -log(-log(noise))
+
+ def gumbel_sample(t, temperature=1., dim=-1):
+ return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+ def sample(
+ model, c, x=None, mask=None, T=12, size=(32, 32),
+ starting_t=0, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=-1,
+ renoise_steps=11, renoise_mode='start'
+ ):
+ with torch.inference_mode():
+ r_range = torch.linspace(0, 1, T+1)[:-1][:, None].expand(-1, c.size(0)).to(c.device)
+ temperatures = torch.linspace(temp_range[0], temp_range[1], T)
+ preds = []
+ if x is None:
+ x = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ elif mask is not None:
+ noise = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ x = noise * mask + (1-mask) * x
+ init_x = x.clone()
+ for i in range(starting_t, T):
+ if renoise_mode == 'prev':
+ prev_x = x.clone()
+ r, temp = r_range[i], temperatures[i]
+ logits = model(x, c, r)
+ if classifier_free_scale >= 0:
+ logits_uncond = model(x, torch.zeros_like(c), r)
+ logits = torch.lerp(logits_uncond, logits, classifier_free_scale)
+ x = logits
+ x_flat = x.permute(0, 2, 3, 1).reshape(-1, x.size(1))
+ if typical_filtering:
+ x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
+ x_flat_norm_p = torch.exp(x_flat_norm)
+ entropy = -(x_flat_norm * x_flat_norm_p).nansum(-1, keepdim=True)
+
+ c_flat_shifted = torch.abs((-x_flat_norm) - entropy)
+ c_flat_sorted, x_flat_indices = torch.sort(c_flat_shifted, descending=False)
+ x_flat_cumsum = x_flat.gather(-1, x_flat_indices).softmax(dim=-1).cumsum(dim=-1)
+
+ last_ind = (x_flat_cumsum < typical_mass).sum(dim=-1)
+ sorted_indices_to_remove = c_flat_sorted > c_flat_sorted.gather(1, last_ind.view(-1, 1))
+ if typical_min_tokens > 1:
+ sorted_indices_to_remove[..., :typical_min_tokens] = 0
+ indices_to_remove = sorted_indices_to_remove.scatter(1, x_flat_indices, sorted_indices_to_remove)
+ x_flat = x_flat.masked_fill(indices_to_remove, -float("Inf"))
+ # x_flat = torch.multinomial(x_flat.div(temp).softmax(-1), num_samples=1)[:, 0]
+ x_flat = gumbel_sample(x_flat, temperature=temp)
+ x = x_flat.view(x.size(0), *x.shape[2:])
+ if mask is not None:
+ x = x * mask + (1-mask) * init_x
+ if i < renoise_steps:
+ if renoise_mode == 'start':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=init_x)
+ elif renoise_mode == 'prev':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=prev_x)
+ else: # 'rand'
+ x, _ = model.add_noise(x, r_range[i+1])
+ preds.append(x.detach())
+ return preds
+
+ return (sample,)
+
+
+@app.cell
+def _(
+ DenoiseUNet,
+ device,
+ get_vae,
+ open_clip,
+ rearrange,
+ text_model_path,
+ torch,
+):
+ vqmodel = get_vae().to(device)
+ vqmodel.eval().requires_grad_(False)
+
+ clip_model, _, _ = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s12b_b42k')
+ clip_model = clip_model.to(device).eval().requires_grad_(False)
+
+
+ def encode(x):
+ return vqmodel.model.encode((2 * x - 1))[-1][-1]
+
+ def decode(img_seq, shape=(32,32)):
+ img_seq = img_seq.view(img_seq.shape[0], -1)
+ b, n = img_seq.shape
+ one_hot_indices = torch.nn.functional.one_hot(img_seq, num_classes=vqmodel.num_tokens).float()
+ z = (one_hot_indices @ vqmodel.model.quantize.embed.weight)
+ z = rearrange(z, 'b (h w) c -> b c h w', h=shape[0], w=shape[1])
+ img = vqmodel.model.decode(z)
+ img = (img.clamp(-1., 1.) + 1) * 0.5
+ return img
+
+ state_dict = torch.load(text_model_path, map_location=device)
+ model = DenoiseUNet(num_labels=8192).to(device)
+ model.load_state_dict(state_dict)
+ model.eval().requires_grad_()
+ print()
+ return clip_model, decode, model
+
+
+@app.cell
+def _(
+ clip_model,
+ config,
+ decode,
+ device,
+ model,
+ sample,
+ time,
+ tokenizer,
+ torch,
+ tqdm,
+ wandb,
+):
+ clip_embedding = torch.zeros(
+ config.batch_size,
+ config.clip_embedding_dim,
+ *config.latent_shape
+ ).to(device)
+
+ last_pos = 0
+ for text, pos in tqdm(config.conditions):
+ tokenized_text = tokenizer.tokenize([text] * config.batch_size).to(device)
+ part_clip_embedding = clip_model.encode_text(tokenized_text).float()[:, :, None, None]
+ clip_embedding[:, :, :, last_pos:pos] = part_clip_embedding
+ last_pos = pos
+ with torch.inference_mode():
+ with torch.autocast(device_type="cuda"):
+ s = time.time()
+ sampled = sample(
+ model, clip_embedding, T=12, size=config.latent_shape, starting_t=0,
+ temp_range=[1.0, 1.0], typical_filtering=True, typical_mass=0.2,
+ typical_min_tokens=1, classifier_free_scale=5, renoise_steps=11, renoise_mode="start"
+ )
+ wandb.log({"Sampling-Time": time.time() - s})
+ sampled = decode(sampled[-1], config.latent_shape).permute(0, 2, 3, 1)
+ return (sampled,)
+
+
+@app.cell
+def _(log_multi_conditioning_results, sampled, wandb):
+ log_multi_conditioning_results(sampled)
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paella-orientation-guided-multi-conditioning/paella_orientation_guided_multi_conditioning.py b/marimo/convert/paella-orientation-guided-multi-conditioning/paella_orientation_guided_multi_conditioning.py
new file mode 100644
index 00000000..c4055e49
--- /dev/null
+++ b/marimo/convert/paella-orientation-guided-multi-conditioning/paella_orientation_guided_multi_conditioning.py
@@ -0,0 +1,302 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Orientation Guided Multi-Conditional Image Generation with Paella + WandB Playground 🪄🐝
+
+
+
+ A demo of Orientation Guided Multi-Conditional Image Generation using [Paella](https://github.com/dome272/Paella) and [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import time
+ import wandb
+ import requests
+ import numpy as np
+ from PIL import Image
+ from io import BytesIO
+ from tqdm.notebook import tqdm
+ import matplotlib.pyplot as plt
+
+ import torch
+ from torch import nn
+ import torchvision
+
+ import open_clip
+ from rudalle import get_vae
+ from einops import rearrange
+ from open_clip import tokenizer
+
+ from Paella.modules import DenoiseUNet
+
+ return (
+ DenoiseUNet,
+ get_vae,
+ np,
+ open_clip,
+ os,
+ rearrange,
+ time,
+ tokenizer,
+ torch,
+ wandb,
+ )
+
+
+@app.cell
+def _(os, torch, wandb):
+ wandb_project = "paella" #@param {"type": "string"}
+ wandb_entity = "geekyrakshit" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, entity=wandb_entity, job_type="orientation-guided-multi-conditioning")
+
+
+ config = wandb.config
+ config.model_artifact = "geekyrakshit/paella/text-model:v0"
+ config.seed = 42
+ config.batch_size = 5
+ config.latent_shape = (32, 32)
+ config.prompt_1 = "a cute portrait of a dog"
+ config.prompt_2 = "a cute portrait of a cat"
+ config.orientation_mode = "horizontal" # ["vertical", "horizontal"]
+ config.interpolation_mode = "spherical-lerp" # ["lerp", "spherical-lerp"]
+
+ # Seed Everything
+ torch.manual_seed(config.seed)
+ torch.random.manual_seed(config.seed)
+ torch.cuda.manual_seed(config.seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = True
+
+ # Download Model from Weights & Biases Artifacts
+ text_model_path = os.path.join(wandb.use_artifact(config.model_artifact, type='model').download(), "model_600000.pt")
+ return config, text_model_path
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print("Using device:", device)
+ return (device,)
+
+
+@app.cell
+def _(config, np, wandb):
+ def log_orientation_guided_multi_conditioning_results(images):
+ images = [wandb.Image(image) for image in (images.cpu().numpy() * 255.0).astype(np.uint8)]
+ table = wandb.Table(
+ columns=[
+ "Seed",
+ "Prompt-1", "Prompt-2",
+ "Orientation-Mode", "Interpolation-Mode", "Latent-Shape",
+ "Generated-Images"
+ ]
+ )
+ table.add_data(
+ config.seed,
+ config.prompt_1, config.prompt_2,
+ config.orientation_mode, config.interpolation_mode, config.latent_shape,
+ images
+ )
+ wandb.log({"Orientation-Guided-Multi-Conditioning-Results": table})
+
+ return (log_orientation_guided_multi_conditioning_results,)
+
+
+@app.cell
+def _(torch):
+ def log(t, eps=1e-20):
+ return torch.log(t + eps)
+
+ def gumbel_noise(t):
+ noise = torch.zeros_like(t).uniform_(0, 1)
+ return -log(-log(noise))
+
+ def gumbel_sample(t, temperature=1., dim=-1):
+ return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+ def sample(
+ model, c, x=None, mask=None, T=12, size=(32, 32),
+ starting_t=0, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=-1,
+ renoise_steps=11, renoise_mode='start'
+ ):
+ with torch.inference_mode():
+ r_range = torch.linspace(0, 1, T+1)[:-1][:, None].expand(-1, c.size(0)).to(c.device)
+ temperatures = torch.linspace(temp_range[0], temp_range[1], T)
+ preds = []
+ if x is None:
+ x = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ elif mask is not None:
+ noise = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ x = noise * mask + (1-mask) * x
+ init_x = x.clone()
+ for i in range(starting_t, T):
+ if renoise_mode == 'prev':
+ prev_x = x.clone()
+ r, temp = r_range[i], temperatures[i]
+ logits = model(x, c, r)
+ if classifier_free_scale >= 0:
+ logits_uncond = model(x, torch.zeros_like(c), r)
+ logits = torch.lerp(logits_uncond, logits, classifier_free_scale)
+ x = logits
+ x_flat = x.permute(0, 2, 3, 1).reshape(-1, x.size(1))
+ if typical_filtering:
+ x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
+ x_flat_norm_p = torch.exp(x_flat_norm)
+ entropy = -(x_flat_norm * x_flat_norm_p).nansum(-1, keepdim=True)
+
+ c_flat_shifted = torch.abs((-x_flat_norm) - entropy)
+ c_flat_sorted, x_flat_indices = torch.sort(c_flat_shifted, descending=False)
+ x_flat_cumsum = x_flat.gather(-1, x_flat_indices).softmax(dim=-1).cumsum(dim=-1)
+
+ last_ind = (x_flat_cumsum < typical_mass).sum(dim=-1)
+ sorted_indices_to_remove = c_flat_sorted > c_flat_sorted.gather(1, last_ind.view(-1, 1))
+ if typical_min_tokens > 1:
+ sorted_indices_to_remove[..., :typical_min_tokens] = 0
+ indices_to_remove = sorted_indices_to_remove.scatter(1, x_flat_indices, sorted_indices_to_remove)
+ x_flat = x_flat.masked_fill(indices_to_remove, -float("Inf"))
+ x_flat = gumbel_sample(x_flat, temperature=temp)
+ x = x_flat.view(x.size(0), *x.shape[2:])
+ if mask is not None:
+ x = x * mask + (1-mask) * init_x
+ if i < renoise_steps:
+ if renoise_mode == 'start':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=init_x)
+ elif renoise_mode == 'prev':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=prev_x)
+ else: # 'rand'
+ x, _ = model.add_noise(x, r_range[i+1])
+ preds.append(x.detach())
+ return preds
+
+ return (sample,)
+
+
+@app.cell
+def _(
+ DenoiseUNet,
+ device,
+ get_vae,
+ open_clip,
+ rearrange,
+ text_model_path,
+ torch,
+):
+ vqmodel = get_vae().to(device)
+ vqmodel.eval().requires_grad_(False)
+
+ clip_model, _, _ = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s12b_b42k')
+ clip_model = clip_model.to(device).eval().requires_grad_(False)
+
+
+ def encode(x):
+ return vqmodel.model.encode((2 * x - 1))[-1][-1]
+
+ def decode(img_seq, shape=(32,32)):
+ img_seq = img_seq.view(img_seq.shape[0], -1)
+ b, n = img_seq.shape
+ one_hot_indices = torch.nn.functional.one_hot(img_seq, num_classes=vqmodel.num_tokens).float()
+ z = (one_hot_indices @ vqmodel.model.quantize.embed.weight)
+ z = rearrange(z, 'b (h w) c -> b c h w', h=shape[0], w=shape[1])
+ img = vqmodel.model.decode(z)
+ img = (img.clamp(-1., 1.) + 1) * 0.5
+ return img
+
+ state_dict = torch.load(text_model_path, map_location=device)
+ model = DenoiseUNet(num_labels=8192).to(device)
+ model.load_state_dict(state_dict)
+ model.eval().requires_grad_()
+ print()
+ return clip_model, decode, model
+
+
+@app.cell
+def _(
+ clip_model,
+ config,
+ decode,
+ device,
+ model,
+ sample,
+ time,
+ tokenizer,
+ torch,
+ wandb,
+):
+ text = tokenizer.tokenize([config.prompt_1, config.prompt_2] * config.batch_size).to(device)
+
+ with torch.inference_mode():
+ with torch.autocast(device_type="cuda"):
+ clip_embeddings = clip_model.encode_text(text).float()
+ clip_embeddings = clip_embeddings[:, :, None, None]
+ clip_embeddings = clip_embeddings.expand(
+ -1, -1, config.latent_shape[0], config.latent_shape[1]
+ )
+
+ if config.orientation_mode == 'vertical':
+ interp_mask = torch.linspace(0, 1, config.latent_shape[0], device=device)
+ interp_mask = interp_mask[None, None, :, None]
+ interp_mask = interp_mask.expand(config.batch_size, 1, -1, config.latent_shape[1])
+ else:
+ interp_mask = torch.linspace(0, 1, config.latent_shape[1], device=device)
+ interp_mask = interp_mask[None, None, None, :]
+ interp_mask = interp_mask.expand(config.batch_size, 1, config.latent_shape[0], -1)
+
+ if config.interpolation_mode == "lerp":
+ clip_embeddings = clip_embeddings[0::2] * (1 - interp_mask) + clip_embeddings[1::2] * interp_mask
+ elif config.interpolation_mode == "spherical-lerp":
+ low, high = clip_embeddings[0::2], clip_embeddings[1::2]
+ low_norm = low / torch.norm(low, dim=1, keepdim=True)
+ high_norm = high / torch.norm(high, dim=1, keepdim=True)
+ omega = torch.acos((low_norm * high_norm).sum(1)).unsqueeze(1)
+ so = torch.sin(omega)
+ clip_embeddings = (torch.sin((1.0 - interp_mask) * omega) / so) * low
+ clip_embeddings = clip_embeddings + (torch.sin(interp_mask * omega) / so) * high
+
+ s = time.time()
+ sampled = sample(
+ model, clip_embeddings, T=12, size=config.latent_shape, starting_t=0,
+ temp_range=[1.0, 1.0], typical_filtering=True, typical_mass=0.2,
+ typical_min_tokens=1, classifier_free_scale=5, renoise_steps=11, renoise_mode="start"
+ )
+ wandb.log({"Sampling-Time": time.time() - s})
+ sampled = decode(sampled[-1], config.latent_shape).permute(0, 2, 3, 1)
+ return (sampled,)
+
+
+@app.cell
+def _(log_orientation_guided_multi_conditioning_results, sampled, wandb):
+ log_orientation_guided_multi_conditioning_results(sampled)
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paella-outpainting/paella_outpainting.py b/marimo/convert/paella-outpainting/paella_outpainting.py
new file mode 100644
index 00000000..c325d256
--- /dev/null
+++ b/marimo/convert/paella-outpainting/paella_outpainting.py
@@ -0,0 +1,334 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Image Outpainting with Paella + WandB Playground 🪄🐝
+
+
+
+ A demo of Image Outpainting using [Paella](https://github.com/dome272/Paella) and [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import time
+ import wandb
+ import requests
+ import numpy as np
+ from PIL import Image
+ from io import BytesIO
+ from tqdm.notebook import tqdm
+ import matplotlib.pyplot as plt
+
+ import torch
+ from torch import nn
+ import torchvision
+
+ import open_clip
+ from rudalle import get_vae
+ from einops import rearrange
+ from open_clip import tokenizer
+
+ from Paella.modules import DenoiseUNet
+
+ return (
+ BytesIO,
+ DenoiseUNet,
+ Image,
+ get_vae,
+ np,
+ open_clip,
+ os,
+ rearrange,
+ requests,
+ time,
+ tokenizer,
+ torch,
+ torchvision,
+ wandb,
+ )
+
+
+@app.cell
+def _(os, torch, wandb):
+ wandb_project = "paella" #@param {"type": "string"}
+ wandb_entity = "geekyrakshit" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, entity=wandb_entity, job_type="outpainting")
+
+
+ config = wandb.config
+ config.model_artifact = "geekyrakshit/paella/text-model:v0"
+ config.seed = 42
+ config.batch_size = 3
+ config.latent_shape = (32, 32)
+ config.image_url = "https://media.istockphoto.com/id/1193591781/photo/obedient-dog-breed-welsh-corgi-pembroke-sitting-and-smiles-on-a-white-background-not-isolate.jpg?s=612x612&w=0&k=20&c=ZDKTgSFQFG9QvuDziGsnt55kvQoqJtIhrmVRkpYqxtQ="
+ config.prompt = "black & white photograph of a rocket from the bottom."
+ config.target_size = 256
+ config.batch_size = 5
+ config.canvas_size = (40, 64)
+ config.top_left = (0, 16)
+
+ # Seed Everything
+ torch.manual_seed(config.seed)
+ torch.random.manual_seed(config.seed)
+ torch.cuda.manual_seed(config.seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = True
+
+ # Download Model from Weights & Biases Artifacts
+ text_model_path = os.path.join(wandb.use_artifact(config.model_artifact, type='model').download(), "model_600000.pt")
+ return config, text_model_path
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print("Using device:", device)
+ return (device,)
+
+
+@app.cell
+def _(config, np, wandb):
+ def log_inoutpainting_results(input_image, generated_images, encoded_tokens, mask):
+ generated_images = [
+ wandb.Image(image)
+ for image in (generated_images.cpu().numpy() * 255.0).astype(np.uint8)
+ ]
+ encoded_tokens = [wandb.Image(
+ token,
+ masks={
+ "mask": {
+ "mask_data": mask,
+ "class_labels": {
+ 0: "non-masked",
+ 1: "region-of-interest"
+ }
+ }
+ }
+ ) for token in encoded_tokens]
+ table = wandb.Table(
+ columns=["Seed", "URL", "Input-Image", "Prompt", "Job-Type", "Encoded-Tokens", "Generated-Image"]
+ )
+ table.add_data(
+ config.seed,
+ config.image_url,
+ wandb.Image(input_image),
+ config.prompt,
+ wandb.run.job_type,
+ encoded_tokens,
+ generated_images
+ )
+ wandb.log({"Inpainting-Outpainting-Results": table})
+
+ return (log_inoutpainting_results,)
+
+
+@app.cell
+def _(torch):
+ def log(t, eps=1e-20):
+ return torch.log(t + eps)
+
+ def gumbel_noise(t):
+ noise = torch.zeros_like(t).uniform_(0, 1)
+ return -log(-log(noise))
+
+ def gumbel_sample(t, temperature=1., dim=-1):
+ return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+ def sample(
+ model, c, x=None, mask=None, T=12, size=(32, 32),
+ starting_t=0, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=-1,
+ renoise_steps=11, renoise_mode='start'
+ ):
+ with torch.inference_mode():
+ r_range = torch.linspace(0, 1, T+1)[:-1][:, None].expand(-1, c.size(0)).to(c.device)
+ temperatures = torch.linspace(temp_range[0], temp_range[1], T)
+ preds = []
+ if x is None:
+ x = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ elif mask is not None:
+ noise = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ x = noise * mask + (1-mask) * x
+ init_x = x.clone()
+ for i in range(starting_t, T):
+ if renoise_mode == 'prev':
+ prev_x = x.clone()
+ r, temp = r_range[i], temperatures[i]
+ logits = model(x, c, r)
+ if classifier_free_scale >= 0:
+ logits_uncond = model(x, torch.zeros_like(c), r)
+ logits = torch.lerp(logits_uncond, logits, classifier_free_scale)
+ x = logits
+ x_flat = x.permute(0, 2, 3, 1).reshape(-1, x.size(1))
+ if typical_filtering:
+ x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
+ x_flat_norm_p = torch.exp(x_flat_norm)
+ entropy = -(x_flat_norm * x_flat_norm_p).nansum(-1, keepdim=True)
+
+ c_flat_shifted = torch.abs((-x_flat_norm) - entropy)
+ c_flat_sorted, x_flat_indices = torch.sort(c_flat_shifted, descending=False)
+ x_flat_cumsum = x_flat.gather(-1, x_flat_indices).softmax(dim=-1).cumsum(dim=-1)
+
+ last_ind = (x_flat_cumsum < typical_mass).sum(dim=-1)
+ sorted_indices_to_remove = c_flat_sorted > c_flat_sorted.gather(1, last_ind.view(-1, 1))
+ if typical_min_tokens > 1:
+ sorted_indices_to_remove[..., :typical_min_tokens] = 0
+ indices_to_remove = sorted_indices_to_remove.scatter(1, x_flat_indices, sorted_indices_to_remove)
+ x_flat = x_flat.masked_fill(indices_to_remove, -float("Inf"))
+ x_flat = gumbel_sample(x_flat, temperature=temp)
+ x = x_flat.view(x.size(0), *x.shape[2:])
+ if mask is not None:
+ x = x * mask + (1-mask) * init_x
+ if i < renoise_steps:
+ if renoise_mode == 'start':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=init_x)
+ elif renoise_mode == 'prev':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=prev_x)
+ else: # 'rand'
+ x, _ = model.add_noise(x, r_range[i+1])
+ preds.append(x.detach())
+ return preds
+
+ return (sample,)
+
+
+@app.cell
+def _(
+ DenoiseUNet,
+ device,
+ get_vae,
+ open_clip,
+ rearrange,
+ text_model_path,
+ torch,
+):
+ vqmodel = get_vae().to(device)
+ vqmodel.eval().requires_grad_(False)
+
+ clip_model, _, _ = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s12b_b42k')
+ clip_model = clip_model.to(device).eval().requires_grad_(False)
+
+
+ def encode(x):
+ return vqmodel.model.encode((2 * x - 1))[-1][-1]
+
+ def decode(img_seq, shape=(32,32)):
+ img_seq = img_seq.view(img_seq.shape[0], -1)
+ b, n = img_seq.shape
+ one_hot_indices = torch.nn.functional.one_hot(img_seq, num_classes=vqmodel.num_tokens).float()
+ z = (one_hot_indices @ vqmodel.model.quantize.embed.weight)
+ z = rearrange(z, 'b (h w) c -> b c h w', h=shape[0], w=shape[1])
+ img = vqmodel.model.decode(z)
+ img = (img.clamp(-1., 1.) + 1) * 0.5
+ return img
+
+ state_dict = torch.load(text_model_path, map_location=device)
+ model = DenoiseUNet(num_labels=8192).to(device)
+ model.load_state_dict(state_dict)
+ model.eval().requires_grad_()
+ print()
+ return clip_model, decode, encode, model
+
+
+@app.cell
+def _(BytesIO, Image, config, device, requests, torchvision):
+ response = requests.get(config.image_url)
+ original_image = Image.open(BytesIO(response.content)).convert("RGB")
+
+ preprocess = torchvision.transforms.Compose([
+ torchvision.transforms.Resize(config.target_size),
+ torchvision.transforms.ToTensor(),
+ ])
+
+ images = preprocess(original_image).unsqueeze(0).expand(config.batch_size, -1, -1, -1).to(device)[:, :3]
+ return images, original_image
+
+
+@app.cell
+def _(
+ clip_model,
+ config,
+ decode,
+ device,
+ encode,
+ images,
+ model,
+ sample,
+ time,
+ tokenizer,
+ torch,
+ wandb,
+):
+ tokenized_text_and_image = tokenizer.tokenize([config.prompt] * images.shape[0]).to(device)
+ with torch.inference_mode():
+ with torch.autocast(device_type="cuda"):
+ clip_embeddings = clip_model.encode_text(tokenized_text_and_image).float()
+ encoded_tokens = encode(images)
+ canvas = torch.zeros((images.shape[0], *config.canvas_size), dtype=torch.long).to(device)
+ canvas[
+ :,
+ config.top_left[0]: config.top_left[0] + encoded_tokens.shape[1],
+ config.top_left[1]: config.top_left[1] + encoded_tokens.shape[2]
+ ] = encoded_tokens
+ mask = torch.ones_like(canvas)
+ mask[
+ :,
+ config.top_left[0]: config.top_left[0] + encoded_tokens.shape[1],
+ config.top_left[1]: config.top_left[1] + encoded_tokens.shape[2]
+ ] = 0
+ s = time.time()
+ sampled = sample(
+ model, clip_embeddings, x=canvas, mask=mask, T=12,
+ size=config.canvas_size, starting_t=0, temp_range=[1.0, 1.0],
+ typical_filtering=True, typical_mass=0.2, typical_min_tokens=1,
+ classifier_free_scale=4, renoise_steps=11
+ )
+ wandb.log({"Sampling-Time": time.time() - s})
+ sampled = decode(sampled[-1], config.canvas_size).permute(0, 2, 3, 1)
+ mask = mask[0:1].cpu().numpy()[0]
+ encoded_tokens = encoded_tokens.cpu().numpy()
+ return encoded_tokens, mask, sampled
+
+
+@app.cell
+def _(
+ encoded_tokens,
+ log_inoutpainting_results,
+ mask,
+ original_image,
+ sampled,
+ wandb,
+):
+ log_inoutpainting_results(original_image, sampled, encoded_tokens, mask)
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paella-structural-morphing/paella_structural_morphing.py b/marimo/convert/paella-structural-morphing/paella_structural_morphing.py
new file mode 100644
index 00000000..2814a8de
--- /dev/null
+++ b/marimo/convert/paella-structural-morphing/paella_structural_morphing.py
@@ -0,0 +1,304 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Structural Morphing of Images with Paella + WandB Playground 🪄🐝
+
+
+
+ A demo of Structural Morphing of Images using [Paella](https://github.com/dome272/Paella) and [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import time
+ import wandb
+ import requests
+ import numpy as np
+ from PIL import Image
+ from io import BytesIO
+ from tqdm.notebook import tqdm
+ import matplotlib.pyplot as plt
+
+ import torch
+ from torch import nn
+ import torchvision
+
+ import open_clip
+ from rudalle import get_vae
+ from einops import rearrange
+ from open_clip import tokenizer
+
+ from Paella.modules import DenoiseUNet
+
+ return (
+ BytesIO,
+ DenoiseUNet,
+ Image,
+ get_vae,
+ np,
+ open_clip,
+ os,
+ rearrange,
+ requests,
+ time,
+ tokenizer,
+ torch,
+ torchvision,
+ wandb,
+ )
+
+
+@app.cell
+def _(os, torch, wandb):
+ wandb_project = "paella" #@param {"type": "string"}
+ wandb_entity = "geekyrakshit" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, entity=wandb_entity, job_type="structural-morphing")
+
+
+ config = wandb.config
+ config.model_artifact = "geekyrakshit/paella/text-model:v0"
+ config.seed = 42
+ config.batch_size = 5
+ config.latent_shape = (32, 32)
+ config.image_url = "https://media.istockphoto.com/id/1193591781/photo/obedient-dog-breed-welsh-corgi-pembroke-sitting-and-smiles-on-a-white-background-not-isolate.jpg?s=612x612&w=0&k=20&c=ZDKTgSFQFG9QvuDziGsnt55kvQoqJtIhrmVRkpYqxtQ="
+ config.prompt = "pink dog"
+ config.target_size = 256
+ config.initial_step = 8
+ config.max_step = 24
+
+ # Seed Everything
+ torch.manual_seed(config.seed)
+ torch.random.manual_seed(config.seed)
+ torch.cuda.manual_seed(config.seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = True
+
+ # Download Model from Weights & Biases Artifacts
+ text_model_path = os.path.join(wandb.use_artifact(config.model_artifact, type='model').download(), "model_600000.pt")
+ return config, text_model_path
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print("Using device:", device)
+ return (device,)
+
+
+@app.cell
+def _(config, np, wandb):
+ def log_structural_editing_results(input_image, generated_images):
+ generated_images = [
+ wandb.Image(image)
+ for image in (generated_images.cpu().numpy() * 255.0).astype(np.uint8)
+ ]
+ table = wandb.Table(
+ columns=["Seed", "URL", "Input-Image", "Prompt", "Generated-Image"]
+ )
+ table.add_data(
+ config.seed,
+ config.image_url,
+ wandb.Image(input_image),
+ config.prompt,
+ generated_images
+ )
+ wandb.log({"Structural-Morphing-Results": table})
+
+ return (log_structural_editing_results,)
+
+
+@app.cell
+def _(torch):
+ def log(t, eps=1e-20):
+ return torch.log(t + eps)
+
+ def gumbel_noise(t):
+ noise = torch.zeros_like(t).uniform_(0, 1)
+ return -log(-log(noise))
+
+ def gumbel_sample(t, temperature=1., dim=-1):
+ return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+ def sample(
+ model, c, x=None, mask=None, T=12, size=(32, 32),
+ starting_t=0, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=-1,
+ renoise_steps=11, renoise_mode='start'
+ ):
+ with torch.inference_mode():
+ r_range = torch.linspace(0, 1, T+1)[:-1][:, None].expand(-1, c.size(0)).to(c.device)
+ temperatures = torch.linspace(temp_range[0], temp_range[1], T)
+ preds = []
+ if x is None:
+ x = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ elif mask is not None:
+ noise = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ x = noise * mask + (1-mask) * x
+ init_x = x.clone()
+ for i in range(starting_t, T):
+ if renoise_mode == 'prev':
+ prev_x = x.clone()
+ r, temp = r_range[i], temperatures[i]
+ logits = model(x, c, r)
+ if classifier_free_scale >= 0:
+ logits_uncond = model(x, torch.zeros_like(c), r)
+ logits = torch.lerp(logits_uncond, logits, classifier_free_scale)
+ x = logits
+ x_flat = x.permute(0, 2, 3, 1).reshape(-1, x.size(1))
+ if typical_filtering:
+ x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
+ x_flat_norm_p = torch.exp(x_flat_norm)
+ entropy = -(x_flat_norm * x_flat_norm_p).nansum(-1, keepdim=True)
+
+ c_flat_shifted = torch.abs((-x_flat_norm) - entropy)
+ c_flat_sorted, x_flat_indices = torch.sort(c_flat_shifted, descending=False)
+ x_flat_cumsum = x_flat.gather(-1, x_flat_indices).softmax(dim=-1).cumsum(dim=-1)
+
+ last_ind = (x_flat_cumsum < typical_mass).sum(dim=-1)
+ sorted_indices_to_remove = c_flat_sorted > c_flat_sorted.gather(1, last_ind.view(-1, 1))
+ if typical_min_tokens > 1:
+ sorted_indices_to_remove[..., :typical_min_tokens] = 0
+ indices_to_remove = sorted_indices_to_remove.scatter(1, x_flat_indices, sorted_indices_to_remove)
+ x_flat = x_flat.masked_fill(indices_to_remove, -float("Inf"))
+ x_flat = gumbel_sample(x_flat, temperature=temp)
+ x = x_flat.view(x.size(0), *x.shape[2:])
+ if mask is not None:
+ x = x * mask + (1-mask) * init_x
+ if i < renoise_steps:
+ if renoise_mode == 'start':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=init_x)
+ elif renoise_mode == 'prev':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=prev_x)
+ else: # 'rand'
+ x, _ = model.add_noise(x, r_range[i+1])
+ preds.append(x.detach())
+ return preds
+
+ return (sample,)
+
+
+@app.cell
+def _(
+ DenoiseUNet,
+ device,
+ get_vae,
+ open_clip,
+ rearrange,
+ text_model_path,
+ torch,
+):
+ vqmodel = get_vae().to(device)
+ vqmodel.eval().requires_grad_(False)
+
+ clip_model, _, _ = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s12b_b42k')
+ clip_model = clip_model.to(device).eval().requires_grad_(False)
+
+
+ def encode(x):
+ return vqmodel.model.encode((2 * x - 1))[-1][-1]
+
+ def decode(img_seq, shape=(32,32)):
+ img_seq = img_seq.view(img_seq.shape[0], -1)
+ b, n = img_seq.shape
+ one_hot_indices = torch.nn.functional.one_hot(img_seq, num_classes=vqmodel.num_tokens).float()
+ z = (one_hot_indices @ vqmodel.model.quantize.embed.weight)
+ z = rearrange(z, 'b (h w) c -> b c h w', h=shape[0], w=shape[1])
+ img = vqmodel.model.decode(z)
+ img = (img.clamp(-1., 1.) + 1) * 0.5
+ return img
+
+ state_dict = torch.load(text_model_path, map_location=device)
+ model = DenoiseUNet(num_labels=8192).to(device)
+ model.load_state_dict(state_dict)
+ model.eval().requires_grad_()
+ print()
+ return clip_model, decode, encode, model
+
+
+@app.cell
+def _(BytesIO, Image, config, device, requests, torchvision):
+ response = requests.get(config.image_url)
+ original_image = Image.open(BytesIO(response.content)).convert("RGB")
+
+ preprocess = torchvision.transforms.Compose([
+ torchvision.transforms.Resize(config.target_size),
+ torchvision.transforms.ToTensor(),
+ ])
+
+ images = preprocess(original_image).unsqueeze(0).expand(config.batch_size, -1, -1, -1).to(device)[:, :3]
+ return images, original_image
+
+
+@app.cell
+def _(
+ clip_model,
+ config,
+ decode,
+ device,
+ encode,
+ images,
+ model,
+ sample,
+ time,
+ tokenizer,
+ torch,
+ wandb,
+):
+ with torch.inference_mode():
+ with torch.autocast(device_type="cuda"):
+ latent_image = encode(images)
+ latent_shape = latent_image.shape[-2:]
+ r = torch.ones(latent_image.size(0), device=device) * (config.initial_step / config.max_step)
+ noised_latent_image, _ = model.add_noise(latent_image, r)
+
+ tokenized_text = tokenizer.tokenize([config.prompt] * images.size(0)).to(device)
+ clip_embeddings = clip_model.encode_text(tokenized_text).float()
+
+ s = time.time()
+ sampled = sample(
+ model, clip_embeddings, x=noised_latent_image, T=config.max_step, size=latent_shape,
+ starting_t=config.initial_step, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=6,
+ renoise_steps=config.max_step - 1, renoise_mode="prev"
+ )
+ wandb.log({"Sampling-Time": time.time() - s})
+
+ sampled = decode(sampled[-1], latent_shape).permute(0, 2, 3, 1)
+ return (sampled,)
+
+
+@app.cell
+def _(log_structural_editing_results, original_image, sampled, wandb):
+ log_structural_editing_results(original_image, sampled)
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/paella-text-conditional/paella_text_conditional.py b/marimo/convert/paella-text-conditional/paella_text_conditional.py
new file mode 100644
index 00000000..8789caef
--- /dev/null
+++ b/marimo/convert/paella-text-conditional/paella_text_conditional.py
@@ -0,0 +1,262 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Text-Conditional Image Generation with Paella + WandB Playground 🪄🐝
+
+
+
+ A demo of Text-Conditional Image Generation using [Paella](https://github.com/dome272/Paella) and [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import time
+ import wandb
+ import requests
+ import numpy as np
+ from tqdm import tqdm
+ from PIL import Image
+ from io import BytesIO
+ import matplotlib.pyplot as plt
+
+ import torch
+ from torch import nn
+ import torchvision
+
+ import open_clip
+ from rudalle import get_vae
+ from einops import rearrange
+ from open_clip import tokenizer
+
+ from Paella.modules import DenoiseUNet
+
+ return (
+ DenoiseUNet,
+ get_vae,
+ np,
+ open_clip,
+ os,
+ rearrange,
+ time,
+ tokenizer,
+ torch,
+ wandb,
+ )
+
+
+@app.cell
+def _(os, torch, wandb):
+ wandb_project = "paella" #@param {"type": "string"}
+ wandb_entity = "geekyrakshit" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, entity=wandb_entity, job_type="text-conditional")
+
+
+ config = wandb.config
+ config.model_artifact = "geekyrakshit/paella/text-model:v0"
+ config.seed = 440
+ config.batch_size = 5
+ config.latent_shape = (32, 32)
+ config.prompt = "a beautiful painting of chernobyl by nekro and pascal blanche and syd mead and greg rutkowski and sin jong hun and moebius and simon stalenhag. in style of cg art. ray tracing. cel shading. hyper detailed. realistic. ue 5. maya. octane render."
+
+ # Seed Everything
+ torch.manual_seed(config.seed)
+ torch.cuda.manual_seed(config.seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = True
+
+ # Download Model from Weights & Biases Artifacts
+ text_model_path = os.path.join(wandb.use_artifact(config.model_artifact, type='model').download(), "model_600000.pt")
+ return config, text_model_path
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print("Using device:", device)
+ return (device,)
+
+
+@app.cell
+def _(config, np, wandb):
+ def log_text_conditional_results(images, latent_shape):
+ images = [wandb.Image(image) for image in (images.cpu().numpy() * 255.0).astype(np.uint8)]
+ table = wandb.Table(columns=["Seed", "Latent-Shape", "Prompt", "Generated-Images"])
+ table.add_data(config.seed, latent_shape, config.prompt, images)
+ wandb.log({"Text-Conditional-Results": table})
+
+ return (log_text_conditional_results,)
+
+
+@app.cell
+def _(torch):
+ def log(t, eps=1e-20):
+ return torch.log(t + eps)
+
+ def gumbel_noise(t):
+ noise = torch.zeros_like(t).uniform_(0, 1)
+ return -log(-log(noise))
+
+ def gumbel_sample(t, temperature=1., dim=-1):
+ return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+ def sample(
+ model, c, x=None, mask=None, T=12, size=(32, 32),
+ starting_t=0, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=-1,
+ renoise_steps=11, renoise_mode='start'
+ ):
+ with torch.inference_mode():
+ r_range = torch.linspace(0, 1, T+1)[:-1][:, None].expand(-1, c.size(0)).to(c.device)
+ temperatures = torch.linspace(temp_range[0], temp_range[1], T)
+ preds = []
+ if x is None:
+ x = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ elif mask is not None:
+ noise = torch.randint(0, model.num_labels, size=(c.size(0), *size), device=c.device)
+ x = noise * mask + (1-mask) * x
+ init_x = x.clone()
+ for i in range(starting_t, T):
+ if renoise_mode == 'prev':
+ prev_x = x.clone()
+ r, temp = r_range[i], temperatures[i]
+ logits = model(x, c, r)
+ if classifier_free_scale >= 0:
+ logits_uncond = model(x, torch.zeros_like(c), r)
+ logits = torch.lerp(logits_uncond, logits, classifier_free_scale)
+ x = logits
+ x_flat = x.permute(0, 2, 3, 1).reshape(-1, x.size(1))
+ if typical_filtering:
+ x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
+ x_flat_norm_p = torch.exp(x_flat_norm)
+ entropy = -(x_flat_norm * x_flat_norm_p).nansum(-1, keepdim=True)
+
+ c_flat_shifted = torch.abs((-x_flat_norm) - entropy)
+ c_flat_sorted, x_flat_indices = torch.sort(c_flat_shifted, descending=False)
+ x_flat_cumsum = x_flat.gather(-1, x_flat_indices).softmax(dim=-1).cumsum(dim=-1)
+
+ last_ind = (x_flat_cumsum < typical_mass).sum(dim=-1)
+ sorted_indices_to_remove = c_flat_sorted > c_flat_sorted.gather(1, last_ind.view(-1, 1))
+ if typical_min_tokens > 1:
+ sorted_indices_to_remove[..., :typical_min_tokens] = 0
+ indices_to_remove = sorted_indices_to_remove.scatter(1, x_flat_indices, sorted_indices_to_remove)
+ x_flat = x_flat.masked_fill(indices_to_remove, -float("Inf"))
+ # x_flat = torch.multinomial(x_flat.div(temp).softmax(-1), num_samples=1)[:, 0]
+ x_flat = gumbel_sample(x_flat, temperature=temp)
+ x = x_flat.view(x.size(0), *x.shape[2:])
+ if mask is not None:
+ x = x * mask + (1-mask) * init_x
+ if i < renoise_steps:
+ if renoise_mode == 'start':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=init_x)
+ elif renoise_mode == 'prev':
+ x, _ = model.add_noise(x, r_range[i+1], random_x=prev_x)
+ else: # 'rand'
+ x, _ = model.add_noise(x, r_range[i+1])
+ preds.append(x.detach())
+ return preds
+
+ return (sample,)
+
+
+@app.cell
+def _(
+ DenoiseUNet,
+ device,
+ get_vae,
+ open_clip,
+ rearrange,
+ text_model_path,
+ torch,
+):
+ vqmodel = get_vae().to(device)
+ vqmodel.eval().requires_grad_(False)
+
+ clip_model, _, _ = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s12b_b42k')
+ clip_model = clip_model.to(device).eval().requires_grad_(False)
+
+
+ def encode(x):
+ return vqmodel.model.encode((2 * x - 1))[-1][-1]
+
+ def decode(img_seq, shape=(32,32)):
+ img_seq = img_seq.view(img_seq.shape[0], -1)
+ b, n = img_seq.shape
+ one_hot_indices = torch.nn.functional.one_hot(img_seq, num_classes=vqmodel.num_tokens).float()
+ z = (one_hot_indices @ vqmodel.model.quantize.embed.weight)
+ z = rearrange(z, 'b (h w) c -> b c h w', h=shape[0], w=shape[1])
+ img = vqmodel.model.decode(z)
+ img = (img.clamp(-1., 1.) + 1) * 0.5
+ return img
+
+ state_dict = torch.load(text_model_path, map_location=device)
+ model = DenoiseUNet(num_labels=8192).to(device)
+ model.load_state_dict(state_dict)
+ model.eval().requires_grad_()
+ print()
+ return clip_model, decode, model
+
+
+@app.cell
+def _(
+ clip_model,
+ config,
+ decode,
+ device,
+ model,
+ sample,
+ time,
+ tokenizer,
+ torch,
+ wandb,
+):
+ tokenized_text = tokenizer.tokenize([config.prompt] * config.batch_size).to(device)
+ with torch.inference_mode():
+ with torch.autocast(device_type="cuda"):
+ clip_embeddings = clip_model.encode_text(tokenized_text)
+ s = time.time()
+ sampled = sample(
+ model, clip_embeddings, T=12, size=config.latent_shape,
+ starting_t=0, temp_range=[1.0, 1.0], typical_filtering=True,
+ typical_mass=0.2, typical_min_tokens=1, classifier_free_scale=5,
+ renoise_steps=11, renoise_mode="start"
+ )
+ wandb.log({"Sampling-Time": time.time() - s})
+ sampled = decode(sampled[-1], config.latent_shape).permute(0, 2, 3, 1)
+ return (sampled,)
+
+
+@app.cell
+def _(config, log_text_conditional_results, sampled, wandb):
+ log_text_conditional_results(sampled, config.latent_shape)
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/peft-llama-token-cls/peft_llama_token_cls.py b/marimo/convert/peft-llama-token-cls/peft_llama_token_cls.py
new file mode 100644
index 00000000..dcfc6044
--- /dev/null
+++ b/marimo/convert/peft-llama-token-cls/peft_llama_token_cls.py
@@ -0,0 +1,427 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📦 Packages and Basic Setup
+ ---
+
+ To run the notebooks you'll need two secrets named `W&B` and `HF_TOKEN`. Also, in the configuration section change the `wandb_entity` to your username/workspace.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install -q -U bitsandbytes datasets evaluate ml-collections seqeval wandb
+ # !pip install -q git+https://github.com/huggingface/peft.git
+ return
+
+
+@app.cell
+def _():
+ import evaluate
+ import numpy as np
+ from transformers import AutoTokenizer
+ from datasets import ClassLabel, load_dataset
+ from transformers import TrainingArguments, Trainer
+ from peft import get_peft_model, LoraConfig, TaskType
+ from transformers import DataCollatorForTokenClassification
+
+ return (
+ DataCollatorForTokenClassification,
+ LoraConfig,
+ TaskType,
+ Trainer,
+ TrainingArguments,
+ get_peft_model,
+ np,
+ )
+
+
+@app.cell
+def _():
+ import wandb
+ wandb.login()
+ return (wandb,)
+
+
+@app.cell
+def _():
+ # @title ⚙️ Configuration
+
+ import ml_collections
+
+ def get_config() -> ml_collections.ConfigDict:
+ config = ml_collections.ConfigDict()
+ config.model = "unsloth/llama-2-7b-bnb-4bit" # @param {type: "string"}
+ config.lora_r = 4 # @param {type: "number"}
+ config.lora_alpha = 32 # @param {type: "number"}
+ config.lora_dropout = 0.1 # @param {type: "number"}
+ config.max_length = 32 # @param {type: "number"}
+ config.batch_size = 16 # @param {type: "number"}
+ config.num_epochs = 5 # @param {type: "number"}
+ config.learning_rate = 1e-3 # @param {type: "number"}
+ config.dataset = "conll2003" # @param {type: "string"}
+ config.wandb_entity = None # @param {type: "string"}
+ return config
+
+ config = get_config()
+ return (config,)
+
+
+@app.cell
+def _(config, wandb):
+ import os
+ wandb.init(project='Llama-NER', job_type='train', group=config.model, config=config.to_dict(), entity=config.wandb_entity)
+ os.environ['WANDB_WATCH'] = 'false'
+ os.environ['WANDB_LOG_MODEL'] = 'false'
+ os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💿 The Dataset
+ ---
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # ds = load_dataset(
+ # config.dataset,
+ # cache_dir="/cache/",
+ # )
+ #
+ # seqeval = evaluate.load("seqeval")
+ return
+
+
+@app.cell
+def _(ds):
+ column_names = ds["train"].column_names
+ features = ds["train"].features
+
+ text_column_name = "tokens"
+ label_column_name = "ner_tags"
+
+ label_list = features[label_column_name].feature.names
+ label2id = {i: i for i in range(len(label_list))}
+ id2label = {v: k for k, v in label2id.items()}
+ return (label_list,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🖖 Utility Functions
+ ---
+ """)
+ return
+
+
+@app.cell
+def _(label_list, np, seqeval):
+ def compute_metrics(p):
+ predictions, labels = p
+ predictions = np.argmax(predictions, axis=2)
+
+ true_predictions = [
+ [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
+ for prediction, label in zip(predictions, labels)
+ ]
+ true_labels = [
+ [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
+ for prediction, label in zip(predictions, labels)
+ ]
+
+ results = seqeval.compute(predictions=true_predictions, references=true_labels)
+ return {
+ "precision": results["overall_precision"],
+ "recall": results["overall_recall"],
+ "f1": results["overall_f1"],
+ "accuracy": results["overall_accuracy"],
+ }
+
+ return (compute_metrics,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🏠 Model Architecture
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Implementating `LlamaForTokenClassification`
+
+ [Source: @KoichiYasuoka](https://github.com/huggingface/transformers/issues/26521#issuecomment-1868284434)
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # from typing import List, Optional, Tuple, Union
+ # import torch
+ # from torch import nn
+ # from transformers.modeling_outputs import TokenClassifierOutput
+ # from transformers.file_utils import add_start_docstrings_to_model_forward
+ # from transformers.models.llama.modeling_llama import LlamaModel, LlamaPreTrainedModel, LLAMA_INPUTS_DOCSTRING
+ #
+ # class LlamaForTokenClassification(LlamaPreTrainedModel):
+ # def __init__(self, config):
+ # super().__init__(config)
+ # self.num_labels = config.num_labels
+ # self.model = LlamaModel(config)
+ # if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+ # classifier_dropout = config.classifier_dropout
+ # elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
+ # classifier_dropout = config.hidden_dropout
+ # else:
+ # classifier_dropout = 0.1
+ # self.dropout = nn.Dropout(classifier_dropout)
+ # self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+ #
+ # # Initialize weights and apply final processing
+ # self.post_init()
+ #
+ # def get_input_embeddings(self):
+ # return self.model.embed_tokens
+ #
+ # def set_input_embeddings(self, value):
+ # self.model.embed_tokens = value
+ #
+ # @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
+ # def forward(
+ # self,
+ # input_ids: Optional[torch.LongTensor] = None,
+ # attention_mask: Optional[torch.Tensor] = None,
+ # position_ids: Optional[torch.LongTensor] = None,
+ # past_key_values: Optional[List[torch.FloatTensor]] = None,
+ # inputs_embeds: Optional[torch.FloatTensor] = None,
+ # labels: Optional[torch.LongTensor] = None,
+ # use_cache: Optional[bool] = None,
+ # output_attentions: Optional[bool] = None,
+ # output_hidden_states: Optional[bool] = None,
+ # return_dict: Optional[bool] = None,
+ # ) -> Union[Tuple, TokenClassifierOutput]:
+ #
+ # return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+ #
+ # transformer_outputs = self.model(
+ # input_ids,
+ # attention_mask=attention_mask,
+ # position_ids=position_ids,
+ # past_key_values=past_key_values,
+ # inputs_embeds=inputs_embeds,
+ # use_cache=use_cache,
+ # output_attentions=output_attentions,
+ # output_hidden_states=output_hidden_states,
+ # return_dict=return_dict,
+ # )
+ #
+ # hidden_states = transformer_outputs[0]
+ # hidden_states = self.dropout(hidden_states)
+ # logits = self.classifier(hidden_states)
+ #
+ # loss = None
+ # if labels is not None:
+ # labels = labels.to(logits.device)
+ # loss_fct = nn.CrossEntropyLoss()
+ # loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ #
+ # if not return_dict:
+ # output = (logits,) + transformer_outputs[2:]
+ # return ((loss,) + output) if loss is not None else output
+ #
+ # return TokenClassifierOutput(
+ # loss=loss,
+ # logits=logits,
+ # hidden_states=transformer_outputs.hidden_states,
+ # attentions=transformer_outputs.attentions
+ # )
+ #
+ # tokenizer = AutoTokenizer.from_pretrained(config.model)
+ #
+ # model = LlamaForTokenClassification.from_pretrained(
+ # config.model,
+ # num_labels=len(label_list),
+ # id2label=id2label,
+ # label2id=label2id,
+ # cache_dir="/cache/",
+ # )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Process Dataset for Token Classification
+ """)
+ return
+
+
+@app.cell
+def _(DataCollatorForTokenClassification, config, ds, tokenizer):
+ def tokenize_and_align_labels(examples):
+ tokenized_inputs = tokenizer(examples["tokens"], is_split_into_words=True, padding='longest', max_length=config.max_length, truncation=True)
+
+ labels = []
+ for i, label in enumerate(examples[f"ner_tags"]):
+ word_ids = tokenized_inputs.word_ids(batch_index=i) # Map tokens to their respective word.
+ previous_word_idx = None
+ label_ids = []
+ for word_idx in word_ids: # Set the special tokens to -100.
+ if word_idx is None:
+ label_ids.append(-100)
+ elif word_idx != previous_word_idx: # Only label the first token of a given word.
+ label_ids.append(label[word_idx])
+ else:
+ label_ids.append(-100)
+ previous_word_idx = word_idx
+ labels.append(label_ids)
+
+ tokenized_inputs["labels"] = labels
+ return tokenized_inputs
+
+ tokenized_ds = ds.map(tokenize_and_align_labels, batched=True)
+ data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
+ return data_collator, tokenized_ds
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Obtaining 🤗 PEFT Model
+ """)
+ return
+
+
+@app.cell
+def _(LoraConfig, TaskType, config, get_peft_model):
+ peft_config = LoraConfig(
+ task_type=TaskType.TOKEN_CLS,
+ inference_mode=False,
+ r=config.lora_r,
+ lora_alpha=config.lora_alpha,
+ lora_dropout=config.lora_dropout
+ )
+
+ model = get_peft_model(model, peft_config)
+ model.print_trainable_parameters()
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ✍️ Training
+ ---
+ """)
+ return
+
+
+@app.cell
+def _(
+ Trainer,
+ TrainingArguments,
+ compute_metrics,
+ config,
+ data_collator,
+ model,
+ tokenized_ds,
+ tokenizer,
+):
+ training_args = TrainingArguments(
+ output_dir="unsloth-llama-2-7b-bnb-4bit-conll2003",
+ learning_rate=config.learning_rate,
+ gradient_accumulation_steps=2,
+ per_device_train_batch_size=config.batch_size,
+ per_device_eval_batch_size=config.batch_size,
+ num_train_epochs=config.num_epochs,
+ logging_steps=100,
+ weight_decay=0.01,
+ evaluation_strategy="epoch",
+ save_strategy="epoch",
+ report_to=["wandb"],
+ optim="paged_adamw_8bit",
+ load_best_model_at_end=True,
+ push_to_hub=True,
+ )
+
+ trainer = Trainer(
+ model=model,
+ args=training_args,
+ train_dataset=tokenized_ds["train"],
+ eval_dataset=tokenized_ds["test"],
+ tokenizer=tokenizer,
+ data_collator=data_collator,
+ compute_metrics=compute_metrics,
+ )
+
+ train_results = trainer.train()
+ return (train_results,)
+
+
+@app.cell
+def _(train_results, wandb):
+ wandb.config.train_results = train_results
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📚 References
+
+ * Github: [`4AI/LS-LLaMA`](https://github.com/4AI/LS-LLaMA)
+ * [Alpaca + Llama 7b example by `@unslothai`](https://colab.research.google.com/drive/1lBzz5KeZJKXjvivbYvmGarix9Ao6Wxe5?usp=sharing)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/prompts-prompts-evaluation/prompts_prompts_evaluation.py b/marimo/convert/prompts-prompts-evaluation/prompts_prompts_evaluation.py
new file mode 100644
index 00000000..20022c7a
--- /dev/null
+++ b/marimo/convert/prompts-prompts-evaluation/prompts_prompts_evaluation.py
@@ -0,0 +1,414 @@
+# /// script
+# dependencies = ["openai", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Iterate and Evaluate LLM applications
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ AI application building is an experimental process where you likely don't know how a given system will perform on your task. To iterate on an application, we need a way to evaluate if it's improving. To do so, a common practice is to test it against the same dataset when there is a change.
+
+ This tutorial will show you how to:
+ - track input prompts and pipeline settings with `wandb.config`
+ - track final evaluation metrics e.g. F1 score or scores from LLM judges, with `wandb.log`
+ - track individual model predictions and metadata in `W&B Tables`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll track F1 score on extracting named entities from an example news headlines dataset from `explosion/prodigy-recipes` from the https://prodi.gy/ team.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ ## Download Data
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! curl -O https://raw.githubusercontent.com/explosion/prodigy-recipes/master/example-datasets/annotated_news_headlines-ORG-PERSON-LOCATION-ner.jsonl
+ subprocess.call(['curl', '-O', 'https://raw.githubusercontent.com/explosion/prodigy-recipes/master/example-datasets/annotated_news_headlines-ORG-PERSON-LOCATION-ner.jsonl'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installation
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb openai !pip install wandb openai
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Create a W&B account and log in
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ wandb.login()
+ return (wandb,)
+
+
+@app.cell
+def _():
+ import json
+ from functools import partial
+ import timeit
+ import openai
+ from concurrent.futures import ThreadPoolExecutor
+ data = []
+ with open('annotated_news_headlines-ORG-PERSON-LOCATION-ner.jsonl') as f:
+ for line in f:
+ data.append(json.loads(line))
+ return ThreadPoolExecutor, data, partial, timeit
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Format data
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here we just remove data we're not using and format the examples for our task.
+ """)
+ return
+
+
+@app.cell
+def _(data):
+ def clean_examples():
+ labelled_examples = []
+ for example in data:
+ entities = []
+ if 'spans' in example:
+ for span in example['spans']:
+ start = span['start']
+ end = span['end']
+ label = span['label']
+ text = '' # Extract the corresponding text from tokens
+ for token in example['tokens']:
+ if token['start'] >= start and token['end'] <= end:
+ text = text + (token['text'] + ' ')
+ entities.append(text.rstrip())
+ labelled_examples.append({'text': example['text'], 'entities': entities})
+ return labelled_examples
+ labelled_examples = clean_examples()
+ return (labelled_examples,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Set up LLM boilerplate
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll call `openai` (you'll need to add an OpenAI API key) with a given prompt to extract the entities and replace `` with our input. We'll also grab useful metadata from the openai response for logging.
+ """)
+ return
+
+
+@app.cell
+def _(timeit):
+ def extract_entities_with_template(text, template_prompt, system_prompt, model, temperature):
+ start_time = timeit.default_timer()
+ prompt = template_prompt.replace('', text)
+ from openai import OpenAI
+ client = OpenAI()
+ response = client.chat.completions.create(model=model, messages=[{'role': 'system', 'content': system_prompt}, {'role': 'user', 'content': prompt}], temperature=temperature)
+ text = response.choices[0].message.content
+ entities = list(filter(None, text.split('\n')))
+ usage = response.usage
+ prompt_tokens = usage.prompt_tokens
+ completion_tokens = usage.completion_tokens
+ total_tokens = usage.total_tokens
+ end_time = timeit.default_timer()
+ _elapsed = end_time - start_time
+ return {'entities': entities, 'model': model, 'prompt': prompt, 'elapsed': _elapsed, 'prompt_tokens': prompt_tokens, 'completion_tokens': completion_tokens, 'total_tokens': total_tokens}
+
+ return (extract_entities_with_template,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Calculate Metric
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here, we make an evaluation metric for our task.
+ Note: It's not shown here, but you could also use an LLM to evaluate your task if it's not as straight forward to evaluate as this task.
+ """)
+ return
+
+
+@app.function
+def calculate_f1(extracted_entities, ground_truth_entities):
+ extracted_set = set(map(str.lower, extracted_entities))
+ ground_truth_set = set(map(str.lower, ground_truth_entities))
+ tp_examples = extracted_set & ground_truth_set
+ tp = len(tp_examples)
+ fp_examples = extracted_set - ground_truth_set
+ fp = len(fp_examples)
+ fn_examples = ground_truth_set - extracted_set
+ fn = len(fn_examples)
+ precision = tp / (tp + fp) if (tp + fp) else 0
+ recall = tp / (tp + fn) if (tp + fn) else 0
+ f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0
+ return f1, tp, fp, fn
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Perform inference in parallel
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Running evaluations can be a bit slow. To speed it up, here is a bit of useful code to gather your examples in parallel. None of this is specific to W&B, but it's useful to have nonetheless.
+ """)
+ return
+
+
+@app.cell
+def _(
+ ThreadPoolExecutor,
+ extract_entities_with_template,
+ labelled_examples,
+ partial,
+ timeit,
+):
+ def inference(examples, system_prompt, template_prompt, model, temperature):
+ _extracted = []
+ openai_func = partial(extract_entities_with_template, model=model, system_prompt=system_prompt, template_prompt=template_prompt, temperature=temperature) # making a new function to openai which has the template
+ start_time = timeit.default_timer() # this is needed because exectutor.map wants a func with one arg
+ with ThreadPoolExecutor(max_workers=8) as executor:
+ for i in executor.map(openai_func, [t['text'] for t in examples]):
+ _extracted.append(i)
+ end_time = timeit.default_timer() # Run the model to extract the entities
+ _elapsed = end_time - start_time
+ return (_extracted, _elapsed)
+ model = 'gpt-3.5-turbo'
+ temperature = 0.7
+ template = '\ntext: \nReturn the entities as a list with a new line between each entity.\n'
+ system_prompt = 'You are an excellent entity extractor reading newspapers and extracting orgs, people and locations. Extract the entities from the follow sentence.'
+ _extracted, _elapsed = inference(labelled_examples[:1], system_prompt, template, model, temperature)
+ print(_extracted[0])
+ print(labelled_examples[0]['text'])
+ return inference, model, system_prompt, temperature, template
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Evaluate extracted entities, save in W&B Table for inspection later
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here, we calcualte our metric across all of our predictions and log them to a `wandb.Table` for later inspection.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ def evaluate(extracted, labelled_examples):
+ total_tp, total_fp, total_fn = (0, 0, 0)
+ eval_table = wandb.Table(columns=['pred', 'truth', 'f1', 'tp', 'fp', 'fn', 'prompt_tokens', 'completion_tokens', 'total_tokens'])
+ for pred, gt in zip(_extracted, labelled_examples):
+ f1, tp, fp, fn = calculate_f1(pred['entities'], gt['entities'])
+ total_tp = total_tp + tp
+ total_fp = total_fp + fp
+ total_fn = total_fn + f1
+ eval_table.add_data(pred['entities'], gt['entities'], f1, tp, fp, fn, pred['prompt_tokens'], pred['completion_tokens'], pred['total_tokens'])
+ wandb.log({'eval_table': eval_table})
+ _overall_precision = total_tp / (total_tp + total_fp) if total_tp + total_fp else 0
+ _overall_recall = total_tp / (total_tp + total_fn) if total_tp + total_fn else 0
+ _overall_f1 = 2 * _overall_precision * _overall_recall / (_overall_precision + _overall_recall) if _overall_precision + _overall_recall else 0
+ return (_overall_precision, _overall_recall, _overall_f1)
+
+ return (evaluate,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Run our pipeline:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To start logging to W&B, you can call `wandb.init` and pass in the config to track the configurations you're experimenting with currently.
+
+ As you experiment, you can call `wandb.log` to track your work. This will log the metrics to W&B. Finally, we'll call `wandb.finish` to stop tracking. This will be tracked as one "Run" in W&B.
+
+ You'll be given a link to W&B to see all of your logs.
+ """)
+ return
+
+
+@app.cell
+def _(
+ evaluate,
+ inference,
+ labelled_examples,
+ model,
+ system_prompt,
+ temperature,
+ template,
+ wandb,
+):
+ NUM_EXAMPLES = 50
+ wandb.init(project='prompts_eval', config={'system_prompt': system_prompt, 'template': template, 'model': model, 'temperature': temperature})
+ _extracted, _elapsed = inference(labelled_examples[:NUM_EXAMPLES], system_prompt, template, model, temperature)
+ _overall_precision, _overall_recall, _overall_f1 = evaluate(_extracted, labelled_examples[:NUM_EXAMPLES])
+ _total_tokens_sum = sum([pred['total_tokens'] for pred in _extracted])
+ _completion_tokens_sum = sum([pred['completion_tokens'] for pred in _extracted])
+ _prompt_tokens_sum = sum([pred['prompt_tokens'] for pred in _extracted])
+ wandb.log({'precision': _overall_precision, 'recall': _overall_recall, 'f1': _overall_f1, 'time_elapsed_total': _elapsed, 'prompt_tokens': _prompt_tokens_sum, 'completion_tokens': _completion_tokens_sum, 'total_tokens': _total_tokens_sum})
+ wandb.finish()
+ return (NUM_EXAMPLES,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Set up experiments
+
+ Start a W&B run per experiment with `wandb.init`, store experiment details in `config` arg. Log results with `wandb.log`. Call `wandb.finish` to end experiment. Loop over all options in grid search to find best configuration.
+ """)
+ return
+
+
+@app.cell
+def _(NUM_EXAMPLES, evaluate, inference, labelled_examples, template, wandb):
+ system_prompts = ['Extract the entities from the follow sentence.', 'You are an excellent entity extractor reading newspapers and extracting orgs, people and locations. Extract the entities from the follow sentence.']
+ for system_prompt_1 in system_prompts:
+ for temperature_1 in [0.2, 0.6, 0.9]:
+ for model_1 in ['gpt-3.5-turbo', 'gpt-3.5-turbo-1106']:
+ wandb.init(project='prompts_eval', config={'system_prompt': system_prompt_1, 'template': template, 'model': model_1, 'temperature': temperature_1})
+ _extracted, _elapsed = inference(labelled_examples[:NUM_EXAMPLES], system_prompt_1, template, model_1, temperature_1)
+ _overall_precision, _overall_recall, _overall_f1 = evaluate(_extracted, labelled_examples[:NUM_EXAMPLES])
+ _total_tokens_sum = sum([pred['total_tokens'] for pred in _extracted])
+ _completion_tokens_sum = sum([pred['completion_tokens'] for pred in _extracted])
+ _prompt_tokens_sum = sum([pred['prompt_tokens'] for pred in _extracted])
+ wandb.log({'precision': _overall_precision, 'recall': _overall_recall, 'f1': _overall_f1, 'time_elapsed_total': _elapsed, 'prompt_tokens': _prompt_tokens_sum, 'completion_tokens': _completion_tokens_sum, 'total_tokens': _total_tokens_sum})
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Conclusion
+
+ You've learned how to use W&B to track evaluations of your LLM applications.
+ You've used `wandb.init` to start tracking, `wandb.log` to log summary evaluation metrics and `wandb.Table` to track individual predictions & scores.
+ We've also shared some best practices to format your code to make it easier to run evaluations in parallel and track every iteration.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Trace your LLM application
+
+ If you want to learn more and you're using complex pipelines of LLM calls, you can leverage W&B Prompts to view traces of your application and see inputs & ouputs of each LLM or function call.
+
+ Learn more about W&B Prompts in the documentation here: [https://docs.wandb.ai/guides/prompts](https://docs.wandb.ai/guides/prompts)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/prompts-w-b-prompts-with-custom-columns/prompts_w_b_prompts_with_custom_columns.py b/marimo/convert/prompts-w-b-prompts-with-custom-columns/prompts_w_b_prompts_with_custom_columns.py
new file mode 100644
index 00000000..626e023e
--- /dev/null
+++ b/marimo/convert/prompts-w-b-prompts-with-custom-columns/prompts_w_b_prompts_with_custom_columns.py
@@ -0,0 +1,395 @@
+# /// script
+# dependencies = ["langchain", "openai", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **[Weights & Biases Prompts](https://docs.wandb.ai/guides/prompts?utm_source=code&utm_medium=colab&utm_campaign=prompts)** is a suite of LLMOps tools built for the development of LLM-powered applications.
+
+ Use W&B Prompts to visualize and inspect the execution flow of your LLMs, analyze the inputs and outputs of your LLMs, view the intermediate results and securely store and manage your prompts and LLM chain configurations.
+
+ #### [🪄 View Prompts In Action](https://wandb.ai/timssweeney/prompts-demo/)
+
+ **In this notebook we will demostrate W&B Prompts:**
+
+ - Using our 1-line LangChain integration
+ - Using our Trace class when building your own LLM Pipelines
+
+ See here for the full [W&B Prompts documentation](https://docs.wandb.ai/guides/prompts)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installation
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb>=0.15.4 !pip install "wandb>=0.15.4" -qqq
+ # packages added via marimo's package management: langchain>=0.0.218 openai !pip install "langchain>=0.0.218" openai -qqq
+ return
+
+
+@app.cell
+def _():
+ import langchain
+ assert langchain.__version__ >= "0.0.218", "Please ensure you are using LangChain v0.0.188 or higher"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup
+
+ This demo requires that you have an [OpenAI key](https://platform.openai.com)
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ from getpass import getpass
+
+ if os.getenv("OPENAI_API_KEY") is None:
+ os.environ["OPENAI_API_KEY"] = getpass("Paste your OpenAI key from: https://platform.openai.com/account/api-keys\n")
+ assert os.getenv("OPENAI_API_KEY", "").startswith("sk-"), "This doesn't look like a valid OpenAI API key"
+ print("OpenAI API key configured")
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Prompts
+
+ W&B Prompts consists of three main components:
+
+ **Trace table**: Overview of the inputs and outputs of a chain.
+
+ **Trace timeline**: Displays the execution flow of the chain and is color-coded according to component types.
+
+ **Model architecture**: View details about the structure of the chain and the parameters used to initialize each component of the chain.
+
+ After running this section, you will see a new panel automatically created in your workspace, showing each execution, the trace, and the model architecture
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Maths with LangChain
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Set the `LANGCHAIN_WANDB_TRACING` environment variable as well as any other relevant [W&B environment variables](https://docs.wandb.ai/guides/track/environment-variables). This could includes a W&B project name, team name, and more. See [wandb.init](https://docs.wandb.ai/ref/python/init) for a full list of arguments.
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ os.environ["LANGCHAIN_WANDB_TRACING"] = "true"
+ os.environ["WANDB_PROJECT"] = "langchain-testing"
+ return
+
+
+@app.cell
+def _():
+ from langchain.chat_models import ChatOpenAI
+ from langchain.agents import load_tools, initialize_agent, AgentType
+
+ return AgentType, ChatOpenAI, initialize_agent, load_tools
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Create a standard math Agent using LangChain
+ """)
+ return
+
+
+@app.cell
+def _(AgentType, ChatOpenAI, initialize_agent, load_tools):
+ llm = ChatOpenAI(temperature=0)
+ tools = load_tools(["llm-math"], llm=llm)
+ math_agent = initialize_agent(tools,
+ llm,
+ agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)
+ return (math_agent,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use LangChain as normal by calling your Agent.
+
+ You will see a Weights & Biases run start and you will be prompted to create a new API key at [wandb.ai/settings](https://wandb.ai/settings) if you haven't already. Store your API key securely. It can only be viewed once when created. Once you enter your API key, the inputs and outputs of your Agent calls will start to be streamed to the Weights & Biases App.
+ """)
+ return
+
+
+@app.cell
+def _(math_agent):
+ # some sample maths questions
+ questions = [
+ "Find the square root of 5.4.",
+ "What is 3 divided by 7.34 raised to the power of pi?",
+ "What is the sin of 0.47 radians, divided by the cube root of 27?"
+ ]
+
+ for question in questions:
+ try:
+ # call your Agent as normal
+ answer = math_agent.run(question)
+ print(answer)
+ except Exception as e:
+ # any errors will be also logged to Weights & Biases
+ print(e)
+ pass
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once each Agent execution completes, all calls in your LangChain object will be logged to Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### LangChain Context Manager
+ Depending on your use case, you might instead prefer to use a context manager to manage your logging to W&B.
+
+ **✨ New: Custom columns** can be logged directly to W&B to display in the same Trace Table with this snippet:
+ ```python
+ import wandb
+ wandb.log(custom_metrics_dict, commit=False})
+ ```
+ Use `commit=False` to make sure that metadata is logged to the same row of the Trace Table as the LangChain output.
+ """)
+ return
+
+
+@app.cell
+def _(math_agent, os):
+ from langchain.callbacks import wandb_tracing_enabled
+ import wandb # To enable custom column logging with wandb.run.log()
+
+ # unset the environment variable and use a context manager instead
+ if "LANGCHAIN_WANDB_TRACING" in os.environ:
+ del os.environ["LANGCHAIN_WANDB_TRACING"]
+
+ # enable tracing using a context manager
+ with wandb_tracing_enabled():
+ for i in range (10):
+ # Log any custom columns you'd like to add to the Trace Table
+ wandb.log({"custom_column": i}, commit=False)
+ try:
+ math_agent.run(f"What is {i} raised to .123243 power?") # this should be traced
+ except:
+ pass
+
+ math_agent.run("What is 2 raised to .123243 power?") # this should not be traced
+ return (wandb,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Non-Lang Chain Implementation
+
+ A W&B Trace is created by logging 1 or more "spans". A root span is expected, which can accept nested child spans, which can in turn accept their own child spans. A Span represents a unit of work, Spans can have type `AGENT`, `TOOL`, `LLM` or `CHAIN`
+
+ When logging with Trace, a single W&B run can have multiple calls to a LLM, Tool, Chain or Agent logged to it, there is no need to start a new W&B run after each generation from your model or pipeline, instead each call will be appended to the Trace Table.
+
+ In this quickstart, we will how to log a single call to an OpenAI model to W&B Trace as a single span. Then we will show how to log a more complex series of nested spans.
+
+ ## Logging with W&B Trace
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Call wandb.init to start a W&B run. Here you can pass a W&B project name as well as an entity name (if logging to a W&B Team), as well as a config and more. See wandb.init for the full list of arguments.
+
+ You will see a Weights & Biases run start and be prompted to create a new API key at [wandb.ai/settings](https://wandb.ai/settings) if you haven't already. Store your API key securely. It can only be viewed once when created. Once you enter your API key, the inputs and outputs of your Agent calls will start to be streamed to the Weights & Biases App.
+
+ **Note:** A W&B run supports logging as many traces you needed to a single run, i.e. you can make multiple calls of `run.log` without the need to create a new run each time
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # start a wandb run to log to
+ wandb.init(project='trace-example')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can also set the entity argument in wandb.init if logging to a W&B Team.
+
+ ### Logging a single Span
+ Now we will query OpenAI times and log the results to a W&B Trace. We will log the inputs and outputs, start and end times, whether the OpenAI call was successful, the token usage, and additional metadata.
+
+ You can see the full description of the arguments to the Trace class [here](https://soumik12345.github.io/wandb-addons/prompts/tracer/).
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ import openai
+ import datetime
+ from wandb.sdk.data_types.trace_tree import Trace
+ openai.api_key = os.environ['OPENAI_API_KEY']
+ model_name = 'gpt-3.5-turbo'
+ temperature = 0.7
+ # define your conifg
+ system_message = 'You are a helpful assistant that always replies in 3 concise bullet points using markdown.'
+ queries_ls = ['What is the capital of France?', 'How do I boil an egg?' * 10000, 'What to do if the aliens arrive?']
+ for _query in queries_ls:
+ _messages = [{'role': 'system', 'content': system_message}, {'role': 'user', 'content': _query}]
+ _start_time_ms = datetime.datetime.now().timestamp() * 1000
+ try:
+ _response = openai.ChatCompletion.create(model=model_name, messages=_messages, temperature=temperature) # deliberately trigger an openai error
+ end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ status = 'success'
+ status_message = (None,)
+ _response_text = _response['choices'][0]['message']['content']
+ _token_usage = _response['usage'].to_dict()
+ except Exception as e:
+ end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ status = 'error'
+ status_message = str(e)
+ _response_text = ''
+ _token_usage = {}
+ _root_span = Trace(name='root_span', kind='llm', status_code=status, status_message=status_message, metadata={'temperature': temperature, 'token_usage': _token_usage, 'model_name': model_name}, start_time_ms=_start_time_ms, end_time_ms=end_time_ms, inputs={'system_prompt': system_message, 'query': _query}, outputs={'response': _response_text})
+ _root_span.log(name='openai_trace') # logged in milliseconds # create a span in wandb # kind can be "llm", "chain", "agent" or "tool" # log the span to wandb
+ return Trace, datetime, model_name, openai, system_message, temperature
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Logging a LLM pipeline using nested Spans
+
+ In this example we will simulate an Agent being called, which then calls a LLM Chain, which calls an OpenAI LLM and then the Agent "calls" a Calculator tool.
+
+ The inputs, outputs and metadata for each step in the execution of our "Agent" is logged in its own span. Spans can have child
+ """)
+ return
+
+
+@app.cell
+def _(Trace, datetime, model_name, openai, os, system_message, temperature):
+ import time
+ openai.api_key = os.environ['OPENAI_API_KEY']
+ _query = 'How many days until the next US election?'
+ _start_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ # The query our agent has to answer
+ _root_span = Trace(name='MyAgent', kind='agent', start_time_ms=_start_time_ms, metadata={'user': 'optimus_12'})
+ chain_span = Trace(name='LLMChain', kind='chain', start_time_ms=_start_time_ms)
+ # part 1 - an Agent is started...
+ _root_span.add_child(chain_span)
+ _messages = [{'role': 'system', 'content': system_message}, {'role': 'user', 'content': _query}]
+ _response = openai.ChatCompletion.create(model=model_name, messages=_messages, temperature=temperature)
+ llm_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ _response_text = _response['choices'][0]['message']['content']
+ _token_usage = _response['usage'].to_dict()
+ llm_span = Trace(name='OpenAI', kind='llm', status_code='success', metadata={'temperature': temperature, 'token_usage': _token_usage, 'model_name': model_name}, start_time_ms=_start_time_ms, end_time_ms=llm_end_time_ms, inputs={'system_prompt': system_message, 'query': _query}, outputs={'response': _response_text})
+ chain_span.add_child(llm_span)
+ chain_span.add_inputs_and_outputs(inputs={'query': _query}, outputs={'response': _response_text})
+ # part 2 - The Agent calls into a LLMChain..
+ chain_span._span.end_time_ms = llm_end_time_ms
+ time.sleep(3)
+ days_to_election = 117
+ tool_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ tool_span = Trace(name='Calculator', kind='tool', status_code='success', start_time_ms=llm_end_time_ms, end_time_ms=tool_end_time_ms, inputs={'input': _response_text}, outputs={'result': days_to_election})
+ # add the Chain span as a child of the root
+ _root_span.add_child(tool_span)
+ _root_span.add_inputs_and_outputs(inputs={'query': _query}, outputs={'result': days_to_election})
+ _root_span._span.end_time_ms = tool_end_time_ms
+ # part 3 - the LLMChain calls an OpenAI LLM...
+ # add the LLM span as a child of the Chain span...
+ # update the end time of the Chain span
+ # update the Chain span's end time
+ # part 4 - the Agent then calls a Tool...
+ # create a Tool span
+ # add the TOOL span as a child of the root
+ # part 5 - the final results from the tool are added
+ # part 6 - log all spans to W&B by logging the root span
+ _root_span.log(name='openai_trace')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once each Agent execution completes, all calls in your LangChain object will be logged to Weights & Biases
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/prompts-wandb-llm-qa-bot/prompts_wandb_llm_qa_bot.py b/marimo/convert/prompts-wandb-llm-qa-bot/prompts_wandb_llm_qa_bot.py
new file mode 100644
index 00000000..16ceecb1
--- /dev/null
+++ b/marimo/convert/prompts-wandb-llm-qa-bot/prompts_wandb_llm_qa_bot.py
@@ -0,0 +1,529 @@
+# /// script
+# dependencies = ["chromadb", "langchain", "openai", "pytube", "tiktoken", "wandb", "youtube-transcript-api"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Building an LLM App for Document Retrieval / Extraction
+
+ This tutorial runs through [this report](https://wandb.ai/gladiator/gradient_dissent_qabot/reports/Building-a-Q-A-Bot-for-Weights-Biases-Gradient-Dissent-Podcast--Vmlldzo0MTcyMDQz) on how to build a basic LLM App for retrieval-augmented question-answering.
+ - Track datasets and embeddings as artifacts
+ - Track prompts and chain executions
+ - Log token counts and cost
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb langchain pytube tiktoken openai youtube-transcript-api chromadb !pip install -qqq wandb langchain pytube tiktoken openai youtube-transcript-api chromadb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Set up OpenAI API Key
+ """)
+ return
+
+
+@app.cell
+def _():
+ from getpass import getpass
+ import os
+
+ if os.getenv("OPENAI_API_KEY") is None:
+ if any(['VSCODE' in x for x in os.environ.keys()]):
+ print('Please enter password in the VS Code prompt at the top of your VS Code window!')
+ os.environ["OPENAI_API_KEY"] = getpass("Paste your OpenAI key from: https://platform.openai.com/account/api-keys\n")
+
+ assert os.getenv("OPENAI_API_KEY", "").startswith("sk-"), "This doesn't look like a valid OpenAI API key"
+ print("OpenAI API key configured")
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Set up config and environment variables
+ - NOTE: set the `entity` to your username or team name
+ - Set wandb [environment variables](https://docs.wandb.ai/guides/track/environment-variables) to change behavior of logging
+ - `ENTITY` - username or team where your projects live
+ - `PROJECT` - project where your runs will live
+ - `LANGCHAIN_WANDB_TRACING` - automatically logs langchain traces, inputs and outputs as part of runs in Weights and Biases
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ from dataclasses import dataclass
+ from pathlib import Path
+ project_name = 'gradient-dissent-qabot'
+ entity = 'wandb'
+ TOTAL_EPISODES = 5 #@param
+ playlist_url = 'https://www.youtube.com/playlist?list=PLD80i8An1OEEb1jP0sjEyiLG8ULRXFob_' #@param
+ root_data_dir = Path('/contents/data')
+ root_artifact_dir = Path('downloaded_artifacts')
+ yt_podcast_data_artifact = f'{entity}/{project_name}/yt_podcast_transcript:latest'
+ summarized_data_artifact = f'{entity}/{project_name}/summarized_podcasts:latest'
+ summarized_que_data_artifact = f'{entity}/{project_name}/summarized_que_podcasts:latest'
+ transcript_embeddings_artifact = f'{entity}/{project_name}/transcript_embeddings:latest'
+ os.makedirs('/contents/data', exist_ok=True)
+ os.environ['LANGCHAIN_WANDB_TRACING'] = 'true'
+ os.environ['WANDB_PROJECT'] = project_name
+ os.environ['WANDB_ENTITY'] = entity
+ return (
+ TOTAL_EPISODES,
+ entity,
+ playlist_url,
+ project_name,
+ root_artifact_dir,
+ root_data_dir,
+ summarized_data_artifact,
+ transcript_embeddings_artifact,
+ yt_podcast_data_artifact,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log in to W&B
+ - You can explicitly login using `wandb login` or `wandb.login()` (See below)
+ - Alternatively you can set environment variables. There are several env variables which you can set to change the behavior of W&B logging. The most important are:
+ - `WANDB_API_KEY` - create a new API key in your "Settings" section under your profile at [wandb.ai/settings](https://wandb.ai/settings)
+ - `WANDB_BASE_URL` - this is the url of the W&B server (You only need this if you are using a private instance)
+ - Create a new API key in "Profile" -> "Settings" in the W&B App. Store your API key securely. It can only be viewed once when created.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import time
+ import pandas as pd
+ from langchain.document_loaders import YoutubeLoader
+ from pytube import Playlist, YouTube
+ from tqdm import tqdm
+
+ def retry_access_yt_object(url, max_retries=5, interval_secs=5):
+ """
+ Retries creating a YouTube object with the given URL and accessing its title several times
+ with a given interval in seconds, until it succeeds or the maximum number of attempts is reached.
+ If the object still cannot be created or the title cannot be accessed after the maximum number
+ of attempts, the last exception is raised.
+ """
+ last_exception = None
+ for i in range(max_retries):
+ try:
+ yt = YouTube(url)
+ title = yt.title
+ return yt
+ except Exception as err: # Access the title of the YouTube object.
+ last_exception = err # Return the YouTube object if successful.
+ print(f'Failed to create YouTube object or access title. Retrying... ({i + 1}/{max_retries})')
+ time.sleep(interval_secs) # Keep track of the last exception raised.
+ raise last_exception # Wait for the specified interval before retrying. # If the YouTube object still cannot be created or the title cannot be accessed after the maximum number of attempts, raise the last exception.
+
+ return Playlist, YoutubeLoader, pd, retry_access_yt_object, tqdm
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log Data Snapshots as Artifacts
+
+ W&B is very unopinionated with regard to how you track your experiments. We could log data in any number of ways.
+ * Log one artifact which represents all the data - training, validation, and test data to one artifact
+ * Log several artifacts - one for each of the training, validation, and test data loaders.
+
+ It is a matter of what best suites your needs and workflows and expectations.
+
+ ### Anatomy of an artifact
+
+ The `Artifact` class will correspond to an entry in the W&B Artifact registry. The artifact has
+ * a name
+ * a type
+ * metadata
+ * description
+ * files, directory of files, or references
+
+ Example usage
+ ```
+ run = wandb.init(project = "my-project")
+ artifact = wandb.Artifact(name = "my_artifact", type = "data")
+ artifact.add_file("/path/to/my/file.txt")
+ run.log_artifact(artifact)
+ run.finish()
+ ```
+ """)
+ return
+
+
+@app.cell
+def _(
+ Playlist,
+ TOTAL_EPISODES,
+ YoutubeLoader,
+ entity,
+ pd,
+ playlist_url,
+ project_name,
+ retry_access_yt_object,
+ root_data_dir,
+ tqdm,
+ wandb,
+):
+ run = wandb.init(project=project_name, entity=entity, job_type='dataset')
+ playlist = Playlist(playlist_url)
+ playlist_video_urls = playlist.video_urls[0:TOTAL_EPISODES]
+ print(f'There are total {len(playlist_video_urls)} videos in the playlist.')
+ video_data = []
+ for video in tqdm(playlist_video_urls, total=len(playlist_video_urls)):
+ try:
+ curr_video_data = {}
+ yt = retry_access_yt_object(video, max_retries=25, interval_secs=2)
+ curr_video_data['title'] = yt.title
+ curr_video_data['url'] = video
+ curr_video_data['duration'] = yt.length
+ curr_video_data['publish_date'] = yt.publish_date.strftime('%Y-%m-%d')
+ loader = YoutubeLoader.from_youtube_url(video)
+ transcript = loader.load()[0].page_content
+ transcript = ' '.join(transcript.split())
+ curr_video_data['transcript'] = transcript
+ curr_video_data['total_words'] = len(transcript.split())
+ video_data.append(curr_video_data)
+ except Exception as inst:
+ print(type(inst))
+ print(inst.args)
+ print(inst)
+ print(f'Failed to scrape {video}') # the exception type
+ print(f'Total podcast episodes scraped: {len(video_data)}') # arguments stored in .args
+ df = pd.DataFrame(video_data)
+ data_path = root_data_dir / 'yt_podcast_transcript.csv'
+ df.to_csv(data_path, index=False)
+ _artifact = wandb.Artifact('yt_podcast_transcript', type='dataset')
+ _artifact.add_file(data_path)
+ # save the scraped data to a csv file
+ # upload the scraped data to wandb
+ run.log_artifact(_artifact)
+ return df, run
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Log a wandb Table to interact with your data
+ - Here we log the dataframe of metadata about the youtube transcripts (urls, length, transcripts)
+ - This allows us to interrogate the original data (filtering, grouping, etc.)
+ """)
+ return
+
+
+@app.cell
+def _(df, run, wandb):
+ # create wandb table
+ _table = wandb.Table(dataframe=df)
+ run.log({'yt_podcast_transcript': _table})
+ run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Summarize YouTube Transcripts
+ - Here we summarize the transcripts in chunks, summarizing each chunk and then summarizing the summaries using the LangChain `load_summarize_chain`
+ - We can do this in parallel since each chunk of a transcript can be summarized independently so we employ `map_reduce`
+ """)
+ return
+
+
+@app.cell
+def _(os, pd, root_artifact_dir, wandb):
+ from langchain.callbacks import get_openai_callback
+ from langchain.chains.summarize import load_summarize_chain
+ from langchain.chat_models import ChatOpenAI
+ from langchain.document_loaders import DataFrameLoader
+ from langchain.prompts import PromptTemplate
+ from langchain.text_splitter import TokenTextSplitter
+
+ def get_data(artifact_name: str, total_episodes: int=None):
+ podcast_artifact = wandb.use_artifact(artifact_name)
+ podcast_artifact_dir = podcast_artifact.download(root_artifact_dir)
+ filename = artifact_name.split(':')[0].split('/')[-1]
+ df = pd.read_csv(os.path.join(podcast_artifact_dir, f'{filename}.csv'))
+ if total_episodes is not None:
+ df = df.iloc[:total_episodes]
+ return df
+
+ def summarize_episode(episode_df: pd.DataFrame):
+ loader = DataFrameLoader(episode_df, page_content_column='transcript')
+ data = loader.load()
+ text_splitter = TokenTextSplitter.from_tiktoken_encoder(chunk_size=1000, chunk_overlap=0)
+ docs = text_splitter.split_documents(data)
+ print(f"Number of documents for podcast {data[0].metadata['title']}: {len(docs)}")
+ llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0)
+ map_prompt = "Write a concise summary of the following short transcript from a podcast.\n Don't add your opinions or interpretations.\n\n {text}\n\n CONCISE SUMMARY:"
+ combine_prompt = 'You have been provided with summaries of chunks of transcripts from a podcast.\n Your task is to merge these intermediate summaries to create a brief and comprehensive summary of the entire podcast.\n The summary should encompass all the crucial points of the podcast.\n Ensure that the summary is atleast 2 paragraph long and effectively captures the essence of the podcast.\n {text}\n\n SUMMARY:' # load docs into langchain format
+ map_prompt_template = PromptTemplate(template=map_prompt, input_variables=['text'])
+ combine_prompt_template = PromptTemplate(template=combine_prompt, input_variables=['text'])
+ chain = load_summarize_chain(llm, chain_type='map_reduce', return_intermediate_steps=True, map_prompt=map_prompt_template, combine_prompt=combine_prompt_template)
+ summary = chain({'input_documents': docs}) # split the documents
+ return summary # initialize LLM # define map prompt # define combine prompt # initialize the summarizer chain
+
+ return (
+ ChatOpenAI,
+ DataFrameLoader,
+ PromptTemplate,
+ TokenTextSplitter,
+ get_data,
+ get_openai_callback,
+ summarize_episode,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Execute Summary Chain and log results
+ - You can instantiate a `WandbTracer` and pass in additional config about this LangChain run.
+ - Log the outputs of the chain like tokens used, cost, etc.
+ - Log the resulting summaries as artifacts
+ """)
+ return
+
+
+@app.cell
+def _(
+ TOTAL_EPISODES,
+ get_data,
+ get_openai_callback,
+ os,
+ root_data_dir,
+ summarize_episode,
+ tqdm,
+ wandb,
+ yt_podcast_data_artifact,
+):
+ from langchain.callbacks.tracers import WandbTracer
+ _tracer = WandbTracer(run_args={'job_type': 'summarize'})
+ df_1 = get_data(artifact_name=yt_podcast_data_artifact, total_episodes=TOTAL_EPISODES)
+ summaries = []
+ with get_openai_callback() as _cb:
+ for _episode in tqdm(df_1.iterrows(), total=len(df_1), desc='Summarizing episodes'):
+ _episode_data = _episode[1].to_frame().T
+ summary = summarize_episode(_episode_data)
+ summaries.append(summary['output_text'])
+ print('*' * 25)
+ print(_cb)
+ print('*' * 25)
+ wandb.log({'total_prompt_tokens': _cb.prompt_tokens, 'total_completion_tokens': _cb.completion_tokens, 'total_tokens': _cb.total_tokens, 'total_cost': _cb.total_cost})
+ df_1['summary'] = summaries
+ path_to_save = os.path.join(root_data_dir, 'summarized_podcasts.csv')
+ df_1.to_csv(path_to_save, index=False)
+ _artifact = wandb.Artifact('summarized_podcasts', type='dataset')
+ _artifact.add_file(path_to_save)
+ wandb.log_artifact(_artifact)
+ _table = wandb.Table(dataframe=df_1)
+ wandb.log({'summarized_podcasts': _table})
+ _tracer.finish()
+ return (WandbTracer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Embed the contents of the YouTube transcripts
+ - Here we use OpenAI embeddings and [ChromaDB](https://www.trychroma.com/) to embed the summaries to make them queriable via vector similarity search when we ask contextual questions to the LLM
+ - Use `wandb.log` and artifacts to log the resulting ChromaDB serialized embeddings.
+ """)
+ return
+
+
+@app.cell
+def _(
+ DataFrameLoader,
+ TokenTextSplitter,
+ os,
+ pd,
+ root_artifact_dir,
+ root_data_dir,
+ wandb,
+):
+ from dataclasses import asdict
+ from langchain.embeddings.openai import OpenAIEmbeddings
+ from langchain.vectorstores import Chroma
+
+ def get_data_1(artifact_name: str, total_episodes=None):
+ podcast_artifact = wandb.use_artifact(artifact_name, type='dataset')
+ podcast_artifact_dir = podcast_artifact.download(root_artifact_dir)
+ filename = artifact_name.split(':')[0].split('/')[-1]
+ df = pd.read_csv(os.path.join(podcast_artifact_dir, f'{filename}.csv'))
+ if total_episodes is not None:
+ df = df.iloc[:total_episodes]
+ return df
+
+ def create_embeddings(episode_df: pd.DataFrame, index: int):
+ loader = DataFrameLoader(episode_df, page_content_column='transcript')
+ data = loader.load()
+ text_splitter = TokenTextSplitter.from_tiktoken_encoder(chunk_size=1000, chunk_overlap=0)
+ docs = text_splitter.split_documents(data)
+ title = data[0].metadata['title']
+ print(f'Number of documents for podcast {title}: {len(docs)}')
+ embeddings = OpenAIEmbeddings()
+ db = Chroma.from_documents(docs, embeddings, persist_directory=os.path.join(root_data_dir / 'chromadb', str(index)))
+ db.persist() # load docs into langchain format # split the documents # initialize embedding engine
+
+ return Chroma, OpenAIEmbeddings, create_embeddings, get_data_1
+
+
+@app.cell
+def _(
+ TOTAL_EPISODES,
+ WandbTracer,
+ create_embeddings,
+ get_data_1,
+ get_openai_callback,
+ root_data_dir,
+ summarized_data_artifact,
+ tqdm,
+ wandb,
+):
+ _tracer = WandbTracer(run_args={'job_type': 'embed_transcripts'})
+ df_2 = get_data_1(artifact_name=summarized_data_artifact, total_episodes=TOTAL_EPISODES)
+ with get_openai_callback() as _cb:
+ for _episode in tqdm(df_2.iterrows(), total=len(df_2), desc='Embedding transcripts'):
+ _episode_data = _episode[1].to_frame().T
+ create_embeddings(_episode_data, index=_episode[0])
+ print('*' * 25)
+ print(_cb)
+ print('*' * 25)
+ wandb.log({'total_prompt_tokens': _cb.prompt_tokens, 'total_completion_tokens': _cb.completion_tokens, 'total_tokens': _cb.total_tokens, 'total_cost': _cb.total_cost})
+ _artifact = wandb.Artifact('transcript_embeddings', type='dataset')
+ _artifact.add_dir(root_data_dir / 'chromadb')
+ wandb.log_artifact(_artifact)
+ _tracer.finish()
+ return (df_2,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Ask Questions Against your Summarized Documents
+
+ Finally we tie everything together:
+ 1. We can pull down our ChromaDB embeddings from W&B
+ 2. Pass them along with a prompt template for QA to the `RetrievalQA` chain and start asking questions!
+ """)
+ return
+
+
+@app.cell
+def _(
+ ChatOpenAI,
+ Chroma,
+ OpenAIEmbeddings,
+ PromptTemplate,
+ chromadb_dir,
+ df_2,
+ get_openai_callback,
+ os,
+):
+ from langchain.chains import RetrievalQA
+
+ def get_answer(podcast: str, question: str):
+ index = df_2[df_2['title'] == podcast].index[0]
+ db_dir = os.path.join(chromadb_dir, str(index))
+ embeddings = OpenAIEmbeddings()
+ db = Chroma(persist_directory=db_dir, embedding_function=embeddings)
+ prompt_template = "Use the following pieces of context to answer the question.\n If you don't know the answer, just say that you don't know, don't try to make up an answer.\n Don't add your opinions or interpretations. Ensure that you complete the answer.\n If the question is not relevant to the context, just say that it is not relevant.\n\n CONTEXT:\n {context}\n\n QUESTION: {question}\n\n ANSWER:"
+ prompt = PromptTemplate(template=prompt_template, input_variables=['context', 'question'])
+ retriever = db.as_retriever()
+ retriever.search_kwargs['k'] = 2
+ qa = RetrievalQA.from_chain_type(llm=ChatOpenAI(temperature=0), chain_type='stuff', retriever=retriever, chain_type_kwargs={'prompt': prompt}, return_source_documents=True)
+ with get_openai_callback() as _cb:
+ result = qa({'query': question})
+ print(_cb)
+ answer = result['result']
+ return answer
+
+ return (get_answer,)
+
+
+@app.cell
+def _(
+ pd,
+ root_data_dir,
+ summarized_data_artifact,
+ transcript_embeddings_artifact,
+ wandb,
+):
+ # download and read data
+ api = wandb.Api()
+ artifact_df = api.artifact(summarized_data_artifact)
+ artifact_df.download(root_data_dir)
+ artifact_embeddings = api.artifact(transcript_embeddings_artifact)
+ chromadb_dir = artifact_embeddings.download(root_data_dir / 'chromadb')
+ df_path = root_data_dir / 'summarized_podcasts.csv'
+ df_3 = pd.read_csv(df_path)
+ return chromadb_dir, df_3
+
+
+@app.cell
+def _(TOTAL_EPISODES, df_3):
+ df_3['title'].tolist()[0:TOTAL_EPISODES]
+ return
+
+
+@app.cell
+def _(WandbTracer, get_answer):
+ _tracer = WandbTracer(run_args={'job_type': 'retriealQA'})
+ answer = get_answer('Enabling LLM-Powered Applications with Harrison Chase of LangChain', 'What did Harrison Chase say?')
+ print(answer)
+ _tracer.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/prompts-wandb-prompts-quickstart/prompts_wandb_prompts_quickstart.py b/marimo/convert/prompts-wandb-prompts-quickstart/prompts_wandb_prompts_quickstart.py
new file mode 100644
index 00000000..c091b4e7
--- /dev/null
+++ b/marimo/convert/prompts-wandb-prompts-quickstart/prompts_wandb_prompts_quickstart.py
@@ -0,0 +1,383 @@
+# /// script
+# dependencies = ["langchain", "openai", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **[Weights & Biases Prompts](https://docs.wandb.ai/guides/prompts?utm_source=code&utm_medium=colab&utm_campaign=prompts)** is a suite of LLMOps tools built for the development of LLM-powered applications.
+
+ Use W&B Prompts to visualize and inspect the execution flow of your LLMs, analyze the inputs and outputs of your LLMs, view the intermediate results and securely store and manage your prompts and LLM chain configurations.
+
+ #### [🪄 View Prompts In Action](https://wandb.ai/timssweeney/prompts-demo/)
+
+ **In this notebook we will demostrate W&B Prompts:**
+
+ - Using our 1-line LangChain integration
+ - Using our Trace class when building your own LLM Pipelines
+
+ See here for the full [W&B Prompts documentation](https://docs.wandb.ai/guides/prompts)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installation
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb>=0.15.4 !pip install "wandb>=0.15.4" -qqq
+ # packages added via marimo's package management: langchain>=0.0.218 openai !pip install "langchain>=0.0.218" openai -qqq
+ return
+
+
+@app.cell
+def _():
+ import langchain
+ assert langchain.__version__ >= "0.0.218", "Please ensure you are using LangChain v0.0.188 or higher"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup
+
+ This demo requires that you have an [OpenAI key](https://platform.openai.com)
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ from getpass import getpass
+
+ if os.getenv("OPENAI_API_KEY") is None:
+ os.environ["OPENAI_API_KEY"] = getpass("Paste your OpenAI key from: https://platform.openai.com/account/api-keys\n")
+ assert os.getenv("OPENAI_API_KEY", "").startswith("sk-"), "This doesn't look like a valid OpenAI API key"
+ print("OpenAI API key configured")
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Prompts
+
+ W&B Prompts consists of three main components:
+
+ **Trace table**: Overview of the inputs and outputs of a chain.
+
+ **Trace timeline**: Displays the execution flow of the chain and is color-coded according to component types.
+
+ **Model architecture**: View details about the structure of the chain and the parameters used to initialize each component of the chain.
+
+ After running this section, you will see a new panel automatically created in your workspace, showing each execution, the trace, and the model architecture
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Maths with LangChain
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Set the `LANGCHAIN_WANDB_TRACING` environment variable as well as any other relevant [W&B environment variables](https://docs.wandb.ai/guides/track/environment-variables). This could includes a W&B project name, team name, and more. See [wandb.init](https://docs.wandb.ai/ref/python/init) for a full list of arguments.
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ os.environ["LANGCHAIN_WANDB_TRACING"] = "true"
+ os.environ["WANDB_PROJECT"] = "langchain-testing"
+ return
+
+
+@app.cell
+def _():
+ from langchain.chat_models import ChatOpenAI
+ from langchain.agents import load_tools, initialize_agent, AgentType
+
+ return AgentType, ChatOpenAI, initialize_agent, load_tools
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Create a standard math Agent using LangChain
+ """)
+ return
+
+
+@app.cell
+def _(AgentType, ChatOpenAI, initialize_agent, load_tools):
+ llm = ChatOpenAI(temperature=0)
+ tools = load_tools(["llm-math"], llm=llm)
+ math_agent = initialize_agent(tools,
+ llm,
+ agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)
+ return (math_agent,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use LangChain as normal by calling your Agent.
+
+ You will see a Weights & Biases run start and you will be prompted to create a new API key at [wandb.ai/settings](https://wandb.ai/settings) if you haven't already. Store your API key securely. It can only be viewed once when created. Once you enter your API key, the inputs and outputs of your Agent calls will start to be streamed to the Weights & Biases App.
+ """)
+ return
+
+
+@app.cell
+def _(math_agent):
+ # some sample maths questions
+ questions = [
+ "Find the square root of 5.4.",
+ "What is 3 divided by 7.34 raised to the power of pi?",
+ "What is the sin of 0.47 radians, divided by the cube root of 27?"
+ ]
+
+ for question in questions:
+ try:
+ # call your Agent as normal
+ answer = math_agent.run(question)
+ print(answer)
+ except Exception as e:
+ # any errors will be also logged to Weights & Biases
+ print(e)
+ pass
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once each Agent execution completes, all calls in your LangChain object will be logged to Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### LangChain Context Manager
+ Depending on your use case, you might instead prefer to use a context manager to manage your logging to W&B:
+ """)
+ return
+
+
+@app.cell
+def _(math_agent, os):
+ from langchain.callbacks import wandb_tracing_enabled
+
+ # unset the environment variable and use a context manager instead
+ if "LANGCHAIN_WANDB_TRACING" in os.environ:
+ del os.environ["LANGCHAIN_WANDB_TRACING"]
+
+ # enable tracing using a context manager
+ with wandb_tracing_enabled():
+ math_agent.run("What is 5 raised to .123243 power?") # this should be traced
+
+ math_agent.run("What is 2 raised to .123243 power?") # this should not be traced
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Non-Lang Chain Implementation
+
+ A W&B Trace is created by logging 1 or more "spans". A root span is expected, which can accept nested child spans, which can in turn accept their own child spans. A Span represents a unit of work, Spans can have type `AGENT`, `TOOL`, `LLM` or `CHAIN`
+
+ When logging with Trace, a single W&B run can have multiple calls to a LLM, Tool, Chain or Agent logged to it, there is no need to start a new W&B run after each generation from your model or pipeline, instead each call will be appended to the Trace Table.
+
+ In this quickstart, we will how to log a single call to an OpenAI model to W&B Trace as a single span. Then we will show how to log a more complex series of nested spans.
+
+ ## Logging with W&B Trace
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Call wandb.init to start a W&B run. Here you can pass a W&B project name as well as an entity name (if logging to a W&B Team), as well as a config and more. See wandb.init for the full list of arguments.
+
+ You will see a Weights & Biases run start and be prompted to create a new API key at [wandb.ai/settings](https://wandb.ai/settings) if you haven't already. Store your API key securely. It can only be viewed once when created. Once you enter your API key, the inputs and outputs of your Agent calls will start to be streamed to the Weights & Biases App.
+
+ **Note:** A W&B run supports logging as many traces you needed to a single run, i.e. you can make multiple calls of `run.log` without the need to create a new run each time
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ # start a wandb run to log to
+ wandb.init(project="trace-example")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can also set the entity argument in wandb.init if logging to a W&B Team.
+
+ ### Logging a single Span
+ Now we will query OpenAI times and log the results to a W&B Trace. We will log the inputs and outputs, start and end times, whether the OpenAI call was successful, the token usage, and additional metadata.
+
+ You can see the full description of the arguments to the Trace class [here](https://soumik12345.github.io/wandb-addons/prompts/tracer/).
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ import openai
+ import datetime
+ from wandb.sdk.data_types.trace_tree import Trace
+ openai.api_key = os.environ['OPENAI_API_KEY']
+ model_name = 'gpt-3.5-turbo'
+ temperature = 0.7
+ # define your conifg
+ system_message = 'You are a helpful assistant that always replies in 3 concise bullet points using markdown.'
+ queries_ls = ['What is the capital of France?', 'How do I boil an egg?' * 10000, 'What to do if the aliens arrive?']
+ for _query in queries_ls:
+ _messages = [{'role': 'system', 'content': system_message}, {'role': 'user', 'content': _query}]
+ _start_time_ms = datetime.datetime.now().timestamp() * 1000
+ try:
+ _response = openai.ChatCompletion.create(model=model_name, messages=_messages, temperature=temperature) # deliberately trigger an openai error
+ end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ status = 'success'
+ status_message = (None,)
+ _response_text = _response['choices'][0]['message']['content']
+ _token_usage = _response['usage'].to_dict()
+ except Exception as e:
+ end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ status = 'error'
+ status_message = str(e)
+ _response_text = ''
+ _token_usage = {}
+ _root_span = Trace(name='root_span', kind='llm', status_code=status, status_message=status_message, metadata={'temperature': temperature, 'token_usage': _token_usage, 'model_name': model_name}, start_time_ms=_start_time_ms, end_time_ms=end_time_ms, inputs={'system_prompt': system_message, 'query': _query}, outputs={'response': _response_text})
+ _root_span.log(name='openai_trace') # logged in milliseconds # create a span in wandb # kind can be "llm", "chain", "agent" or "tool" # log the span to wandb
+ return Trace, datetime, model_name, openai, system_message, temperature
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Logging a LLM pipeline using nested Spans
+
+ In this example we will simulate an Agent being called, which then calls a LLM Chain, which calls an OpenAI LLM and then the Agent "calls" a Calculator tool.
+
+ The inputs, outputs and metadata for each step in the execution of our "Agent" is logged in its own span. Spans can have child
+ """)
+ return
+
+
+@app.cell
+def _(Trace, datetime, model_name, openai, os, system_message, temperature):
+ import time
+ openai.api_key = os.environ['OPENAI_API_KEY']
+ _query = 'How many days until the next US election?'
+ _start_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ # The query our agent has to answer
+ _root_span = Trace(name='MyAgent', kind='agent', start_time_ms=_start_time_ms, metadata={'user': 'optimus_12'})
+ chain_span = Trace(name='LLMChain', kind='chain', start_time_ms=_start_time_ms)
+ # part 1 - an Agent is started...
+ _root_span.add_child(chain_span)
+ _messages = [{'role': 'system', 'content': system_message}, {'role': 'user', 'content': _query}]
+ _response = openai.ChatCompletion.create(model=model_name, messages=_messages, temperature=temperature)
+ llm_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ _response_text = _response['choices'][0]['message']['content']
+ _token_usage = _response['usage'].to_dict()
+ llm_span = Trace(name='OpenAI', kind='llm', status_code='success', metadata={'temperature': temperature, 'token_usage': _token_usage, 'model_name': model_name}, start_time_ms=_start_time_ms, end_time_ms=llm_end_time_ms, inputs={'system_prompt': system_message, 'query': _query}, outputs={'response': _response_text})
+ chain_span.add_child(llm_span)
+ chain_span.add_inputs_and_outputs(inputs={'query': _query}, outputs={'response': _response_text})
+ # part 2 - The Agent calls into a LLMChain..
+ chain_span._span.end_time_ms = llm_end_time_ms
+ time.sleep(3)
+ days_to_election = 117
+ tool_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+ tool_span = Trace(name='Calculator', kind='tool', status_code='success', start_time_ms=llm_end_time_ms, end_time_ms=tool_end_time_ms, inputs={'input': _response_text}, outputs={'result': days_to_election})
+ # add the Chain span as a child of the root
+ _root_span.add_child(tool_span)
+ _root_span.add_inputs_and_outputs(inputs={'query': _query}, outputs={'result': days_to_election})
+ _root_span._span.end_time_ms = tool_end_time_ms
+ # part 3 - the LLMChain calls an OpenAI LLM...
+ # add the LLM span as a child of the Chain span...
+ # update the end time of the Chain span
+ # update the Chain span's end time
+ # part 4 - the Agent then calls a Tool...
+ # create a Tool span
+ # add the TOOL span as a child of the root
+ # part 5 - the final results from the tool are added
+ # part 6 - log all spans to W&B by logging the root span
+ _root_span.log(name='openai_trace')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once each Agent execution completes, all calls in your LangChain object will be logged to Weights & Biases
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-8-node-classification-with-w-b/pyg_8_node_classification_with_w_b.py b/marimo/convert/pyg-8-node-classification-with-w-b/pyg_8_node_classification_with_w_b.py
new file mode 100644
index 00000000..44966877
--- /dev/null
+++ b/marimo/convert/pyg-8-node-classification-with-w-b/pyg_8_node_classification_with_w_b.py
@@ -0,0 +1,695 @@
+# /// script
+# dependencies = ["https://data-pyg-org/whl/torch-${torch}-html", "pytorch_geometric @ git+https://github.com/pyg-team/pytorch_geometric.git", "torch-scatter", "torch-sparse", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install required packages.
+
+ # packages added via marimo's package management: torch-scatter https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-sparse https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: git+https://github.com/pyg-team/pytorch_geometric.git !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ # packages added via marimo's package management: wandb !pip install -qqq wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Setup and login to Weights & Biases
+ """)
+ return
+
+
+@app.cell
+def _():
+ enable_wandb = True
+ if enable_wandb:
+ import wandb
+ return enable_wandb, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import os
+ import pdb
+ import torch
+ import pandas
+
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return pandas, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Helper function for visualization.
+ """)
+ return
+
+
+@app.cell
+def _(pandas, wandb):
+ import matplotlib.pyplot as plt
+ from sklearn.manifold import TSNE
+
+ def visualize(h, color):
+ z = TSNE(n_components=2).fit_transform(h.detach().cpu().numpy())
+ plt.figure(figsize=(10,10))
+ plt.xticks([])
+ plt.yticks([])
+ plt.scatter(z[:, 0], z[:, 1], s=70, c=color, cmap="Set2")
+ plt.show()
+
+ def embedding_to_wandb(h, color, key="embedding"):
+ num_components = h.shape[-1]
+ df = pandas.DataFrame(data=h.detach().cpu().numpy(),
+ columns=[f"c_{i}" for i in range(num_components)])
+ df["target"] = color.detach().cpu().numpy().astype("str")
+ cols = df.columns.tolist()
+ df = df[cols[-1:] + cols[:-1]]
+ wandb.log({key: df})
+
+ return embedding_to_wandb, visualize
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Node Classification with Graph Neural Networks
+
+ [Previous: Introduction: Hands-on Graph Neural Networks](https://colab.research.google.com/drive/1h3-vJGRVloF5zStxL5I0rSy4ZUPNsjy8)
+
+ This tutorial will teach you how to apply **Graph Neural Networks (GNNs) to the task of node classification**.
+ Here, we are given the ground-truth labels of only a small subset of nodes, and want to infer the labels for all the remaining nodes (*transductive learning*).
+
+ To demonstrate, we make use of the `Cora` dataset, which is a **citation network** where nodes represent documents.
+ Each node is described by a 1433-dimensional bag-of-words feature vector.
+ Two documents are connected if there exists a citation link between them.
+ The task is to infer the category of each document (7 in total).
+
+ This dataset was first introduced by [Yang et al. (2016)](https://arxiv.org/abs/1603.08861) as one of the datasets of the `Planetoid` benchmark suite.
+ We again can make use [PyTorch Geometric](https://github.com/rusty1s/pytorch_geometric) for an easy access to this dataset via [`torch_geometric.datasets.Planetoid`](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.Planetoid):
+ """)
+ return
+
+
+@app.cell
+def _():
+ from torch_geometric.datasets import Planetoid
+ from torch_geometric.transforms import NormalizeFeatures
+
+
+
+ dataset = Planetoid(root='data/Planetoid', name='Cora', transform=NormalizeFeatures())
+
+ print()
+ print(f'Dataset: {dataset}:')
+ print('======================')
+ print(f'Number of graphs: {len(dataset)}')
+ print(f'Number of features: {dataset.num_features}')
+ print(f'Number of classes: {dataset.num_classes}')
+
+ data = dataset[0] # Get the first graph object.
+
+ print()
+ print(data)
+ print('===========================================================================================================')
+
+ # Gather some statistics about the graph.
+ print(f'Number of nodes: {data.num_nodes}')
+ print(f'Number of edges: {data.num_edges}')
+ print(f'Average node degree: {data.num_edges / data.num_nodes:.2f}')
+ print(f'Number of training nodes: {data.train_mask.sum()}')
+ print(f'Training node label rate: {int(data.train_mask.sum()) / data.num_nodes:.2f}')
+ print(f'Has isolated nodes: {data.has_isolated_nodes()}')
+ print(f'Has self-loops: {data.has_self_loops()}')
+ print(f'Is undirected: {data.is_undirected()}')
+ return data, dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Overall, this dataset is quite similar to the previously used [`KarateClub`](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.KarateClub) network.
+ We can see that the `Cora` network holds 2,708 nodes and 10,556 edges, resulting in an average node degree of 3.9.
+ For training this dataset, we are given the ground-truth categories of 140 nodes (20 for each class).
+ This results in a training node label rate of only 5%.
+
+ In contrast to `KarateClub`, this graph holds the additional attributes `val_mask` and `test_mask`, which denotes which nodes should be used for validation and testing.
+ Furthermore, we make use of **[data transformations](https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html#data-transforms) via `transform=NormalizeFeatures()`**.
+ Transforms can be used to modify your input data before inputting them into a neural network, *e.g.*, for normalization or data augmentation.
+ Here, we [row-normalize](https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.NormalizeFeatures) the bag-of-words input feature vectors.
+
+ We can further see that this network is undirected, and that there exists no isolated nodes (each document has at least one citation).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training a Multi-layer Perception Network (MLP)
+
+ In theory, we should be able to infer the category of a document solely based on its content, *i.e.* its bag-of-words feature representation, without taking any relational information into account.
+
+ Let's verify that by constructing a simple MLP that solely operates on input node features (using shared weights across all nodes):
+ """)
+ return
+
+
+@app.cell
+def _(dataset, torch):
+ from torch.nn import Linear
+ import torch.nn.functional as F
+
+ class MLP(torch.nn.Module):
+
+ def __init__(self, hidden_channels):
+ super().__init__()
+ torch.manual_seed(12345)
+ self.lin1 = Linear(dataset.num_features, hidden_channels)
+ self.lin2 = Linear(hidden_channels, dataset.num_classes)
+
+ def forward(self, x):
+ x = self.lin1(x)
+ x = x.relu()
+ x = F.dropout(x, p=0.5, training=self.training)
+ x = self.lin2(x)
+ return x
+ model = MLP(hidden_channels=16)
+ print(model)
+ return F, MLP
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ (optionally) logging the data attributes to W&B summary.
+ """)
+ return
+
+
+@app.cell
+def _(data, dataset, enable_wandb, wandb):
+ if enable_wandb:
+ wandb.init(project='node-classification')
+ summary = dict()
+ summary["data"] = dict()
+ summary["data"]["num_features"] = dataset.num_features
+ summary["data"]["num_classes"] = dataset.num_classes
+ summary["data"]["num_nodes"] = data.num_nodes
+ summary["data"]["num_edges"] = data.num_edges
+ summary["data"]["has_isolated_nodes"] = data.has_isolated_nodes()
+ summary["data"]["has_self_nodes"] = data.has_self_loops()
+ summary["data"]["is_undirected"] = data.is_undirected()
+ summary["data"]["num_training_nodes"] = data.train_mask.sum()
+ wandb.summary = summary
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Our MLP is defined by two linear layers and enhanced by [ReLU](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html?highlight=relu#torch.nn.ReLU) non-linearity and [dropout](https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html?highlight=dropout#torch.nn.Dropout).
+ Here, we first reduce the 1433-dimensional feature vector to a low-dimensional embedding (`hidden_channels=16`), while the second linear layer acts as a classifier that should map each low-dimensional node embedding to one of the 7 classes.
+
+ Let's train our simple MLP by following a similar procedure as described in [the first part of this tutorial](https://colab.research.google.com/drive/1h3-vJGRVloF5zStxL5I0rSy4ZUPNsjy8).
+ We again make use of the **cross entropy loss** and **Adam optimizer**.
+ This time, we also define a **`test` function** to evaluate how well our final model performs on the test node set (which labels have not been observed during training).
+
+ We also visualize the embeddings of the untrained model to in visually comparing the progress made by the training process below.
+
+ **NOTE**: *For W&B mode, please set up the embedding projector from the setting panel of the logged table. More information can be found here: https://docs.wandb.ai/ref/app/features/panels/weave/embedding-projector*
+ """)
+ return
+
+
+@app.cell
+def _(
+ MLP,
+ data,
+ display,
+ embedding_to_wandb,
+ enable_wandb,
+ torch,
+ visualize,
+ wandb,
+):
+ from IPython.display import Javascript
+ display(Javascript('google.colab.output.setIframeHeight(0, true, {maxHeight: 300})'))
+ model_1 = MLP(hidden_channels=16)
+ with torch.no_grad():
+ _out = model_1(data.x)
+ if enable_wandb:
+ embedding_to_wandb(_out, color=data.y, key='mlp/embedding/init')
+ else:
+ visualize(_out, data.y)
+ _criterion = torch.nn.CrossEntropyLoss()
+ _optimizer = torch.optim.Adam(model_1.parameters(), lr=0.01, weight_decay=0.0005)
+
+ def _train():
+ model_1.train()
+ _optimizer.zero_grad()
+ _out = model_1(data.x)
+ _loss = _criterion(_out[data.train_mask], data.y[data.train_mask])
+ _loss.backward()
+ _optimizer.step()
+ return _loss
+
+ def test():
+ model_1.eval()
+ _out = model_1(data.x)
+ pred = _out.argmax(dim=1)
+ test_correct = pred[data.test_mask] == data.y[data.test_mask]
+ test_acc = int(test_correct.sum()) / int(data.test_mask.sum())
+ return test_acc
+ for _epoch in range(1, 201):
+ _loss = _train()
+ if enable_wandb:
+ wandb.log({'mlp/loss': _loss})
+ print(f'Epoch: {_epoch:03d}, Loss: {_loss:.4f}')
+ return Javascript, model_1, test
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ After training the model, we can call the `test` function to see how well our model performs on unseen labels.
+ Here, we are interested in the accuracy of the model, *i.e.*, the ratio of correctly classified nodes:
+
+ We also visualize the embeddings of the output. This will give us a visual hint as to how good the model is performing, when compared to the embeddings of the geometric models defined below.
+ """)
+ return
+
+
+@app.cell
+def _(data, embedding_to_wandb, enable_wandb, model_1, test, visualize, wandb):
+ test_acc = test()
+ _out = model_1(data.x)
+ if enable_wandb:
+ embedding_to_wandb(_out, color=data.y, key='mlp/embedding/trained')
+ wandb.summary['mlp/accuracy'] = test_acc
+ wandb.log({'mlp/accuracy': test_acc})
+ else:
+ visualize(_out, data.y)
+ print(f'Test Accuracy: {test_acc:.4f}')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ As one can see, our MLP performs rather bad with only about 59% test accuracy.
+ But why does the MLP do not perform better?
+ The main reason for that is that this model suffers from heavy overfitting due to only having access to a **small amount of training nodes**, and therefore generalizes poorly to unseen node representations.
+
+ It also fails to incorporate an important bias into the model: **Cited papers are very likely related to the category of a document**.
+ That is exactly where Graph Neural Networks come into play and can help to boost the performance of our model.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training a Graph Neural Network (GNN)
+
+ We can easily convert our MLP to a GNN by swapping the `torch.nn.Linear` layers with PyG's GNN operators.
+
+ Following-up on [the first part of this tutorial](https://colab.research.google.com/drive/1h3-vJGRVloF5zStxL5I0rSy4ZUPNsjy8), we replace the linear layers by the [`GCNConv`](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GCNConv) module.
+ To recap, the **GCN layer** ([Kipf et al. (2017)](https://arxiv.org/abs/1609.02907)) is defined as
+
+ $$
+ \mathbf{x}_v^{(\ell + 1)} = \mathbf{W}^{(\ell + 1)} \sum_{w \in \mathcal{N}(v) \, \cup \, \{ v \}} \frac{1}{c_{w,v}} \cdot \mathbf{x}_w^{(\ell)}
+ $$
+
+ where $\mathbf{W}^{(\ell + 1)}$ denotes a trainable weight matrix of shape `[num_output_features, num_input_features]` and $c_{w,v}$ refers to a fixed normalization coefficient for each edge.
+ In contrast, a single `Linear` layer is defined as
+
+ $$
+ \mathbf{x}_v^{(\ell + 1)} = \mathbf{W}^{(\ell + 1)} \mathbf{x}_v^{(\ell)}
+ $$
+
+ which does not make use of neighboring node information.
+ """)
+ return
+
+
+@app.cell
+def _(F, dataset, torch):
+ from torch_geometric.nn import GCNConv
+
+ class GCN(torch.nn.Module):
+
+ def __init__(self, hidden_channels):
+ super().__init__()
+ torch.manual_seed(1234567)
+ self.conv1 = GCNConv(dataset.num_features, hidden_channels)
+ self.conv2 = GCNConv(hidden_channels, dataset.num_classes)
+
+ def forward(self, x, edge_index):
+ x = self.conv1(x, edge_index)
+ x = x.relu()
+ x = F.dropout(x, p=0.5, training=self.training)
+ x = self.conv2(x, edge_index)
+ return x
+ model_2 = GCN(hidden_channels=16)
+ print(model_2)
+ return (GCN,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's visualize the node embeddings of our **untrained** GCN network.
+ For visualization, we make use of [**TSNE**](https://scikit-learn.org/stable/modules/generated/sklearn.manifold.TSNE.html) to embed our 7-dimensional node embeddings onto a 2D plane.
+ """)
+ return
+
+
+@app.cell
+def _(GCN, data, embedding_to_wandb, enable_wandb, visualize):
+ model_3 = GCN(hidden_channels=16)
+ model_3.eval()
+ _out = model_3(data.x, data.edge_index)
+ if enable_wandb:
+ embedding_to_wandb(_out, color=data.y, key='gcn/embedding/init')
+ else:
+ visualize(_out, data.y)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We certainly can do better by training our model.
+ The training and testing procedure is once again the same, but this time we make use of the node features `x` **and** the graph connectivity `edge_index` as input to our GCN model.
+ """)
+ return
+
+
+@app.cell
+def _(GCN, Javascript, data, display, enable_wandb, torch, wandb):
+ display(Javascript('google.colab.output.setIframeHeight(0, true, {maxHeight: 300})'))
+ model_4 = GCN(hidden_channels=16)
+ if enable_wandb:
+ wandb.watch(model_4)
+ _optimizer = torch.optim.Adam(model_4.parameters(), lr=0.01, weight_decay=0.0005)
+ _criterion = torch.nn.CrossEntropyLoss()
+
+ def _train():
+ model_4.train()
+ _optimizer.zero_grad()
+ _out = model_4(data.x, data.edge_index)
+ _loss = _criterion(_out[data.train_mask], data.y[data.train_mask])
+ _loss.backward()
+ _optimizer.step()
+ return _loss
+
+ def test_1():
+ model_4.eval()
+ _out = model_4(data.x, data.edge_index)
+ pred = _out.argmax(dim=1)
+ test_correct = pred[data.test_mask] == data.y[data.test_mask]
+ test_acc = int(test_correct.sum()) / int(data.test_mask.sum())
+ return test_acc
+ for _epoch in range(1, 101):
+ _loss = _train()
+ if enable_wandb:
+ wandb.log({'gcn/loss': _loss})
+ print(f'Epoch: {_epoch:03d}, Loss: {_loss:.4f}')
+ return model_4, test_1
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ After training the model, we can check its test accuracy:
+ """)
+ return
+
+
+@app.cell
+def _(test_1):
+ test_acc_1 = test_1()
+ print(f'Test Accuracy: {test_acc_1:.4f}')
+ return (test_acc_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **There it is!**
+ By simply swapping the linear layers with GNN layers, we can reach **81.5% of test accuracy**!
+ This is in stark contrast to the 59% of test accuracy obtained by our MLP, indicating that relational information plays a crucial role in obtaining better performance.
+
+ We can also verify that once again by looking at the output embeddings of our **trained** model, which now produces a far better clustering of nodes of the same category.
+ """)
+ return
+
+
+@app.cell
+def _(
+ data,
+ embedding_to_wandb,
+ enable_wandb,
+ model_4,
+ test_acc_1,
+ visualize,
+ wandb,
+):
+ model_4.eval()
+ _out = model_4(data.x, data.edge_index)
+ if enable_wandb:
+ wandb.summary['gcn/accuracy'] = test_acc_1
+ wandb.log({'gcn/accuracy': test_acc_1})
+ embedding_to_wandb(_out, color=data.y, key='gcn/embedding/trained')
+ wandb.finish()
+ else:
+ visualize(_out, data.y)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Using W&B Sweeps
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this section, we'll look into how we can use [W&B Sweeps](https://wandb.ai/site/sweeps/) to perform a hyper-parameter search for the GCN. For this to work, it is essential for wandb to be enabled, i.e., `enable_wandb` should be set to `True`.
+ """)
+ return
+
+
+@app.cell
+def _(enable_wandb):
+ assert enable_wandb, "W&B not enabled. Please, enable W&B and restart the notebook"
+ return
+
+
+@app.cell
+def _(GCN, data, embedding_to_wandb, torch, wandb):
+ import tqdm
+
+ def agent_fn():
+ wandb.init()
+ model = GCN(hidden_channels=wandb.config.hidden_channels)
+ wandb.watch(model)
+ with torch.no_grad():
+ _out = model(data.x, data.edge_index)
+ embedding_to_wandb(_out, color=data.y, key='gcn/embedding/init')
+ _optimizer = torch.optim.Adam(model.parameters(), lr=wandb.config.lr, weight_decay=wandb.config.weight_decay)
+ _criterion = torch.nn.CrossEntropyLoss()
+
+ def _train():
+ model.train()
+ _optimizer.zero_grad()
+ _out = model(data.x, data.edge_index)
+ _loss = _criterion(_out[data.train_mask], data.y[data.train_mask]) # Clear gradients.
+ _loss.backward() # Perform a single forward pass.
+ _optimizer.step() # Compute the loss solely based on the training nodes.
+ return _loss # Derive gradients.
+ # Update parameters based on gradients.
+ def test():
+ model.eval()
+ _out = model(data.x, data.edge_index)
+ pred = _out.argmax(dim=1)
+ test_correct = pred[data.test_mask] == data.y[data.test_mask]
+ test_acc = int(test_correct.sum()) / int(data.test_mask.sum()) # Use the class with highest probability.
+ return test_acc # Check against ground-truth labels.
+ for _epoch in tqdm.tqdm(range(1, 101)): # Derive ratio of correct predictions.
+ _loss = _train()
+ wandb.log({'gcn/loss': _loss})
+ model.eval()
+ _out = model(data.x, data.edge_index)
+ test_acc = test()
+ wandb.summary['gcn/accuracy'] = test_acc
+ wandb.log({'gcn/accuracy': test_acc})
+ embedding_to_wandb(_out, color=data.y, key='gcn/embedding/trained')
+ wandb.finish()
+
+ return (agent_fn,)
+
+
+@app.cell
+def _(wandb):
+ sweep_config = {
+ "name": "gcn-sweep",
+ "method": "bayes",
+ "metric": {
+ "name": "gcn/accuracy",
+ "goal": "maximize",
+ },
+ "parameters": {
+ "hidden_channels": {
+ "values": [8, 16, 32]
+ },
+ "weight_decay": {
+ "distribution": "normal",
+ "mu": 5e-4,
+ "sigma": 1e-5,
+ },
+ "lr": {
+ "min": 1e-4,
+ "max": 1e-3
+ }
+ }
+ }
+
+ # Register the Sweep with W&B
+ sweep_id = wandb.sweep(sweep_config, project="node-classification")
+ return (sweep_id,)
+
+
+@app.cell
+def _(agent_fn, sweep_id, wandb):
+ # Run the Sweeps agent
+ wandb.agent(sweep_id, project="node-classification", function=agent_fn, count=50)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Conclusion
+
+ In this chapter, you have seen how to apply GNNs to real-world problems, and, in particular, how they can effectively be used for boosting a model's performance.
+ In the next section, we will look into how GNNs can be used for the task of graph classification.
+
+ [Next: Graph Classification with Graph Neural Networks](https://colab.research.google.com/drive/1I8a0DfQ3fI7Njc62__mVXUlcAleUclnb)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## (Optional) Exercises
+
+ 1. To achieve better model performance and to avoid overfitting, it is usually a good idea to select the best model based on an additional validation set.
+ The `Cora` dataset provides a validation node set as `data.val_mask`, but we haven't used it yet.
+ Can you modify the code to select and test the model with the highest validation performance?
+ This should bring test performance to **82% accuracy**.
+
+ 2. How does `GCN` behave when increasing the hidden feature dimensionality or the number of layers?
+ Does increasing the number of layers help at all?
+
+ 3. You can try to use different GNN layers to see how model performance changes. What happens if you swap out all `GCNConv` instances with [`GATConv`](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GATConv) layers that make use of attention? Try to write a 2-layer `GAT` model that makes use of 8 attention heads in the first layer and 1 attention head in the second layer, uses a `dropout` ratio of `0.6` inside and outside each `GATConv` call, and uses a `hidden_channels` dimensions of `8` per head.
+ """)
+ return
+
+
+@app.cell
+def _(F, data, torch):
+ from torch_geometric.nn import GATConv
+
+ class GAT(torch.nn.Module):
+
+ def __init__(self, hidden_channels, heads):
+ super().__init__()
+ torch.manual_seed(1234567)
+ self.conv1 = GATConv(...)
+ self.conv2 = GATConv(...)
+
+ def forward(self, x, edge_index):
+ x = F.dropout(x, p=0.6, training=self.training)
+ x = self.conv1(x, edge_index)
+ x = F.elu(x)
+ x = F.dropout(x, p=0.6, training=self.training)
+ x = self.conv2(x, edge_index)
+ return x
+ model_5 = GAT(hidden_channels=8, heads=8)
+ print(model_5)
+ _optimizer = torch.optim.Adam(model_5.parameters(), lr=0.005, weight_decay=0.0005)
+ _criterion = torch.nn.CrossEntropyLoss()
+
+ def _train():
+ model_5.train()
+ _optimizer.zero_grad()
+ _out = model_5(data.x, data.edge_index)
+ _loss = _criterion(_out[data.train_mask], data.y[data.train_mask])
+ _loss.backward()
+ _optimizer.step()
+ return _loss
+
+ def test_2(mask):
+ model_5.eval()
+ _out = model_5(data.x, data.edge_index)
+ pred = _out.argmax(dim=1)
+ correct = pred[mask] == data.y[mask]
+ acc = int(correct.sum()) / int(mask.sum())
+ return acc
+ for _epoch in range(1, 201):
+ _loss = _train()
+ val_acc = test_2(data.val_mask)
+ test_acc_2 = test_2(data.test_mask)
+ print(f'Epoch: {_epoch:03d}, Loss: {_loss:.4f}, Val: {val_acc:.4f}, Test: {test_acc_2:.4f}')
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-graph-classification-with-pyg-and-w-b/pyg_graph_classification_with_pyg_and_w_b.py b/marimo/convert/pyg-graph-classification-with-pyg-and-w-b/pyg_graph_classification_with_pyg_and_w_b.py
new file mode 100644
index 00000000..7af17a1e
--- /dev/null
+++ b/marimo/convert/pyg-graph-classification-with-pyg-and-w-b/pyg_graph_classification_with_pyg_and_w_b.py
@@ -0,0 +1,535 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+
+ return (os,)
+
+
+@app.cell
+def _(os):
+ import torch
+ torch_version = torch.__version__.split("+")
+ os.environ["TORCH"] = torch_version[0]
+ os.environ["CUDA"] = torch_version[1]
+ return (torch,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will start by installing the basic packages i.e. PyTorch Geometric for implementing our graph neural networks, plotly for easier visualization and W&B for tracking our experiments.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install pyg-lib torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
+ # !pip install torch-geometric
+ # !pip install wandb
+ # !pip install plotly
+ # !pip install --upgrade scipy
+ # !wget "https://gist.githubusercontent.com/mogproject/50668d3ca60188c50e6ef3f5f3ace101/raw/e11d5ac2b83fb03c0e5a9448ee3670b9dfcd5bf9/visualize.py"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now let us import `wandb` and log in to your W&B account.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now we will import all the packages that we will use as we progress through this example.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # General imports
+ import json
+ import collections
+ import numpy as np
+ import pandas as pd
+ # Data science imports
+ import matplotlib.pyplot as plt
+ import seaborn as sns
+ import plotly
+ import scipy.sparse as sp
+ import wandb
+ from torch import Tensor
+ import torch.nn.functional as F
+ # Import Weights & Biases for Experiment Tracking
+ import torch_geometric
+ from torch_geometric.nn import GCNConv
+ # Graph imports
+ from torch_geometric.utils import to_networkx
+ import networkx as nx
+ from networkx.algorithms import community
+ from tqdm.auto import trange
+
+ return F, GCNConv, json, nx, plotly, to_networkx, trange, wandb
+
+
+@app.cell
+def _():
+ #External helper packages
+ # magic command not supported in marimo; please file an issue to add support
+ # %run visualize.py
+ from visualize import GraphVisualization
+
+ return (GraphVisualization,)
+
+
+@app.cell
+def _(torch):
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Dataset and Exploratory Data Analysis
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To start logging and tracking information with W&B, we will first initialize a run. This run will reside inside a project in your W&B profile and it will store everything related to this experiment which in this case is EDA.
+
+ **Note:** Usage of Weights & Biases for training is optional, in case you don't wish to use it, simply uncheck `use_wandb`.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ use_wandb = True #@param {type:"boolean"}
+ _wandb_project = 'intro_to_pyg' #@param {type:"string"}
+ _wandb_run_name = 'upload_and_analyze_dataset' #@param {type:"string"}
+ if use_wandb:
+ wandb.init(project=_wandb_project, name=_wandb_run_name)
+ return (use_wandb,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will work on the [MUTAG](https://chrsmrrs.github.io/datasets/docs/datasets/) dataset for classifying graphs into one of two classes. Downloading and loading the graphs and labels from this dataset is supported internally by PyG.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from torch_geometric.datasets import TUDataset
+
+
+ dataset_path = "data/TUDataset"
+ dataset = TUDataset(root=dataset_path, name='MUTAG')
+
+ dataset.download()
+ return dataset, dataset_path
+
+
+@app.cell
+def _(dataset, json, use_wandb, wandb):
+ data_details = {
+ "num_node_features": dataset.num_node_features,
+ "num_edge_features": dataset.num_edge_features,
+ "num_classes": dataset.num_classes,
+ "num_node_labels": dataset.num_node_labels,
+ "num_edge_labels": dataset.num_edge_labels
+ }
+
+ if use_wandb:
+ # Log all the details about the data to W&B.
+ wandb.log(data_details) #🪄🐝
+ else:
+ print(json.dumps(data_details, sort_keys=True, indent=4))
+ return (data_details,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In the snippet below, we convert the PyG graph to a plotly figure for visualization.
+ """)
+ return
+
+
+@app.cell
+def _(GraphVisualization, dataset, nx, to_networkx):
+ def create_graph(graph):
+ g = to_networkx(graph)
+ pos = nx.spring_layout(g)
+ vis = GraphVisualization(g, pos, node_text_position='top left', node_size=20)
+ _fig = vis.create_figure()
+ return _fig
+ _fig = create_graph(dataset[0])
+ _fig.show()
+ return (create_graph,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In order to perform EDA, we will log the input graphs along with the number of nodes, edges and the ground truth labels as a W&B Table.
+ """)
+ return
+
+
+@app.cell
+def _(create_graph, dataset, plotly, use_wandb, wandb):
+ if use_wandb:
+ table = wandb.Table(columns=['Graph', 'Number of Nodes', 'Number of Edges', 'Label']) # Log exploratory visualizations for each data point to W&B
+ for graph in dataset:
+ _fig = create_graph(graph)
+ n_nodes = graph.num_nodes
+ n_edges = graph.num_edges
+ label = graph.y.item()
+ table.add_data(wandb.Html(plotly.io.to_html(_fig)), n_nodes, n_edges, label)
+ wandb.log({'data': table})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We would also like to log and track our dataset so we will add the downloaded dataset to a W&B artifact and log that as well.
+ """)
+ return
+
+
+@app.cell
+def _(data_details, dataset_path, use_wandb, wandb):
+ if use_wandb:
+ # Log the dataset to W&B as an artifact.
+ dataset_artifact = wandb.Artifact(name="MUTAG", type="dataset", metadata=data_details)
+ dataset_artifact.add_dir(dataset_path)
+ wandb.log_artifact(dataset_artifact)
+
+ # End the W&B run
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training the Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Creating Training and Testing Data
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We shuffle the dataset and split it into training and testing sets
+ """)
+ return
+
+
+@app.cell
+def _(dataset, torch):
+ torch.manual_seed(12345)
+ dataset_1 = dataset.shuffle()
+ train_dataset = dataset_1[:150]
+ test_dataset = dataset_1[150:]
+ print(f'Number of training graphs: {len(train_dataset)}')
+ print(f'Number of test graphs: {len(test_dataset)}')
+ return dataset_1, test_dataset, train_dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let us create the dataloaders to effectively batch the graph inputs. Batching graphs is implemented extremely effectively in PyG and that has been used in this example.
+ """)
+ return
+
+
+@app.cell
+def _(test_dataset, train_dataset):
+ from torch_geometric.loader import DataLoader
+
+ train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+ test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
+
+ for step, data in enumerate(train_loader):
+ print(f'Step {step + 1}:')
+ print('=======')
+ print(f'Number of graphs in the current batch: {data.num_graphs}')
+ print(data)
+ print()
+ return test_loader, train_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Implementing the Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To perform the classification, we use a very simple model with three graph convolution layers implemented in PyTorch Geometric.
+ """)
+ return
+
+
+@app.cell
+def _(F, GCNConv, dataset_1, torch):
+ from torch.nn import Linear
+ from torch_geometric.nn import global_mean_pool
+
+ class GCN(torch.nn.Module):
+
+ def __init__(self, hidden_channels):
+ super(GCN, self).__init__()
+ torch.manual_seed(12345)
+ self.conv1 = GCNConv(dataset_1.num_node_features, hidden_channels)
+ self.conv2 = GCNConv(hidden_channels, hidden_channels)
+ self.conv3 = GCNConv(hidden_channels, hidden_channels)
+ self.lin = Linear(hidden_channels, dataset_1.num_classes)
+
+ def forward(self, x, edge_index, batch):
+ x = self.conv1(x, edge_index)
+ x = x.relu()
+ x = self.conv2(x, edge_index) # 1. Obtain node embeddings
+ x = x.relu()
+ x = self.conv3(x, edge_index)
+ x = global_mean_pool(x, batch)
+ x = F.dropout(x, p=0.5, training=self.training)
+ x = self.lin(x)
+ return x
+ model = GCN(hidden_channels=64) # 2. Readout layer
+ print(model) # [batch_size, hidden_channels] # 3. Apply a final classifier
+ return (GCN,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Training
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In the following cell, start a new W&B run and using the `use_artifact` function, we tell W&B that the given artifact is being used as input to this run.
+
+ **Note:** Usage of Weights & Biases for training is optional, in case you don't wish to use it, simply uncheck `use_wandb`.
+ """)
+ return
+
+
+@app.cell
+def _(use_wandb, wandb):
+ _wandb_project = 'intro_to_pyg' #@param {type:"string"}
+ _wandb_run_name = 'upload_and_analyze_dataset' #@param {type:"string"}
+ if use_wandb:
+ # Initialize W&B run for training
+ wandb.init(project='intro_to_pyg')
+ wandb.use_artifact('manan-goel/intro_to_pyg/MUTAG:v0')
+ return
+
+
+@app.cell
+def _(GCN, display, torch, train_loader):
+ from IPython.display import Javascript
+ display(Javascript('google.colab.output.setIframeHeight(0, true, {maxHeight: 300})'))
+ model_1 = GCN(hidden_channels=64)
+ optimizer = torch.optim.Adam(model_1.parameters(), lr=0.001)
+ criterion = torch.nn.CrossEntropyLoss()
+
+ def train():
+ model_1.train()
+ for data in train_loader:
+ out = model_1(data.x, data.edge_index, data.batch)
+ loss = criterion(out, data.y)
+ loss.backward() # Iterate in batches over the training dataset.
+ optimizer.step() # Perform a single forward pass.
+ optimizer.zero_grad() # Compute the loss. # Derive gradients. # Update parameters based on gradients. # Clear gradients.
+
+ return criterion, model_1, train
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In the function in the cell below, we implement a function to calculate accuracy and loss on the test set but along with that we also create a W&B table which can be used to see the graph along with the output from the model and the ground truth label.
+ 
+ """)
+ return
+
+
+@app.cell
+def _(create_graph, criterion, model_1, plotly, use_wandb, wandb):
+ def test(loader, create_table=False):
+ model_1.eval()
+ table = wandb.Table(columns=['graph', 'ground truth', 'prediction']) if use_wandb else None
+ correct = 0
+ loss_ = 0
+ for data in loader: # Iterate in batches over the training/test dataset.
+ out = model_1(data.x, data.edge_index, data.batch)
+ loss = criterion(out, data.y)
+ loss_ = loss_ + loss.item()
+ pred = out.argmax(dim=1) # Use the class with highest probability.
+ if create_table and use_wandb:
+ table.add_data(wandb.Html(plotly.io.to_html(create_graph(data))), data.y.item(), pred.item())
+ correct = correct + int((pred == data.y).sum())
+ return (correct / len(loader.dataset), loss_ / len(loader.dataset), table) # Check against ground-truth labels. # Derive ratio of correct predictions.
+
+ return (test,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The great thing about PyG is that all the code is very similar to PyTorch and you can use the same training logic which is what has been done in the following cell. We calculate the accuracy and loss on the training and validation sets and long them to our W&B dashboard.
+
+ We also log the table that was mentioned above to the run dashboard for debugging.
+ """)
+ return
+
+
+@app.cell
+def _(
+ model_1,
+ test,
+ test_loader,
+ torch,
+ train,
+ train_loader,
+ trange,
+ use_wandb,
+ wandb,
+):
+ for epoch in trange(1, 171):
+ train()
+ train_acc, train_loss, _ = test(train_loader)
+ test_acc, test_loss, test_table = test(test_loader, create_table=True)
+ if use_wandb:
+ wandb.log({'train/loss': train_loss, 'train/acc': train_acc, 'test/acc': test_acc, 'test/loss': test_loss, 'test/table': test_table}) # Log metrics to W&B
+ torch.save(model_1, 'graph_classification_model.pt')
+ if use_wandb:
+ artifact = wandb.Artifact(name='graph_classification_model', type='model')
+ artifact.add_file('graph_classification_model.pt')
+ wandb.log_artifact(artifact)
+ if use_wandb:
+ # Finish the W&B run
+ wandb.finish() # Log model checkpoint as an artifact to W&B
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This is what the project dashboard looks like and you can interact with it i.e. make the panels larger or smaller and more!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📚 Resources
+
+ * [Node Classification (with W&B)](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/pyg/8_Node_Classification_(with_W%26B).ipynb) contains a few tips for taking most advantage of W&B with PyG
+ * [Point Cloud Classification using PyTorch Geometric](https://wandb.ai/geekyrakshit/pyg-point-cloud/reports/Point-Cloud-Classification-using-PyTorch-Geometric--VmlldzozMTExMTE3)
+ * [Recommending Amazon Products using Graph Neural Networks in PyTorch Geometric](https://wandb.ai/manan-goel/gnn-recommender/reports/Recommending-Amazon-Products-using-Graph-Neural-Networks-in-PyTorch-Geometric--VmlldzozMTA3MzYw)
+ """)
+ 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, please contact support@wandb.com
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-point-cloud-segmentation-00-eda/pyg_point_cloud_segmentation_00_eda.py b/marimo/convert/pyg-point-cloud-segmentation-00-eda/pyg_point_cloud_segmentation_00_eda.py
new file mode 100644
index 00000000..51c8086e
--- /dev/null
+++ b/marimo/convert/pyg-point-cloud-segmentation-00-eda/pyg_point_cloud_segmentation_00_eda.py
@@ -0,0 +1,223 @@
+# /// script
+# dependencies = ["https://data-pyg-org/whl/torch-${torch}-html", "pytorch_geometric @ git+https://github.com/pyg-team/pytorch_geometric.git", "torch-cluster", "torch-scatter", "torch-sparse", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Explore ShapeNet Dataset using PyTorch Geometric and Weights & Biases 🪄🐝
+
+
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/pyg/point-cloud-segmentation/colabs/pyg/point-cloud-segmentation/00_eda.ipynb)
+
+ This notebook demonstrates how to fetch and load the ShapeNet dataset for point cloud classification and segmentation tasks using [PyTorch Geometric](https://www.pyg.org/) and explore the dataset using [Weights & Biases](https://wandb.ai/site).
+
+ If you wish to know how to train and evaluate the model on the ShapeNetCore dataset using Weights & Biases, you can check out the following notebooks:
+
+ **Train DGCNN:** [](https://colab.research.google.com/github/wandb/examples/blob/pyg/point-cloud-segmentation/colabs/pyg/point-cloud-segmentation/01_dgcnn_train.ipynb)
+
+ **Evaluate DGCNN:** [](https://colab.research.google.com/github/wandb/examples/blob/pyg/point-cloud-segmentation/colabs/pyg/point-cloud-segmentation/02_dgcnn_evaluate.ipynb)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install Required Packages
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import torch
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return (os,)
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: torch-scatter https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-sparse https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-cluster https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: git+https://github.com/pyg-team/pytorch_geometric.git !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Import Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import numpy as np
+ from tqdm.auto import tqdm
+ import torch.nn.functional as F
+ from torch_scatter import scatter
+ from torchmetrics.functional import jaccard_index
+ import torch_geometric.transforms as T
+ from torch_geometric.datasets import ShapeNet
+ from torch_geometric.loader import DataLoader
+ from torch_geometric.nn import MLP, DynamicEdgeConv
+
+ return ShapeNet, T, np, tqdm, wandb
+
+
+@app.cell
+def _(ShapeNet, T, os, wandb):
+ wandb_project = "pyg-point-cloud" #@param {"type": "string"}
+ wandb_run_name = "evaluate-dgcnn" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, name=wandb_run_name, job_type="visualize")
+
+ config = wandb.config
+ config.category = 'Airplane' #@param ["Bag", "Cap", "Car", "Chair", "Earphone", "Guitar", "Knife", "Lamp", "Laptop", "Motorbike", "Mug", "Pistol", "Rocket", "Skateboard", "Table"] {type:"raw"}
+
+ path = os.path.join('ShapeNet', config.category)
+ pre_transform = T.NormalizeScale()
+ train_dataset = ShapeNet(path, config.category, split='trainval', pre_transform=pre_transform)
+ test_dataset = ShapeNet(path, config.category, split='test', pre_transform=pre_transform)
+ return config, test_dataset, train_dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Visualize Train-Val Dataset
+ """)
+ return
+
+
+@app.cell
+def _(tqdm, train_dataset):
+ segmentation_class_frequency = {}
+ for _idx in tqdm(range(len(train_dataset))):
+ _pc_viz = train_dataset[_idx].pos.numpy().tolist()
+ _segmentation_label = train_dataset[_idx].y.numpy().tolist()
+ for _label in set(_segmentation_label):
+ segmentation_class_frequency[_label] = _segmentation_label.count(_label)
+ class_offset = min(list(segmentation_class_frequency.keys()))
+ return class_offset, segmentation_class_frequency
+
+
+@app.cell
+def _(
+ class_offset,
+ config,
+ np,
+ segmentation_class_frequency,
+ tqdm,
+ train_dataset,
+ wandb,
+):
+ table = wandb.Table(columns=['Point-Cloud', 'Segmentation-Class-Frequency', 'Model-Category', 'Split'])
+ for _idx in tqdm(range(len(train_dataset))):
+ _pc_viz = train_dataset[_idx].pos.numpy().tolist()
+ _segmentation_label = train_dataset[_idx].y.numpy().tolist()
+ _frequency_dict = {key: 0 for key in segmentation_class_frequency.keys()}
+ for _label in set(_segmentation_label):
+ _frequency_dict[_label] = _segmentation_label.count(_label)
+ for _j in range(len(_pc_viz)):
+ _pc_viz[_j] = _pc_viz[_j] + [_segmentation_label[_j] + 1 - class_offset]
+ table.add_data(wandb.Object3D(np.array(_pc_viz)), _frequency_dict, config.category, 'Train-Val')
+ return (table,)
+
+
+@app.cell
+def _(config, segmentation_class_frequency, wandb):
+ _data = [[key, segmentation_class_frequency[key]] for key in segmentation_class_frequency.keys()]
+ wandb.log({f'ShapeNet Class-Frequency Distribution for {config.category} Train-Val Set': wandb.plot.bar(wandb.Table(data=_data, columns=['Class', 'Frequency']), 'Class', 'Frequency', title=f'ShapeNet Class-Frequency Distribution for {config.category} Train-Val Set')})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Visualize Test Dataset
+ """)
+ return
+
+
+@app.cell
+def _(test_dataset, tqdm, train_dataset):
+ segmentation_class_frequency_1 = {}
+ for _idx in tqdm(range(len(test_dataset))):
+ _pc_viz = train_dataset[_idx].pos.numpy().tolist()
+ _segmentation_label = train_dataset[_idx].y.numpy().tolist()
+ for _label in set(_segmentation_label):
+ segmentation_class_frequency_1[_label] = _segmentation_label.count(_label)
+ return (segmentation_class_frequency_1,)
+
+
+@app.cell
+def _(
+ class_offset,
+ config,
+ np,
+ segmentation_class_frequency_1,
+ table,
+ test_dataset,
+ tqdm,
+ train_dataset,
+ wandb,
+):
+ for _idx in tqdm(range(len(test_dataset))):
+ _pc_viz = train_dataset[_idx].pos.numpy().tolist()
+ _segmentation_label = train_dataset[_idx].y.numpy().tolist()
+ _frequency_dict = {key: 0 for key in segmentation_class_frequency_1.keys()}
+ for _label in set(_segmentation_label):
+ _frequency_dict[_label] = _segmentation_label.count(_label)
+ for _j in range(len(_pc_viz)):
+ _pc_viz[_j] = _pc_viz[_j] + [_segmentation_label[_j] + 1 - class_offset]
+ table.add_data(wandb.Object3D(np.array(_pc_viz)), _frequency_dict, config.category, 'Test')
+ wandb.log({'ShapeNet-Dataset': table})
+ return
+
+
+@app.cell
+def _(segmentation_class_frequency_1, wandb):
+ _data = [[key, segmentation_class_frequency_1[key]] for key in segmentation_class_frequency_1.keys()]
+ wandb.log({f'ShapeNet Class-Frequency Distribution for Test Set': wandb.plot.bar(wandb.Table(data=_data, columns=['Class', 'Frequency']), 'Class', 'Frequency', title=f'ShapeNet Class-Frequency Distribution for Test Set')})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-point-cloud-segmentation-01-dgcnn-train/pyg_point_cloud_segmentation_01_dgcnn_train.py b/marimo/convert/pyg-point-cloud-segmentation-01-dgcnn-train/pyg_point_cloud_segmentation_01_dgcnn_train.py
new file mode 100644
index 00000000..60208afe
--- /dev/null
+++ b/marimo/convert/pyg-point-cloud-segmentation-01-dgcnn-train/pyg_point_cloud_segmentation_01_dgcnn_train.py
@@ -0,0 +1,583 @@
+# /// script
+# dependencies = ["https://data-pyg-org/whl/torch-${torch}-html", "pytorch_geometric @ git+https://github.com/pyg-team/pytorch_geometric.git", "torch-cluster", "torch-scatter", "torch-sparse", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Train DGCNN Model using PyTorch Geometric and Weights & Biases 🪄🐝
+
+
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/pyg/point-cloud-segmentation/colabs/pyg/point-cloud-segmentation/01_dgcnn_train.ipynb)
+
+ This notebook demonstrates an implementation of the [Dynamic Graph CNN](https://arxiv.org/pdf/1801.07829.pdf) for point cloud segmnetation implemented using [PyTorch Geometric](https://www.pyg.org/) and experiment tracked and visualized using [Weights & Biases](https://wandb.ai/site). The code here is inspired by [this](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/dgcnn_segmentation.py) original implementation.
+
+ If you wish to know how to evaluate the model on the ShapeNetCore dataset using Weights & Biases, you can check out the following notebook:
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/pyg/point-cloud-segmentation/colabs/pyg/point-cloud-segmentation/02_dgcnn_evaluate.ipynb)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install Required Packages
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import torch
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return os, torch
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: torch-scatter https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-sparse https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-cluster https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: git+https://github.com/pyg-team/pytorch_geometric.git !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Import Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import random
+ import numpy as np
+ from tqdm.auto import tqdm
+ import torch.nn.functional as F
+ from torch_scatter import scatter
+ from torchmetrics.functional import jaccard_index
+ import torch_geometric.transforms as T
+ from torch_geometric.datasets import ShapeNet
+ from torch_geometric.loader import DataLoader
+ from torch_geometric.nn import MLP, DynamicEdgeConv
+
+ return (
+ DataLoader,
+ DynamicEdgeConv,
+ F,
+ MLP,
+ ShapeNet,
+ T,
+ jaccard_index,
+ np,
+ random,
+ scatter,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Initialize Weights & Biases
+
+ We need to call [`wandb.init()`](https://docs.wandb.ai/ref/python/init) once at the beginning of our program to initialize a new job. This creates a new run in W&B and launches a background process to sync data.
+ """)
+ return
+
+
+@app.cell
+def _(random, torch, wandb):
+ wandb_project = "pyg-point-cloud" #@param {"type": "string"}
+ wandb_run_name = "train-dgcnn" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, name=wandb_run_name, job_type="train")
+
+ config = wandb.config
+
+ config.seed = 42
+ config.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+ random.seed(config.seed)
+ torch.manual_seed(config.seed)
+ device = torch.device(config.device)
+
+ config.category = 'Airplane' #@param ["Bag", "Cap", "Car", "Chair", "Earphone", "Guitar", "Knife", "Lamp", "Laptop", "Motorbike", "Mug", "Pistol", "Rocket", "Skateboard", "Table"] {type:"raw"}
+ config.random_jitter_translation = 1e-2
+ config.random_rotation_interval_x = 15
+ config.random_rotation_interval_y = 15
+ config.random_rotation_interval_z = 15
+ config.validation_split = 0.2
+ config.batch_size = 16
+ config.num_workers = 6
+
+ config.num_nearest_neighbours = 30
+ config.aggregation_operator = "max"
+ config.dropout = 0.5
+ config.initial_lr = 1e-3
+ config.lr_scheduler_step_size = 5
+ config.gamma = 0.8
+
+ config.epochs = 1
+ return config, device
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Load ShapeNet Dataset using PyTorch Geometric
+
+ We now load, preprocess and batch the ModelNet dataset for training, validation/testing and visualization.
+ """)
+ return
+
+
+@app.cell
+def _(T, config):
+ transform = T.Compose([
+ T.RandomJitter(config.random_jitter_translation),
+ T.RandomRotate(config.random_rotation_interval_x, axis=0),
+ T.RandomRotate(config.random_rotation_interval_y, axis=1),
+ T.RandomRotate(config.random_rotation_interval_z, axis=2)
+ ])
+ pre_transform = T.NormalizeScale()
+ return pre_transform, transform
+
+
+@app.cell
+def _(ShapeNet, config, os, pre_transform, transform):
+ dataset_path = os.path.join('ShapeNet', config.category)
+
+ train_val_dataset = ShapeNet(
+ dataset_path, config.category, split='trainval',
+ transform=transform, pre_transform=pre_transform
+ )
+ return (train_val_dataset,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we need to offset the segmentation labels
+ """)
+ return
+
+
+@app.cell
+def _(tqdm, train_val_dataset):
+ segmentation_class_frequency = {}
+ for idx in tqdm(range(len(train_val_dataset))):
+ pc_viz = train_val_dataset[idx].pos.numpy().tolist()
+ segmentation_label = train_val_dataset[idx].y.numpy().tolist()
+ for label in set(segmentation_label):
+ segmentation_class_frequency[label] = segmentation_label.count(label)
+ class_offset = min(list(segmentation_class_frequency.keys()))
+ print("Class Offset:", class_offset)
+
+ for idx in range(len(train_val_dataset)):
+ train_val_dataset[idx].y -= class_offset
+ return (segmentation_class_frequency,)
+
+
+@app.cell
+def _(config, train_val_dataset):
+ num_train_examples = int((1 - config.validation_split) * len(train_val_dataset))
+ train_dataset = train_val_dataset[:num_train_examples]
+ val_dataset = train_val_dataset[num_train_examples:]
+ return train_dataset, val_dataset
+
+
+@app.cell
+def _(DataLoader, config, train_dataset, val_dataset):
+ train_loader = DataLoader(
+ train_dataset, batch_size=config.batch_size,
+ shuffle=True, num_workers=config.num_workers
+ )
+ val_loader = DataLoader(
+ val_dataset, batch_size=config.batch_size,
+ shuffle=False, num_workers=config.num_workers
+ )
+ visualization_loader = DataLoader(
+ val_dataset[:10], batch_size=1,
+ shuffle=False, num_workers=config.num_workers
+ )
+ return train_loader, val_loader, visualization_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Implementing the DGCNN Model using PyTorch Geometric
+ """)
+ return
+
+
+@app.cell
+def _(DynamicEdgeConv, F, MLP, torch):
+ class DGCNN(torch.nn.Module):
+ def __init__(self, out_channels, k=30, aggr='max'):
+ super().__init__()
+
+ self.conv1 = DynamicEdgeConv(MLP([2 * 6, 64, 64]), k, aggr)
+ self.conv2 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr)
+ self.conv3 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr)
+
+ self.mlp = MLP(
+ [3 * 64, 1024, 256, 128, out_channels],
+ dropout=0.5, norm=None
+ )
+
+ def forward(self, data):
+ x, pos, batch = data.x, data.pos, data.batch
+ x0 = torch.cat([x, pos], dim=-1)
+
+ x1 = self.conv1(x0, batch)
+ x2 = self.conv2(x1, batch)
+ x3 = self.conv3(x2, batch)
+
+ out = self.mlp(torch.cat([x1, x2, x3], dim=1))
+ return F.log_softmax(out, dim=1)
+
+ return (DGCNN,)
+
+
+@app.cell
+def _(DGCNN, config, device, torch, train_dataset):
+ config.num_classes = train_dataset.num_classes
+
+ model = DGCNN(
+ out_channels=train_dataset.num_classes,
+ k=config.num_nearest_neighbours,
+ aggr=config.aggregation_operator
+ ).to(device)
+ optimizer = torch.optim.Adam(model.parameters(), lr=config.initial_lr)
+ scheduler = torch.optim.lr_scheduler.StepLR(
+ optimizer, step_size=config.lr_scheduler_step_size, gamma=config.gamma
+ )
+ return model, optimizer, scheduler
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training DGCNN and Logging Metrics on Weights & Biases
+ """)
+ return
+
+
+@app.cell
+def _(
+ F,
+ ShapeNet,
+ config,
+ device,
+ jaccard_index,
+ model,
+ optimizer,
+ scatter,
+ torch,
+ tqdm,
+ train_loader,
+):
+ def train_step(epoch):
+ model.train()
+
+ ious, categories = [], []
+ total_loss = correct_nodes = total_nodes = 0
+ y_map = torch.empty(
+ train_loader.dataset.num_classes, device=device
+ ).long()
+ num_train_examples = len(train_loader)
+
+ progress_bar = tqdm(
+ train_loader, desc=f"Training Epoch {epoch}/{config.epochs}"
+ )
+
+ for data in progress_bar:
+ data = data.to(device)
+
+ optimizer.zero_grad()
+ outs = model(data)
+ loss = F.nll_loss(outs, data.y)
+ loss.backward()
+ optimizer.step()
+
+ total_loss += loss.item()
+
+ correct_nodes += outs.argmax(dim=1).eq(data.y).sum().item()
+ total_nodes += data.num_nodes
+
+ sizes = (data.ptr[1:] - data.ptr[:-1]).tolist()
+ for out, y, category in zip(outs.split(sizes), data.y.split(sizes),
+ data.category.tolist()):
+ category = list(ShapeNet.seg_classes.keys())[category]
+ part = ShapeNet.seg_classes[category]
+ part = torch.tensor(part, device=device)
+
+ y_map[part] = torch.arange(part.size(0), device=device)
+
+ iou = jaccard_index(
+ out[:, part].argmax(dim=-1), y_map[y],
+ task="multiclass", num_classes=part.size(0)
+ )
+ ious.append(iou)
+
+ categories.append(data.category)
+
+ iou = torch.tensor(ious, device=device)
+ category = torch.cat(categories, dim=0)
+ mean_iou = float(scatter(iou, category, reduce='mean').mean())
+
+ return {
+ "Train/Loss": total_loss / num_train_examples,
+ "Train/Accuracy": correct_nodes / total_nodes,
+ "Train/IoU": mean_iou
+ }
+
+ return (train_step,)
+
+
+@app.cell
+def _(
+ F,
+ ShapeNet,
+ config,
+ device,
+ jaccard_index,
+ model,
+ scatter,
+ torch,
+ tqdm,
+ val_loader,
+):
+ @torch.no_grad()
+ def val_step(epoch):
+ model.eval()
+
+ ious, categories = [], []
+ total_loss = correct_nodes = total_nodes = 0
+ y_map = torch.empty(
+ val_loader.dataset.num_classes, device=device
+ ).long()
+ num_val_examples = len(val_loader)
+
+ progress_bar = tqdm(
+ val_loader, desc=f"Validating Epoch {epoch}/{config.epochs}"
+ )
+
+ for data in progress_bar:
+ data = data.to(device)
+ outs = model(data)
+
+ loss = F.nll_loss(outs, data.y)
+ total_loss += loss.item()
+
+ correct_nodes += outs.argmax(dim=1).eq(data.y).sum().item()
+ total_nodes += data.num_nodes
+
+ sizes = (data.ptr[1:] - data.ptr[:-1]).tolist()
+ for out, y, category in zip(outs.split(sizes), data.y.split(sizes),
+ data.category.tolist()):
+ category = list(ShapeNet.seg_classes.keys())[category]
+ part = ShapeNet.seg_classes[category]
+ part = torch.tensor(part, device=device)
+
+ y_map[part] = torch.arange(part.size(0), device=device)
+
+ iou = jaccard_index(
+ out[:, part].argmax(dim=-1), y_map[y],
+ task="multiclass", num_classes=part.size(0)
+ )
+ ious.append(iou)
+
+ categories.append(data.category)
+
+ iou = torch.tensor(ious, device=device)
+ category = torch.cat(categories, dim=0)
+ mean_iou = float(scatter(iou, category, reduce='mean').mean())
+
+ return {
+ "Validation/Loss": total_loss / num_val_examples,
+ "Validation/Accuracy": correct_nodes / total_nodes,
+ "Validation/IoU": mean_iou
+ }
+
+ return (val_step,)
+
+
+@app.cell
+def _(
+ ShapeNet,
+ device,
+ jaccard_index,
+ model,
+ np,
+ scatter,
+ segmentation_class_frequency,
+ torch,
+ tqdm,
+ visualization_loader,
+ wandb,
+):
+ @torch.no_grad()
+ def visualization_step(epoch, table):
+ model.eval()
+ for data in tqdm(visualization_loader):
+ data = data.to(device)
+ outs = model(data)
+
+ predicted_labels = outs.argmax(dim=1)
+ accuracy = predicted_labels.eq(data.y).sum().item() / data.num_nodes
+
+ sizes = (data.ptr[1:] - data.ptr[:-1]).tolist()
+ ious, categories = [], []
+ y_map = torch.empty(
+ visualization_loader.dataset.num_classes, device=device
+ ).long()
+ for out, y, category in zip(
+ outs.split(sizes), data.y.split(sizes), data.category.tolist()
+ ):
+ category = list(ShapeNet.seg_classes.keys())[category]
+ part = ShapeNet.seg_classes[category]
+ part = torch.tensor(part, device=device)
+ y_map[part] = torch.arange(part.size(0), device=device)
+ iou = jaccard_index(
+ out[:, part].argmax(dim=-1), y_map[y],
+ task="multiclass", num_classes=part.size(0)
+ )
+ ious.append(iou)
+ categories.append(data.category)
+ iou = torch.tensor(ious, device=device)
+ category = torch.cat(categories, dim=0)
+ mean_iou = float(scatter(iou, category, reduce='mean').mean())
+
+ gt_pc_viz = data.pos.cpu().numpy().tolist()
+ segmentation_label = data.y.cpu().numpy().tolist()
+ frequency_dict = {key: 0 for key in segmentation_class_frequency.keys()}
+ for label in set(segmentation_label):
+ frequency_dict[label] = segmentation_label.count(label)
+ for j in range(len(gt_pc_viz)):
+ # gt_pc_viz[j] += [segmentation_label[j] + 1 - class_offset]
+ gt_pc_viz[j] += [segmentation_label[j] + 1]
+
+ predicted_pc_viz = data.pos.cpu().numpy().tolist()
+ segmentation_label = data.y.cpu().numpy().tolist()
+ frequency_dict = {key: 0 for key in segmentation_class_frequency.keys()}
+ for label in set(segmentation_label):
+ frequency_dict[label] = segmentation_label.count(label)
+ for j in range(len(predicted_pc_viz)):
+ # predicted_pc_viz[j] += [segmentation_label[j] + 1 - class_offset]
+ predicted_pc_viz[j] += [segmentation_label[j] + 1]
+
+ table.add_data(
+ epoch, wandb.Object3D(np.array(gt_pc_viz)),
+ wandb.Object3D(np.array(predicted_pc_viz)),
+ accuracy, mean_iou
+ )
+
+ return table
+
+ return (visualization_step,)
+
+
+@app.cell
+def _(model, optimizer, torch, wandb):
+ def save_checkpoint(epoch):
+ """Save model checkpoints as Weights & Biases artifacts"""
+ torch.save({
+ 'epoch': epoch,
+ 'model_state_dict': model.state_dict(),
+ 'optimizer_state_dict': optimizer.state_dict()
+ }, "checkpoint.pt")
+
+ artifact_name = wandb.util.make_artifact_name_safe(
+ f"{wandb.run.name}-{wandb.run.id}-checkpoint"
+ )
+
+ checkpoint_artifact = wandb.Artifact(artifact_name, type="checkpoint")
+ checkpoint_artifact.add_file("checkpoint.pt")
+ wandb.log_artifact(
+ checkpoint_artifact, aliases=["latest", f"epoch-{epoch}"]
+ )
+
+ return (save_checkpoint,)
+
+
+@app.cell
+def _(
+ config,
+ save_checkpoint,
+ scheduler,
+ train_step,
+ val_step,
+ visualization_step,
+ wandb,
+):
+ table = wandb.Table(columns=["Epoch", "Ground-Truth", "Prediction", "Accuracy", "IoU"])
+
+ for epoch in range(1, config.epochs + 1):
+ train_metrics = train_step(epoch)
+ val_metrics = val_step(epoch)
+
+ metrics = {**train_metrics, **val_metrics}
+ metrics["learning_rate"] = scheduler.get_last_lr()[-1]
+ wandb.log(metrics)
+
+ table = visualization_step(epoch, table)
+
+ scheduler.step()
+ save_checkpoint(epoch)
+
+ wandb.log({"Evaluation": table})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, you can check out the following notebook to learn how to evaluate the model on the ShapeNetCore dataset using Weights & Biases, you can check out the following notebook:
+
+ []()
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-point-cloud-segmentation-02-dgcnn-evaluate/pyg_point_cloud_segmentation_02_dgcnn_evaluate.py b/marimo/convert/pyg-point-cloud-segmentation-02-dgcnn-evaluate/pyg_point_cloud_segmentation_02_dgcnn_evaluate.py
new file mode 100644
index 00000000..1702c842
--- /dev/null
+++ b/marimo/convert/pyg-point-cloud-segmentation-02-dgcnn-evaluate/pyg_point_cloud_segmentation_02_dgcnn_evaluate.py
@@ -0,0 +1,412 @@
+# /// script
+# dependencies = ["https://data-pyg-org/whl/torch-${torch}-html", "pytorch_geometric @ git+https://github.com/pyg-team/pytorch_geometric.git", "torch-cluster", "torch-scatter", "torch-sparse", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Evaluate DGCNN Model Weights & Biases 🪄🐝
+
+
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/pyg/point-cloud-segmentation/colabs/pyg/point-cloud-segmentation/02_dgcnn_evaluate.ipynb)
+
+ This notebook demonstrates the evaluation of [Dynamic Graph CNN](https://arxiv.org/pdf/1801.07829.pdf) for point cloud segmnetation. You can check the following notebook for referring to the training code:
+
+ [](https://colab.research.google.com/github/wandb/examples/blob/pyg/point-cloud-segmentation/colabs/pyg/point-cloud-segmentation/01_dgcnn_train.ipynb)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install Required Packages
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import torch
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return os, torch
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: torch-scatter https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-sparse https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-cluster https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: git+https://github.com/pyg-team/pytorch_geometric.git !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Import Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import random
+ import numpy as np
+ from tqdm.auto import tqdm
+ import torch.nn.functional as F
+ from torch_scatter import scatter
+ from torchmetrics.functional import jaccard_index
+ import torch_geometric.transforms as T
+ from torch_geometric.datasets import ShapeNet
+ from torch_geometric.loader import DataLoader
+ from torch_geometric.nn import MLP, DynamicEdgeConv
+
+ return (
+ DataLoader,
+ DynamicEdgeConv,
+ F,
+ MLP,
+ ShapeNet,
+ T,
+ jaccard_index,
+ np,
+ random,
+ scatter,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Initialize Weights & Biases
+
+ We need to call [`wandb.init()`](https://docs.wandb.ai/ref/python/init) once at the beginning of our program to initialize a new job. This creates a new run in W&B and launches a background process to sync data.
+ """)
+ return
+
+
+@app.cell
+def _(random, torch, wandb):
+ wandb_project = "pyg-point-cloud" #@param {"type": "string"}
+ wandb_run_name = "evaluate-dgcnn" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, name=wandb_run_name, job_type="evaluate")
+
+ config = wandb.config
+
+ config.seed = 42
+ config.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+ random.seed(config.seed)
+ torch.manual_seed(config.seed)
+ device = torch.device(config.device)
+
+ config.category = 'Airplane' #@param ["Bag", "Cap", "Car", "Chair", "Earphone", "Guitar", "Knife", "Lamp", "Laptop", "Motorbike", "Mug", "Pistol", "Rocket", "Skateboard", "Table"] {type:"raw"}
+ config.random_jitter_translation = 1e-2
+ config.random_rotation_interval_x = 15
+ config.random_rotation_interval_y = 15
+ config.random_rotation_interval_z = 15
+ config.batch_size = 1
+ config.num_workers = 6
+
+ config.num_nearest_neighbours = 30
+ config.aggregation_operator = "max"
+ config.dropout = 0.5
+ config.initial_lr = 1e-3
+ config.lr_scheduler_step_size = 20
+ config.gamma = 0.8
+
+ config.artifact_address = 'wandb/point-cloud-segmentation/dgcnn-3n97rfrv-checkpoint:v29'
+ config.epochs = 30
+ return config, device
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Load ShapeNet Dataset using PyTorch Geometric
+
+ We now load, preprocess and batch the ModelNet dataset for training, validation/testing and visualization.
+ """)
+ return
+
+
+@app.cell
+def _(T, config):
+ transform = T.Compose([
+ T.RandomJitter(config.random_jitter_translation),
+ T.RandomRotate(config.random_rotation_interval_x, axis=0),
+ T.RandomRotate(config.random_rotation_interval_y, axis=1),
+ T.RandomRotate(config.random_rotation_interval_z, axis=2)
+ ])
+ pre_transform = T.NormalizeScale()
+ return pre_transform, transform
+
+
+@app.cell
+def _(ShapeNet, config, os, pre_transform, transform):
+ dataset_path = os.path.join('ShapeNet', config.category)
+
+ train_dataset = ShapeNet(
+ dataset_path, config.category, split='trainval',
+ transform=transform, pre_transform=pre_transform
+ )
+ test_dataset = ShapeNet(
+ dataset_path, config.category, split='test',
+ pre_transform=pre_transform
+ )
+ return test_dataset, train_dataset
+
+
+@app.cell
+def _(test_dataset, tqdm, train_dataset):
+ segmentation_class_frequency = {}
+ for _idx in tqdm(range(len(train_dataset))):
+ pc_viz = train_dataset[_idx].pos.numpy().tolist()
+ segmentation_label = train_dataset[_idx].y.numpy().tolist()
+ for label in set(segmentation_label):
+ segmentation_class_frequency[label] = segmentation_label.count(label)
+ for _idx in tqdm(range(len(test_dataset))):
+ pc_viz = train_dataset[_idx].pos.numpy().tolist()
+ segmentation_label = train_dataset[_idx].y.numpy().tolist()
+ for label in set(segmentation_label):
+ segmentation_class_frequency[label] = segmentation_label.count(label)
+ class_offset = min(list(segmentation_class_frequency.keys()))
+ class_offset
+ return class_offset, segmentation_class_frequency
+
+
+@app.cell
+def _(class_offset, test_dataset, tqdm, train_dataset):
+ for _idx in tqdm(range(len(train_dataset))):
+ train_dataset[_idx].y -= class_offset
+ for _idx in tqdm(range(len(test_dataset))):
+ test_dataset[_idx].y -= class_offset
+ return
+
+
+@app.cell
+def _(DataLoader, config, test_dataset, train_dataset):
+ train_loader = DataLoader(
+ train_dataset, batch_size=config.batch_size,
+ shuffle=True, num_workers=config.num_workers
+ )
+ test_loader = DataLoader(
+ test_dataset, batch_size=config.batch_size,
+ shuffle=False, num_workers=config.num_workers
+ )
+ return test_loader, train_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Load Checkpoint
+ """)
+ return
+
+
+@app.cell
+def _(DynamicEdgeConv, F, MLP, torch):
+ class DGCNN(torch.nn.Module):
+ def __init__(self, out_channels, k=30, aggr='max'):
+ super().__init__()
+
+ self.conv1 = DynamicEdgeConv(MLP([2 * 6, 64, 64]), k, aggr)
+ self.conv2 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr)
+ self.conv3 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr)
+
+ self.mlp = MLP(
+ [3 * 64, 1024, 256, 128, out_channels],
+ dropout=0.5, norm=None
+ )
+
+ def forward(self, data):
+ x, pos, batch = data.x, data.pos, data.batch
+ x0 = torch.cat([x, pos], dim=-1)
+
+ x1 = self.conv1(x0, batch)
+ x2 = self.conv2(x1, batch)
+ x3 = self.conv3(x2, batch)
+
+ out = self.mlp(torch.cat([x1, x2, x3], dim=1))
+ return F.log_softmax(out, dim=1)
+
+ return (DGCNN,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Since we saved the checkpoints as artifacts on our Weights & Biases workspace, we can now fetch and load them.
+ """)
+ return
+
+
+@app.cell
+def _(DGCNN, config, device, os, torch, train_dataset, wandb):
+ config.num_classes = train_dataset.num_classes
+
+ model = DGCNN(
+ out_channels=train_dataset.num_classes,
+ k=config.num_nearest_neighbours,
+ aggr=config.aggregation_operator
+ ).to(device)
+
+ model_artifact = wandb.use_artifact(config.artifact_address, type='checkpoint')
+ artifact_dir = model_artifact.download()
+ model_checkpoint_path = os.path.join(artifact_dir, "checkpoint.pt")
+
+ model.load_state_dict(torch.load(model_checkpoint_path)["model_state_dict"])
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Evaluation
+ """)
+ return
+
+
+@app.cell
+def _(
+ ShapeNet,
+ class_offset,
+ device,
+ jaccard_index,
+ model,
+ np,
+ scatter,
+ segmentation_class_frequency,
+ torch,
+ tqdm,
+ wandb,
+):
+ def evaluate(loader, split, table):
+ total_accuracy, total_iou = 0, 0
+ for data in tqdm(loader):
+ data = data.to(device)
+ with torch.no_grad():
+ model.eval()
+ outs = model(data)
+
+ predicted_labels = outs.argmax(dim=1)
+ accuracy = predicted_labels.eq(data.y).sum().item() / data.num_nodes
+
+ sizes = (data.ptr[1:] - data.ptr[:-1]).tolist()
+ ious, categories = [], []
+ y_map = torch.empty(
+ loader.dataset.num_classes, device=device
+ ).long()
+ for out, y, category in zip(
+ outs.split(sizes), data.y.split(sizes), data.category.tolist()
+ ):
+ category = list(ShapeNet.seg_classes.keys())[category]
+ part = ShapeNet.seg_classes[category]
+ part = torch.tensor(part, device=device)
+ y_map[part] = torch.arange(part.size(0), device=device)
+ iou = jaccard_index(
+ out[:, part].argmax(dim=-1), y_map[y],
+ task="multiclass", num_classes=part.size(0)
+ )
+ ious.append(iou)
+ categories.append(data.category)
+ iou = torch.tensor(ious, device=device)
+ category = torch.cat(categories, dim=0)
+ mean_iou = float(scatter(iou, category, reduce='mean').mean())
+
+ gt_pc_viz = data.pos.cpu().numpy().tolist()
+ segmentation_label = data.y.cpu().numpy().tolist()
+ frequency_dict = {key: 0 for key in segmentation_class_frequency.keys()}
+ for label in set(segmentation_label):
+ frequency_dict[label] = segmentation_label.count(label)
+ for j in range(len(gt_pc_viz)):
+ gt_pc_viz[j] += [segmentation_label[j] + 1 - class_offset]
+
+ predicted_pc_viz = data.pos.cpu().numpy().tolist()
+ segmentation_label = data.y.cpu().numpy().tolist()
+ frequency_dict = {key: 0 for key in segmentation_class_frequency.keys()}
+ for label in set(segmentation_label):
+ frequency_dict[label] = segmentation_label.count(label)
+ for j in range(len(predicted_pc_viz)):
+ predicted_pc_viz[j] += [segmentation_label[j] + 1 - class_offset]
+
+ table.add_data(
+ wandb.Object3D(np.array(gt_pc_viz)),
+ wandb.Object3D(np.array(predicted_pc_viz)),
+ accuracy, mean_iou, split, "DGCNN"
+ )
+ total_accuracy += accuracy
+ total_iou += mean_iou
+
+ wandb.log({
+ f"{split}/Accuracy": total_accuracy / len(loader),
+ f"{split}/IoU": total_iou / len(loader),
+ })
+
+ return table
+
+ return (evaluate,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We evaluate the results and store them in a Weights & Biases Table.
+ """)
+ return
+
+
+@app.cell
+def _(evaluate, test_loader, train_loader, wandb):
+ table = wandb.Table(columns=["Ground-Truth", "Prediction", "Accuracy", "IoU", "Split", "Model-Name"])
+ evaluate(train_loader, "Train-Val", table)
+ evaluate(test_loader, "Test", table)
+ return (table,)
+
+
+@app.cell
+def _(table, wandb):
+ wandb.log({"Evaluation-Results": table})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-pointnet-classification-00-eda/pyg_pointnet_classification_00_eda.py b/marimo/convert/pyg-pointnet-classification-00-eda/pyg_pointnet_classification_00_eda.py
new file mode 100644
index 00000000..de86b8e1
--- /dev/null
+++ b/marimo/convert/pyg-pointnet-classification-00-eda/pyg_pointnet_classification_00_eda.py
@@ -0,0 +1,222 @@
+# /// script
+# dependencies = ["https://data-pyg-org/whl/torch-${torch}-html", "pytorch_geometric @ git+https://github.com/pyg-team/pytorch_geometric.git", "torch-cluster", "torch-scatter", "torch-sparse", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Explore ModelNet Datasets using PyTorch Geometric and Weights & Biases 🪄🐝
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install Required Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import torch
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We now install PyTorch Geometric according to our PyTorch Version. We also install Weights & Biases.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: torch-scatter https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-sparse https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-cluster https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: git+https://github.com/pyg-team/pytorch_geometric.git !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Import Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ from glob import glob
+ from PIL import Image
+ from tqdm.auto import tqdm
+ import wandb
+ import torch.nn.functional as F
+ import numpy as np
+ import networkx as nx
+ import matplotlib.pyplot as plt
+ from pyvis.network import Network
+ from mpl_toolkits.mplot3d import Axes3D
+ import torch_geometric.transforms as T
+ from torch_geometric.datasets import ModelNet
+ from torch_geometric.loader import DataLoader
+ from torch_geometric.utils import to_networkx
+ from torch_geometric.nn import knn_graph, radius_graph
+
+ return ModelNet, T, glob, tqdm, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Initialize Weights & Biases
+
+ We need to call [`wandb.init()`](https://docs.wandb.ai/ref/python/init) once at the beginning of our program to initialize a new job. This creates a new run in W&B and launches a background process to sync data.
+ """)
+ return
+
+
+@app.cell
+def _(glob, os, wandb):
+ wandb_project = "pyg-point-cloud" #@param {"type": "string"}
+ wandb_run_name = "modelnet10/train/sampling-comparison" #@param {"type": "string"}
+
+
+ wandb.init(project=wandb_project, name=wandb_run_name, job_type="eda")
+
+ # Set experiment configs to be synced with wandb
+ config = wandb.config
+ config.display_sample = 2048 #@param {type:"slider", min:256, max:4096, step:16}
+ config.modelnet_dataset_alias = "ModelNet10" #@param ["ModelNet10", "ModelNet40"] {type:"raw"}
+
+ # Classes for ModelNet10 and ModelNet40
+ categories = sorted([
+ x.split(os.sep)[-2]
+ for x in glob(os.path.join(
+ config.modelnet_dataset_alias, "raw", '*', ''
+ ))
+ ])
+
+
+ config.categories = categories
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Load ModelNet Dataset using PyTorch Geometric
+ """)
+ return
+
+
+@app.cell
+def _(ModelNet, T, config):
+ pre_transform = T.NormalizeScale()
+ transform = T.SamplePoints(config.display_sample)
+ train_dataset = ModelNet(
+ root=config.modelnet_dataset_alias,
+ name=config.modelnet_dataset_alias[-2:],
+ train=True,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+ val_dataset = ModelNet(
+ root=config.modelnet_dataset_alias,
+ name=config.modelnet_dataset_alias[-2:],
+ train=False,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+ return train_dataset, val_dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log Data to [`wandb.Table`](https://docs.wandb.ai/ref/python/data-types/table)
+
+ We now log the dataset using a [Weights & Biases Table](https://docs.wandb.ai/guides/data-vis), which includes visualizing the individual point clouds as W&B's interactive 3D visualization format [`wandb.object3D`](https://docs.wandb.ai/ref/python/data-types/object3d). We also log the frequency distribution of the classes in the dataset using [`wandb.plot`](https://docs.wandb.ai/guides/track/log/plots).
+ """)
+ return
+
+
+@app.cell
+def _(config, tqdm, train_dataset, wandb):
+ _table = wandb.Table(columns=['Model', 'Class', 'Split'])
+ _category_dict = {key: 0 for key in config.categories}
+ for _idx in tqdm(range(len(train_dataset[:20]))):
+ _point_cloud = wandb.Object3D(train_dataset[_idx].pos.numpy())
+ _category = config.categories[int(train_dataset[_idx].y.item())]
+ _category_dict[_category] += 1
+ _table.add_data(_point_cloud, _category, 'Train')
+ _data = [[key, _category_dict[key]] for key in config.categories]
+ wandb.log({f'{config.modelnet_dataset_alias} Class-Frequency Distribution': wandb.plot.bar(wandb.Table(data=_data, columns=['Class', 'Frequency']), 'Class', 'Frequency', title=f'{config.modelnet_dataset_alias} Class-Frequency Distribution')})
+ return
+
+
+@app.cell
+def _(config, tqdm, val_dataset, wandb):
+ _table = wandb.Table(columns=['Model', 'Class', 'Split'])
+ _category_dict = {key: 0 for key in config.categories}
+ for _idx in tqdm(range(len(val_dataset[:100]))):
+ _point_cloud = wandb.Object3D(val_dataset[_idx].pos.numpy())
+ _category = config.categories[int(val_dataset[_idx].y.item())]
+ _category_dict[_category] += 1
+ _table.add_data(_point_cloud, _category, 'Test')
+ wandb.log({config.modelnet_dataset_alias: _table})
+ _data = [[key, _category_dict[key]] for key in config.categories]
+ wandb.log({f'{config.modelnet_dataset_alias} Class-Frequency Distribution': wandb.plot.bar(wandb.Table(data=_data, columns=['Class', 'Frequency']), 'Class', 'Frequency', title=f'{config.modelnet_dataset_alias} Class-Frequency Distribution')})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, you can check out the following notebook to learn how to compare different sampling strategies in PyTorch Geometric using Weights & Biases
+
+ [](http://wandb.me/pyg-sampling)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-pointnet-classification-01-compare-sampling/pyg_pointnet_classification_01_compare_sampling.py b/marimo/convert/pyg-pointnet-classification-01-compare-sampling/pyg_pointnet_classification_01_compare_sampling.py
new file mode 100644
index 00000000..01db47b5
--- /dev/null
+++ b/marimo/convert/pyg-pointnet-classification-01-compare-sampling/pyg_pointnet_classification_01_compare_sampling.py
@@ -0,0 +1,199 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Explore Graph Sampling Techniques using PyTorch Geometric and Weights & Biases 🪄🐝
+
+
+
+ If you wish to know how to explore and visualize point cloud datasets using PyTorch Geometric and Weights & Biases, you can check out the following notebook:
+
+ [](http://wandb.me/pyg-modelnet-eda)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install Required Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import torch
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We now install PyTorch Geometric according to our PyTorch Version. We also install Weights & Biases.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}.html
+ !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ !pip install -q wandb
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Import Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import random
+ import numpy as np
+ from glob import glob
+ from tqdm.auto import tqdm
+ from matplotlib import pyplot as plt
+ import wandb
+ import networkx as nx
+ from pyvis.network import Network
+ import torch.nn.functional as F
+ import torch_geometric.transforms as T
+ from torch_geometric.datasets import ModelNet
+ from torch_geometric.loader import DataLoader
+ from torch_geometric.utils import to_networkx
+ from torch_geometric.nn import knn_graph, radius_graph
+
+ return (
+ ModelNet,
+ Network,
+ T,
+ knn_graph,
+ nx,
+ radius_graph,
+ to_networkx,
+ wandb,
+ )
+
+
+@app.cell
+def _(ModelNet, T):
+ pre_transform = T.NormalizeScale()
+ transform = T.SamplePoints(128)
+ low_train_dataset = ModelNet(
+ root="ModelNet10",
+ name='10',
+ train=True,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+
+ pre_transform = T.NormalizeScale()
+ transform = T.SamplePoints(2048)
+ high_train_dataset = ModelNet(
+ root="ModelNet10",
+ name='10',
+ train=True,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+ return high_train_dataset, low_train_dataset
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We take a single point cloud from the dataset and compare the KNN-sampled subgraph and radius-sampled subgraph by visualizing the subgraphs as [`wandb.Html`](https://docs.wandb.ai/ref/python/data-types/html) on a [Weights & Biases Table](https://docs.wandb.ai/guides/data-vis).
+ """)
+ return
+
+
+@app.cell
+def _(
+ Network,
+ high_train_dataset,
+ knn_graph,
+ low_train_dataset,
+ nx,
+ radius_graph,
+ to_networkx,
+ wandb,
+):
+ with wandb.init(
+ project="pyg-point-cloud",
+ name="sampling/modelnet10",
+ entity="geekyrakshit",
+ job_type="eda"
+ ):
+ table = wandb.Table(columns=[
+ "Model", "KNN-Sampled-Subgraph", "Nearest-Neighbours", "Radius-Sampled-Subgraph", "Radius"
+ ])
+
+ sample_data = low_train_dataset[0]
+
+ sample_data.edge_index = knn_graph(sample_data.pos, k=6)
+ G = to_networkx(sample_data, to_undirected=True)
+ nt = Network('500px', '500px')
+ nt.from_nx(G)
+ knn_sampled = wandb.Html(nt.generate_html())
+
+ sample_data = low_train_dataset[0]
+ sample_data.edge_index = radius_graph(sample_data.pos, r=0.5)
+ G = to_networkx(sample_data, to_undirected=True)
+ nx.draw(G)
+ nt = Network('500px', '500px')
+ nt.from_nx(G)
+ radius_sampled = wandb.Html(nt.generate_html())
+
+ table.add_data(
+ wandb.Object3D(high_train_dataset[0].pos.numpy()), knn_sampled, 6, radius_sampled, 0.
+ )
+
+ wandb.log({"Sampling-Comparison": table})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, you can check out the following notebook to learn how to train the PointNet++ architecture using PyTorch Geometric and Weights & Biases
+
+ [](http://wandb.me/pyg-pointnet2-train)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-pointnet-classification-02-pointnet-plus-plus/pyg_pointnet_classification_02_pointnet_plus_plus.py b/marimo/convert/pyg-pointnet-classification-02-pointnet-plus-plus/pyg_pointnet_classification_02_pointnet_plus_plus.py
new file mode 100644
index 00000000..3831de6b
--- /dev/null
+++ b/marimo/convert/pyg-pointnet-classification-02-pointnet-plus-plus/pyg_pointnet_classification_02_pointnet_plus_plus.py
@@ -0,0 +1,505 @@
+# /// script
+# dependencies = ["https://data-pyg-org/whl/torch-${torch}-html", "pytorch_geometric @ git+https://github.com/pyg-team/pytorch_geometric.git", "torch-cluster", "torch-scatter", "torch-sparse", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Train PointNet++ Model using PyTorch Geometric and Weights & Biases 🪄🐝
+
+
+
+ This notebook demonstrates an implementation of the [PointeNet++](https://arxiv.org/pdf/1706.02413.pdf) architecture implemented using PyTorch Geometric and experiment tracked and visualized using [Weights & Biases](https://wandb.ai/site).
+
+ If you wish to know how to compare and visualize the different sampling strategies used in the PointNet++ implementation, you can check out the following notebook:
+
+ [](http://wandb.me/pyg-pointnet2-train)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install Required Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install required packages.
+ import os
+ import torch
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return os, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We now install PyTorch Geometric according to our PyTorch Version. We also install Weights & Biases.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: torch-scatter https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-sparse https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-cluster https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: git+https://github.com/pyg-team/pytorch_geometric.git !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Import Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import random
+ from glob import glob
+ from tqdm.auto import tqdm
+ import wandb
+ import torch.nn.functional as F
+ import torch_geometric.transforms as T
+ from torch_geometric.datasets import ModelNet
+ from torch_geometric.loader import DataLoader
+ from torch_geometric.nn import MLP, PointConv, fps, global_max_pool, radius
+
+ return (
+ DataLoader,
+ F,
+ MLP,
+ ModelNet,
+ PointConv,
+ T,
+ fps,
+ glob,
+ global_max_pool,
+ radius,
+ random,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Initialize Weights & Biases
+
+ We need to call [`wandb.init()`](https://docs.wandb.ai/ref/python/init) once at the beginning of our program to initialize a new job. This creates a new run in W&B and launches a background process to sync data.
+ """)
+ return
+
+
+@app.cell
+def _(glob, os, random, torch, wandb):
+ wandb_project = "pyg-point-cloud" #@param {"type": "string"}
+ wandb_run_name = "final-experiment/modelnet10/2" #@param {"type": "string"}
+
+ wandb.init(project=wandb_project, name=wandb_run_name, job_type="baseline-train")
+
+ # Set experiment configs to be synced with wandb
+ config = wandb.config
+ config.modelnet_dataset_alias = "ModelNet10" #@param ["ModelNet10", "ModelNet40"] {type:"raw"}
+
+ config.seed = 4242 #@param {type:"number"}
+ random.seed(config.seed)
+ torch.manual_seed(config.seed)
+
+ config.sample_points = 2048 #@param {type:"slider", min:256, max:4096, step:16}
+
+ config.categories = sorted([
+ x.split(os.sep)[-2]
+ for x in glob(os.path.join(
+ config.modelnet_dataset_alias, "raw", '*', ''
+ ))
+ ])
+
+ config.batch_size = 16 #@param {type:"slider", min:4, max:128, step:4}
+ config.num_workers = 6 #@param {type:"slider", min:1, max:10, step:1}
+
+ config.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ device = torch.device(config.device)
+
+ config.set_abstraction_ratio_1 = 0.748 #@param {type:"slider", min:0.1, max:1.0, step:0.01}
+ config.set_abstraction_radius_1 = 0.4817 #@param {type:"slider", min:0.1, max:1.0, step:0.01}
+ config.set_abstraction_ratio_2 = 0.3316 #@param {type:"slider", min:0.1, max:1.0, step:0.01}
+ config.set_abstraction_radius_2 = 0.2447 #@param {type:"slider", min:0.1, max:1.0, step:0.01}
+ config.dropout = 0.1 #@param {type:"slider", min:0.1, max:1.0, step:0.1}
+
+ config.learning_rate = 1e-4 #@param {type:"number"}
+ config.epochs = 10 #@param {type:"slider", min:1, max:100, step:1}
+ config.num_visualization_samples = 20 #@param {type:"slider", min:1, max:100, step:1}
+ return config, device
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Load ModelNet Dataset using PyTorch Geometric
+
+ We now load, preprocess and batch the ModelNet dataset for training, validation/testing and visualization.
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, ModelNet, T, config, random):
+ pre_transform = T.NormalizeScale()
+ transform = T.SamplePoints(config.sample_points)
+
+
+ train_dataset = ModelNet(
+ root=config.modelnet_dataset_alias,
+ name=config.modelnet_dataset_alias[-2:],
+ train=True,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+ train_loader = DataLoader(
+ train_dataset,
+ batch_size=config.batch_size,
+ shuffle=True,
+ num_workers=config.num_workers
+ )
+
+ val_dataset = ModelNet(
+ root=config.modelnet_dataset_alias,
+ name=config.modelnet_dataset_alias[-2:],
+ train=False,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+ val_loader = DataLoader(
+ val_dataset,
+ batch_size=config.batch_size,
+ shuffle=False,
+ num_workers=config.num_workers
+ )
+
+ random_indices = random.sample(
+ list(range(len(val_dataset))),
+ config.num_visualization_samples
+ )
+ vizualization_loader = DataLoader(
+ [val_dataset[idx] for idx in random_indices],
+ batch_size=1,
+ shuffle=False,
+ num_workers=config.num_workers
+ )
+ return train_loader, val_loader, vizualization_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Implementing the PointNet++ Model using PyTorch Geometric
+ """)
+ return
+
+
+@app.cell
+def _(PointConv, fps, radius, torch):
+ class SetAbstraction(torch.nn.Module):
+ def __init__(self, ratio, r, nn):
+ super().__init__()
+ self.ratio = ratio
+ self.r = r
+ self.conv = PointConv(nn, add_self_loops=False)
+
+ def forward(self, x, pos, batch):
+ idx = fps(pos, batch, ratio=self.ratio)
+ row, col = radius(pos, pos[idx], self.r, batch, batch[idx],
+ max_num_neighbors=64)
+ edge_index = torch.stack([col, row], dim=0)
+ x_dst = None if x is None else x[idx]
+ x = self.conv((x, x_dst), (pos, pos[idx]), edge_index)
+ pos, batch = pos[idx], batch[idx]
+ return x, pos, batch
+
+ return (SetAbstraction,)
+
+
+@app.cell
+def _(global_max_pool, torch):
+ class GlobalSetAbstraction(torch.nn.Module):
+ def __init__(self, nn):
+ super().__init__()
+ self.nn = nn
+
+ def forward(self, x, pos, batch):
+ x = self.nn(torch.cat([x, pos], dim=1))
+ x = global_max_pool(x, batch)
+ pos = pos.new_zeros((x.size(0), 3))
+ batch = torch.arange(x.size(0), device=batch.device)
+ return x, pos, batch
+
+ return (GlobalSetAbstraction,)
+
+
+@app.cell
+def _(GlobalSetAbstraction, MLP, SetAbstraction, torch):
+ class PointNet2(torch.nn.Module):
+ def __init__(
+ self,
+ set_abstraction_ratio_1, set_abstraction_ratio_2,
+ set_abstraction_radius_1, set_abstraction_radius_2, dropout
+ ):
+ super().__init__()
+
+ # Input channels account for both `pos` and node features.
+ self.sa1_module = SetAbstraction(
+ set_abstraction_ratio_1,
+ set_abstraction_radius_1,
+ MLP([3, 64, 64, 128])
+ )
+ self.sa2_module = SetAbstraction(
+ set_abstraction_ratio_2,
+ set_abstraction_radius_2,
+ MLP([128 + 3, 128, 128, 256])
+ )
+ self.sa3_module = GlobalSetAbstraction(MLP([256 + 3, 256, 512, 1024]))
+
+ self.mlp = MLP([1024, 512, 256, 10], dropout=dropout, norm=None)
+
+ def forward(self, data):
+ sa0_out = (data.x, data.pos, data.batch)
+ sa1_out = self.sa1_module(*sa0_out)
+ sa2_out = self.sa2_module(*sa1_out)
+ sa3_out = self.sa3_module(*sa2_out)
+ x, pos, batch = sa3_out
+
+ return self.mlp(x).log_softmax(dim=-1)
+
+ return (PointNet2,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training PointNet++ and Logging Metrics on Weights & Biases
+ """)
+ return
+
+
+@app.cell
+def _(PointNet2, config, device, torch):
+ # Define PointNet++ model.
+ model = PointNet2(
+ config.set_abstraction_ratio_1,
+ config.set_abstraction_ratio_2,
+ config.set_abstraction_radius_1,
+ config.set_abstraction_radius_2,
+ config.dropout
+ ).to(device)
+
+ # Define Optimizer
+ optimizer = torch.optim.Adam(
+ model.parameters(), lr=config.learning_rate
+ )
+ return model, optimizer
+
+
+@app.cell
+def _(
+ F,
+ config,
+ device,
+ model,
+ optimizer,
+ torch,
+ tqdm,
+ train_loader,
+ val_loader,
+ vizualization_loader,
+ wandb,
+):
+ def train_step(epoch):
+ """Training Step"""
+ model.train()
+ epoch_loss, correct = 0, 0
+ num_train_examples = len(train_loader)
+
+ progress_bar = tqdm(
+ range(num_train_examples),
+ desc=f"Training Epoch {epoch}/{config.epochs}"
+ )
+ data_iter = iter(train_loader)
+ for batch_idx in progress_bar:
+ data = next(data_iter).to(device)
+
+ optimizer.zero_grad()
+ prediction = model(data)
+ loss = F.nll_loss(prediction, data.y)
+ loss.backward()
+ optimizer.step()
+
+ epoch_loss += loss.item()
+ correct += prediction.max(1)[1].eq(data.y).sum().item()
+
+ epoch_loss = epoch_loss / num_train_examples
+ epoch_accuracy = correct / len(train_loader.dataset)
+
+ wandb.log({
+ "Train/Loss": epoch_loss,
+ "Train/Accuracy": epoch_accuracy
+ })
+
+
+ def val_step(epoch):
+ """Validation Step"""
+ model.eval()
+ epoch_loss, correct = 0, 0
+ num_val_examples = len(val_loader)
+
+ progress_bar = tqdm(
+ range(num_val_examples),
+ desc=f"Validation Epoch {epoch}/{config.epochs}"
+ )
+ data_iter = iter(val_loader)
+ for batch_idx in progress_bar:
+ data = next(data_iter).to(device)
+
+ with torch.no_grad():
+ prediction = model(data)
+
+ loss = F.nll_loss(prediction, data.y)
+ epoch_loss += loss.item()
+ correct += prediction.max(1)[1].eq(data.y).sum().item()
+
+ epoch_loss = epoch_loss / num_val_examples
+ epoch_accuracy = correct / len(val_loader.dataset)
+
+ wandb.log({
+ "Validation/Loss": epoch_loss,
+ "Validation/Accuracy": epoch_accuracy
+ })
+
+
+ def visualize_evaluation(table, epoch):
+ """Visualize validation result in a Weights & Biases Table"""
+ point_clouds, losses, predictions, ground_truths, is_correct = [], [], [], [], []
+ progress_bar = tqdm(
+ range(config.num_visualization_samples),
+ desc=f"Generating Visualizations for Epoch {epoch}/{config.epochs}"
+ )
+
+ for idx in progress_bar:
+ data = next(iter(vizualization_loader)).to(device)
+
+ with torch.no_grad():
+ prediction = model(data)
+
+ point_clouds.append(
+ wandb.Object3D(torch.squeeze(data.pos, dim=0).cpu().numpy())
+ )
+ losses.append(F.nll_loss(prediction, data.y).item())
+ predictions.append(config.categories[int(prediction.max(1)[1].item())])
+ ground_truths.append(config.categories[int(data.y.item())])
+ is_correct.append(prediction.max(1)[1].eq(data.y).sum().item())
+
+ table.add_data(
+ epoch, point_clouds, losses, predictions, ground_truths, is_correct
+ )
+ return table
+
+
+ def save_checkpoint(epoch):
+ """Save model checkpoints as Weights & Biases artifacts"""
+ torch.save({
+ 'epoch': epoch,
+ 'model_state_dict': model.state_dict(),
+ 'optimizer_state_dict': optimizer.state_dict()
+ }, "checkpoint.pt")
+
+ artifact_name = wandb.util.make_artifact_name_safe(
+ f"{wandb.run.name}-{wandb.run.id}-checkpoint"
+ )
+
+ checkpoint_artifact = wandb.Artifact(artifact_name, type="checkpoint")
+ checkpoint_artifact.add_file("checkpoint.pt")
+ wandb.log_artifact(
+ checkpoint_artifact, aliases=["latest", f"epoch-{epoch}"]
+ )
+
+ return save_checkpoint, train_step, val_step, visualize_evaluation
+
+
+@app.cell
+def _(
+ config,
+ save_checkpoint,
+ train_step,
+ val_step,
+ visualize_evaluation,
+ wandb,
+):
+ table = wandb.Table(
+ columns=[
+ "Epoch",
+ "Point-Clouds",
+ "Losses",
+ "Predicted-Classes",
+ "Ground-Truth",
+ "Is-Correct"
+ ]
+ )
+ for epoch in range(1, config.epochs + 1):
+ train_step(epoch)
+ val_step(epoch)
+ visualize_evaluation(table, epoch)
+ save_checkpoint(epoch)
+ wandb.log({"Evaluation": table})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, you can check out the following notebook to learn how to run a hyperparameter sweep on our PointNet++ trainig loop using Weights & Biases:
+
+ |Tune Hyperparameters using Weights & Biases Sweep|[](http://wandb.me/pyg-pointnet2-sweep)|
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pyg-pointnet-classification-03-sweep/pyg_pointnet_classification_03_sweep.py b/marimo/convert/pyg-pointnet-classification-03-sweep/pyg_pointnet_classification_03_sweep.py
new file mode 100644
index 00000000..35d8d6ae
--- /dev/null
+++ b/marimo/convert/pyg-pointnet-classification-03-sweep/pyg_pointnet_classification_03_sweep.py
@@ -0,0 +1,394 @@
+# /// script
+# dependencies = ["https://data-pyg-org/whl/torch-${torch}-html", "pytorch_geometric @ git+https://github.com/pyg-team/pytorch_geometric.git", "torch-cluster", "torch-scatter", "torch-sparse", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Run a Hyperparamter Sweep on PointNet++ 🪄🐝
+
+
+
+ This notebook demonstrates the process of running a [Hyperparameter Sweep using Weights & Biases](https://docs.wandb.ai/guides/sweeps) on our point cloud classification training workflow in order to maximize the performance of our model.
+
+ If you wish to know how to implement the PointNet++ architecture and train it you can check out the following [notebook](http://wandb.me/pyg-sampling).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install Required Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install required packages.
+ import os
+ import torch
+ os.environ['TORCH'] = torch.__version__
+ print(torch.__version__)
+ return os, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We now install PyTorch Geometric according to our PyTorch Version. We also install Weights & Biases.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: torch-scatter https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-sparse https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: torch-cluster https://data.pyg.org/whl/torch-${TORCH}.html !pip install -q torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}.html
+ # packages added via marimo's package management: git+https://github.com/pyg-team/pytorch_geometric.git !pip install -q git+https://github.com/pyg-team/pytorch_geometric.git
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Import Libraries
+ """)
+ return
+
+
+@app.cell
+def _():
+ import gc
+ from glob import glob
+ from tqdm.auto import tqdm
+ import wandb
+ import torch.nn.functional as F
+ import torch_geometric.transforms as T
+ from torch_geometric.datasets import ModelNet
+ from torch_geometric.loader import DataLoader
+ from torch_geometric.nn import MLP, PointConv, fps, global_max_pool, radius
+
+ return (
+ DataLoader,
+ F,
+ MLP,
+ ModelNet,
+ PointConv,
+ T,
+ fps,
+ gc,
+ glob,
+ global_max_pool,
+ radius,
+ tqdm,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Function to Build Data Loaders
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, ModelNet, T):
+ def get_dataset_and_loaders(sample_points, batch_size, num_workers):
+ pre_transform = T.NormalizeScale()
+ transform = T.SamplePoints(sample_points)
+
+ train_dataset = ModelNet(
+ root="ModelNet10",
+ name='10',
+ train=True,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+ train_loader = DataLoader(
+ train_dataset,
+ batch_size=batch_size,
+ shuffle=True,
+ num_workers=num_workers
+ )
+
+ val_dataset = ModelNet(
+ root="ModelNet10",
+ name='10',
+ train=False,
+ transform=transform,
+ pre_transform=pre_transform
+ )
+ val_loader = DataLoader(
+ val_dataset,
+ batch_size=batch_size,
+ shuffle=False,
+ num_workers=num_workers
+ )
+
+ return train_dataset, train_loader, val_dataset, val_loader
+
+ return (get_dataset_and_loaders,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Implementing the PointNet++ Model using PyTorch Geometric
+ """)
+ return
+
+
+@app.cell
+def _(MLP, PointConv, fps, global_max_pool, radius, torch):
+ class SetAbstraction(torch.nn.Module):
+ def __init__(self, ratio, r, nn):
+ super().__init__()
+ self.ratio = ratio
+ self.r = r
+ self.conv = PointConv(nn, add_self_loops=False)
+
+ def forward(self, x, pos, batch):
+ idx = fps(pos, batch, ratio=self.ratio)
+ row, col = radius(pos, pos[idx], self.r, batch, batch[idx], max_num_neighbors=64)
+ edge_index = torch.stack([col, row], dim=0)
+ x_dst = None if x is None else x[idx]
+ x = self.conv((x, x_dst), (pos, pos[idx]), edge_index)
+ pos, batch = pos[idx], batch[idx]
+ return x, pos, batch
+
+
+ class GlobalSetAbstraction(torch.nn.Module):
+ def __init__(self, nn):
+ super().__init__()
+ self.nn = nn
+
+ def forward(self, x, pos, batch):
+ x = self.nn(torch.cat([x, pos], dim=1))
+ x = global_max_pool(x, batch)
+ pos = pos.new_zeros((x.size(0), 3))
+ batch = torch.arange(x.size(0), device=batch.device)
+ return x, pos, batch
+
+
+ class PointNet2(torch.nn.Module):
+ def __init__(self, set_abstraction_ratio_1, set_abstraction_ratio_2, dropout):
+ super().__init__()
+
+ # Input channels account for both `pos` and node features.
+ self.sa1_module = SetAbstraction(
+ set_abstraction_ratio_1, 0.2, MLP([3, 64, 64, 128])
+ )
+ self.sa2_module = SetAbstraction(
+ set_abstraction_ratio_2, 0.4, MLP([128 + 3, 128, 128, 256])
+ )
+ self.sa3_module = GlobalSetAbstraction(MLP([256 + 3, 256, 512, 1024]))
+
+ self.mlp = MLP([1024, 512, 256, 10], dropout=dropout, norm=None)
+
+ def forward(self, data):
+ sa0_out = (data.x, data.pos, data.batch)
+ sa1_out = self.sa1_module(*sa0_out)
+ sa2_out = self.sa2_module(*sa1_out)
+ sa3_out = self.sa3_module(*sa2_out)
+ x, pos, batch = sa3_out
+
+ return self.mlp(x).log_softmax(dim=-1)
+
+ return (PointNet2,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Define a Training Function Instrumented with WandB
+ """)
+ return
+
+
+@app.cell
+def _(F, PointNet2, gc, get_dataset_and_loaders, glob, os, torch, tqdm, wandb):
+ def train():
+ wandb.init(project="pyg-point-cloud", entity="geekyrakshit")
+
+ # Set Default Configs
+ config = wandb.config
+ config.categories = sorted([
+ x.split(os.sep)[-2]
+ for x in glob(os.path.join("ModelNet10", "raw", '*', ''))
+ ])
+ config.num_workers = 6
+ config.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ device = torch.device(config.device)
+ config.learning_rate = 1e-4
+ config.epochs = 5
+
+ # Get tuned configs from sweep
+ batch_size = config.batch_size
+ sample_points = config.sample_points
+ set_abstraction_ratio_1 = config.set_abstraction_ratio_1
+ set_abstraction_ratio_2 = config.set_abstraction_ratio_2
+ dropout = config.dropout
+
+ # Create datasets and dataloaders
+ (
+ train_dataset, train_loader, val_dataset, val_loader
+ ) = get_dataset_and_loaders(
+ sample_points, batch_size, config.num_workers
+ )
+
+ model = PointNet2(
+ set_abstraction_ratio_1, set_abstraction_ratio_2, dropout
+ ).to(device)
+ optimizer = torch.optim.Adam(
+ model.parameters(), lr=config.learning_rate
+ )
+
+ for epoch in range(1, config.epochs + 1):
+
+ # Training Step
+ model.train()
+ epoch_loss, correct = 0, 0
+ num_train_examples = len(train_loader)
+
+ progress_bar = tqdm(
+ range(num_train_examples),
+ desc=f"Training Epoch {epoch}/{config.epochs}"
+ )
+ data_iter = iter(train_loader)
+ for batch_idx in progress_bar:
+ data = next(data_iter).to(device)
+
+ optimizer.zero_grad()
+ prediction = model(data)
+ loss = F.nll_loss(prediction, data.y)
+ loss.backward()
+ optimizer.step()
+
+ epoch_loss += loss.item()
+ correct += prediction.max(1)[1].eq(data.y).sum().item()
+
+ epoch_loss = epoch_loss / num_train_examples
+ epoch_accuracy = correct / len(train_loader.dataset)
+
+ wandb.log({
+ "Train/Loss": epoch_loss,
+ "Train/Accuracy": epoch_accuracy
+ })
+
+ # Validation Step
+ model.eval()
+ epoch_loss, correct = 0, 0
+ num_val_examples = len(val_loader)
+
+ progress_bar = tqdm(
+ range(num_val_examples),
+ desc=f"Validation Epoch {epoch}/{config.epochs}"
+ )
+ data_iter = iter(val_loader)
+ for batch_idx in progress_bar:
+ data = next(data_iter).to(device)
+
+ with torch.no_grad():
+ prediction = model(data)
+
+ loss = F.nll_loss(prediction, data.y)
+ epoch_loss += loss.item()
+ correct += prediction.max(1)[1].eq(data.y).sum().item()
+
+ epoch_loss = epoch_loss / num_val_examples
+ epoch_accuracy = correct / len(val_loader.dataset)
+
+ wandb.log({
+ "Validation/Loss": epoch_loss,
+ "Validation/Accuracy": epoch_accuracy
+ })
+
+ # Save Checkpoint
+ torch.save({
+ 'epoch': epoch,
+ 'model_state_dict': model.state_dict(),
+ 'optimizer_state_dict': optimizer.state_dict()
+ }, "checkpoint.pt")
+
+ artifact_name = wandb.util.make_artifact_name_safe(
+ f"{wandb.run.name}-{wandb.run.id}-checkpoint"
+ )
+
+ checkpoint_artifact = wandb.Artifact(artifact_name, type="checkpoint")
+ checkpoint_artifact.add_file("checkpoint.pt")
+ wandb.log_artifact(
+ checkpoint_artifact, aliases=["latest", f"epoch-{epoch}"]
+ )
+
+ model = model.cpu()
+ del model
+ gc.collect()
+ torch.cuda.empty_cache()
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Start the Hyperparameter Sweep
+ """)
+ return
+
+
+@app.cell
+def _(train, wandb):
+ # Define sweep configuration
+ sweep_configuration = {
+ 'method': 'bayes',
+ 'metric': {'goal': 'maximize', 'name': 'Validation/Accuracy'},
+ 'parameters':
+ {
+ 'batch_size': {'values': [8, 16, 32, 64]},
+ 'sample_points': {'values': [512, 1024, 2048]},
+ 'set_abstraction_ratio_1': {'min': 0.1, 'max': 0.9},
+ 'set_abstraction_ratio_2': {'min': 0.1, 'max': 0.9},
+ 'dropout': {'min': 0.1, 'max': 0.7},
+ }
+ }
+
+ # Get Sweep ID
+ sweep_id = wandb.sweep(
+ sweep=sweep_configuration, project='pyg-point-cloud', entity="geekyrakshit"
+ )
+
+ # Run Sweep
+ wandb.agent(sweep_id, function=train, count=30)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-how-does-adding-dropout-affect-model-performance/pytorch_how_does_adding_dropout_affect_model_performance.py b/marimo/convert/pytorch-how-does-adding-dropout-affect-model-performance/pytorch_how_does_adding_dropout_affect_model_performance.py
new file mode 100644
index 00000000..e3ad04af
--- /dev/null
+++ b/marimo/convert/pytorch-how-does-adding-dropout-affect-model-performance/pytorch_how_does_adding_dropout_affect_model_performance.py
@@ -0,0 +1,445 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this colab, we'll see an example of adding dropout to a PyTorch model and observe the effect dropout has on the model's performance by tracking our models in [Weights & Biases](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk).
+
+ You can read more about using dropout in PyTorch [here](https://wandb.ai/authors/ayusht/reports/Dropout-in-PyTorch-An-Example--VmlldzoxNTgwOTE).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import torch
+ from torch import nn
+ from torch import optim
+ from torch.nn import functional as F
+ import torchvision
+ from torchvision import datasets, transforms
+ from torch.utils.data import DataLoader
+
+ import matplotlib.pyplot as plt
+ import numpy as np
+
+ return F, nn, np, optim, plt, torch, torchvision, transforms
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print(device)
+ return (device,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download dataset and prepare `DataLoader`s
+ """)
+ return
+
+
+@app.cell
+def _(torch, torchvision, transforms):
+ BATCH_SIZE = 32
+
+ transform = transforms.Compose(
+ [transforms.ToTensor(),
+ transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+ trainset = torchvision.datasets.CIFAR10(root="./data", train=True,
+ download=True, transform=transform)
+ trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,
+ shuffle=True, num_workers=2)
+
+ testset = torchvision.datasets.CIFAR10(root="./data", train=False,
+ download=True, transform=transform)
+ testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE,
+ shuffle=False, num_workers=2)
+
+ CLASS_NAMES = ("plane", "car", "bird", "cat",
+ "deer", "dog", "frog", "horse", "ship", "truck")
+ return CLASS_NAMES, testloader, trainloader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Visualize data
+ """)
+ return
+
+
+@app.cell
+def _(CLASS_NAMES, np, plt):
+ def show_batch(image_batch, label_batch):
+ plt.figure(figsize=(10,10))
+ for n in range(25):
+ ax = plt.subplot(5,5,n+1)
+ img = image_batch[n] / 2 + 0.5 # unnormalize
+ img = img.numpy()
+ plt.imshow(np.transpose(img, (1, 2, 0)))
+ plt.title(CLASS_NAMES[label_batch[n]])
+ plt.axis("off")
+
+ return (show_batch,)
+
+
+@app.cell
+def _(show_batch, trainloader):
+ sample_images, sample_labels = next(iter(trainloader))
+ show_batch(sample_images, sample_labels)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Define training
+ """)
+ return
+
+
+@app.cell
+def _(torch, wandb):
+ def train(model, device, train_loader, optimizer, criterion, epoch, steps_per_epoch=20):
+ model.train()
+ train_loss = 0
+ train_total = 0
+ train_correct = 0
+ for batch_idx, (data, target) in enumerate(train_loader, start=0):
+ data, target = (data.to(device), target.to(device))
+ optimizer.zero_grad()
+ output = model(data)
+ loss = criterion(output, target)
+ train_loss = train_loss + loss.item()
+ scores, predictions = torch.max(output.data, 1)
+ train_total = train_total + target.size(0)
+ train_correct = train_correct + int(sum(predictions == target))
+ optimizer.zero_grad()
+ loss.backward()
+ optimizer.step()
+ acc = round(train_correct / train_total * 100, 2)
+ print('Epoch [{}], Loss: {}, Accuracy: {}'.format(_epoch, train_loss / train_total, acc), end='')
+ wandb.log({'Train Loss': train_loss / train_total, 'Train Accuracy': acc, 'Epoch': _epoch})
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Define testing
+ """)
+ return
+
+
+@app.cell
+def _(torch, wandb):
+ def test(model, device, test_loader, criterion, classes):
+ model.eval() # Switch model to evaluation mode. This is necessary for layers like dropout, batchnorm etc which behave differently in training and evaluation mode
+ test_loss = 0
+ test_total = 0
+ test_correct = 0
+ example_images = []
+ with torch.no_grad():
+ for data, target in test_loader:
+ data, target = (data.to(device), target.to(device))
+ output = model(data)
+ test_loss = test_loss + criterion(output, target).item()
+ scores, predictions = torch.max(output.data, 1) # Load the input features and labels from the test dataset
+ test_total = test_total + target.size(0)
+ test_correct = test_correct + int(sum(predictions == target))
+ acc = round(test_correct / test_total * 100, 2) # Make predictions: Pass image data from test dataset, make predictions about class image belongs to (0-9 in this case)
+ print(' Test_loss: {}, Test_accuracy: {}'.format(test_loss / test_total, acc))
+ wandb.log({'Test Loss': test_loss / test_total, 'Test Accuracy': acc}) # Compute the loss sum up batch loss
+
+ return (test,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training the unregularized model
+ """)
+ return
+
+
+@app.cell
+def _(F, nn, torch):
+ class Net(nn.Module):
+ def __init__(self, input_shape=(3,32,32)):
+ super(Net, self).__init__()
+
+ self.conv1 = nn.Conv2d(3, 32, 3)
+ self.conv2 = nn.Conv2d(32, 64, 3)
+ self.conv3 = nn.Conv2d(64, 128, 3)
+
+ self.pool = nn.MaxPool2d(2,2)
+
+ n_size = self._get_conv_output(input_shape)
+
+ self.fc1 = nn.Linear(n_size, 512)
+ self.fc2 = nn.Linear(512, 10)
+
+ def _get_conv_output(self, shape):
+ batch_size = 1
+ input = torch.autograd.Variable(torch.rand(batch_size, *shape))
+ output_feat = self._forward_features(input)
+ n_size = output_feat.data.view(batch_size, -1).size(1)
+ return n_size
+
+ def _forward_features(self, x):
+ x = self.pool(F.relu(self.conv1(x)))
+ x = self.pool(F.relu(self.conv2(x)))
+ x = self.pool(F.relu(self.conv3(x)))
+ return x
+
+ def forward(self, x):
+ x = self._forward_features(x)
+ x = x.view(x.size(0), -1)
+ x = F.relu(self.fc1(x))
+ x = self.fc2(x)
+ return x
+
+ return (Net,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Initialize Model, Loss and Optimizer
+ """)
+ return
+
+
+@app.cell
+def _(Net, device, nn, optim):
+ net = Net().to(device)
+ print(net)
+
+ criterion = nn.CrossEntropyLoss()
+ optimizer = optim.Adam(net.parameters())
+ return criterion, net, optimizer
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Train
+ """)
+ return
+
+
+@app.cell
+def _(
+ CLASS_NAMES,
+ criterion,
+ device,
+ net,
+ optimizer,
+ test,
+ testloader,
+ train,
+ trainloader,
+ wandb,
+):
+ wandb.init(project='dropout')
+ wandb.watch(net, log='all')
+ for _epoch in range(8):
+ train(net, device, trainloader, optimizer, criterion, _epoch)
+ test(net, device, testloader, criterion, CLASS_NAMES)
+ print('Finished Training')
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training a model with dropout regularization
+ """)
+ return
+
+
+@app.cell
+def _(F, Net, nn, torch):
+ class Net_1(nn.Module):
+
+ def __init__(self, input_shape=(3, 32, 32)):
+ super(Net, self).__init__()
+ self.conv1 = nn.Conv2d(3, 32, 3)
+ self.conv2 = nn.Conv2d(32, 64, 3)
+ self.conv3 = nn.Conv2d(64, 128, 3)
+ self.pool = nn.MaxPool2d(2, 2)
+ n_size = self._get_conv_output(input_shape)
+ self.fc1 = nn.Linear(n_size, 512)
+ self.fc2 = nn.Linear(512, 10)
+ self.dropout = nn.Dropout(0.25)
+
+ def _get_conv_output(self, shape):
+ batch_size = 1
+ input = torch.autograd.Variable(torch.rand(batch_size, *shape))
+ output_feat = self._forward_features(input)
+ n_size = output_feat.data.view(batch_size, -1).size(1)
+ return n_size
+
+ def _forward_features(self, x):
+ x = self.pool(F.relu(self.conv1(x)))
+ x = self.pool(F.relu(self.conv2(x)))
+ x = self.pool(F.relu(self.conv3(x)))
+ return x
+
+ def forward(self, x):
+ x = self._forward_features(x)
+ x = x.view(x.size(0), -1)
+ x = self.dropout(x)
+ x = F.relu(self.fc1(x))
+ x = self.dropout(x)
+ x = self.fc2(x)
+ return x
+
+ return (Net_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Initialize Model, Loss and Optimizer
+ """)
+ return
+
+
+@app.cell
+def _(Net_1, device, nn, optim):
+ net_1 = Net_1().to(device)
+ print(net_1)
+ criterion_1 = nn.CrossEntropyLoss()
+ optimizer_1 = optim.Adam(net_1.parameters())
+ return criterion_1, net_1, optimizer_1
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### Train
+ """)
+ return
+
+
+@app.cell
+def _(
+ CLASS_NAMES,
+ criterion_1,
+ device,
+ net_1,
+ optimizer_1,
+ test,
+ testloader,
+ train,
+ trainloader,
+ wandb,
+):
+ wandb.init(project='dropout')
+ wandb.watch(net_1, log='all')
+ for _epoch in range(8):
+ train(net_1, device, trainloader, optimizer_1, criterion_1, _epoch)
+ test(net_1, device, testloader, criterion_1, CLASS_NAMES)
+ print('Finished Training')
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ # Visualize and debug model pipelines with W&B
+ Think of [W&B](https://www.wandb.com/) like GitHub for machine learning models — **save everything you need to debug, compare and reproduce your models** — architecture, hyperparameters, weights, model predictions, GPU usage, git commits, and even datasets — with a few lines of code.
+
+ W&B lightweight integrations work with any Python script, and all you need to do is sign up for a free W&B account to start tracking and visualizing your models.
+
+ Used by the likes of OpenAI, Lyft, Github and researchers at top machine learning labs across the world, W&B is part of the new standard of best practices for machine learning.
+
+ How W&B can help you optimize your machine learning workflows:
+
+ - [Debug](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#Free-2) model performance in real time
+ - Automatically tracked [GPU, CPU usage](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#System-4) and other system metrics
+ - Powerful [custom charts](https://wandb.ai/wandb/customizable-charts/reports/Powerful-Custom-Charts-To-Debug-Model-Peformance--VmlldzoyNzY4ODI)
+ - [Share model insights](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#Share-8) interactively
+ - Efficient [hyperparameter optimization](https://docs.wandb.com/sweeps)
+ - Dataset and model [pipeline tracking](https://docs.wandb.com/artifacts) and production model management
+
+ **W&B is free for individuals, academics and open source projects.**
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-fine-tuning-a-transformer-with-pytorch-lightning/pytorch_lightning_fine_tuning_a_transformer_with_pytorch_lightning.py b/marimo/convert/pytorch-lightning-fine-tuning-a-transformer-with-pytorch-lightning/pytorch_lightning_fine_tuning_a_transformer_with_pytorch_lightning.py
new file mode 100644
index 00000000..6ad62d2b
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-fine-tuning-a-transformer-with-pytorch-lightning/pytorch_lightning_fine_tuning_a_transformer_with_pytorch_lightning.py
@@ -0,0 +1,465 @@
+# /// script
+# dependencies = ["lightning", "pandas", "torch", "transformers", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Train a Model to Check Your Grammar Using W&B, PyTorch Lightning ⚡, and 🤗
+
+ *Based on Ayush Chaurasia's awesome [W&B report](https://wandb.ai/cayush/bert-finetuning/reports/Sentence-Classification-With-Huggingface-BERT-and-W-B--Vmlldzo4MDMwNA) and [colab](https://colab.research.google.com/drive/1SQ-FOgji8AiyrQ08sIVfDiA8OUw4bC12?usp=sharing) which performs the same task using BERT, vanilla PyTorch, and W&B.*
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this notebook, we are going to train a model to detect ungrammatical sentences from the CoLA dataset. To perform the classification, we will be using Pytorch Lightning ⚡ to fine tune [DistilBERT](https://arxiv.org/abs/1910.01108), a transformer model from huggingface 🤗.
+
+ We'll use Weights & Biases to:
+ - Version our model inputs and outputs using [W&B Artifacts](https://docs.wandb.ai/guides/artifacts), including preprocessing steps, train/validation splits, and model checkpoints
+ - Log and visualize training and validation performance using [W&B's Pytorch Lightning integration](https://docs.wandb.ai/guides/integrations/lightning)
+ - Visualize and explore the raw dataset using [W&B Tables](https://docs.wandb.ai/guides/data-vis)
+ - Orchestrate a hyperparameter search using [W&B Sweeps](https://docs.wandb.ai/guides/sweeps)
+
+ Be sure to follow the links that each run outputs to your W&B workspace, where you will be able to see...
+
+ **Your model's performance metrics updating in real time**
+
+ 
+
+ **The raw data as a W&B Table, which you can sort, group, and filter**
+
+ 
+
+ **An awesome artifact graph showing our full pipeline**
+
+ 
+
+ **Interactive visualizations of how our hyperparameter choices effect model performance**
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install some dependencies
+ # packages added via marimo's package management: pandas torch lightning transformers !pip install pandas torch lightning transformers
+ # packages added via marimo's package management: wandb !pip install -Uq wandb
+ return
+
+
+@app.cell
+def _():
+ # Bulk import cell
+ import wandb
+ import random
+ import torch
+ import transformers
+ import numpy as np
+ import pandas as pd
+ import lightning.pytorch as pl
+
+ return pd, pl, torch, transformers, wandb
+
+
+@app.cell
+def _(pl):
+ # Derandomizing cell
+ pl.seed_everything(1234)
+ return
+
+
+@app.cell
+def _():
+ """
+ Note that if you are using W&B local you will need to pass the url of your W&B
+
+ For example:
+ """
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ project = "grammar-checker" # W&B project name here
+ entity = None # your W&B username or teamname here
+ return entity, project
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # The CoLA Dataset 🥤
+
+ We’ll fine tune the model on The Corpus of Linguistic Acceptability (CoLA) dataset for single sentence classification. It’s a set of sentences labeled as grammatically correct or incorrect. It was first published in May of 2018, and is one of the tests included in the “GLUE Benchmark” on which models like DistilBERT are competing.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll use a [reference artifact](https://docs.wandb.ai/guides/artifacts/references) to store a pointer to the source data. The advantages of doing this are:
+ * Any runs that use this artifact reference will be able to trace their lineage back to the true source
+ * We can use W&B to download the raw data in our code.
+
+ The cell below starts a run with job type `register-data`. In the context of this run, we:
+ 1. Create an artifact called `cola-raw`
+ 2. Add a reference to the CoLA dataset to our `cola-raw` artifact
+ 3. Log the `cola-raw` artifact to Weights & Biases.
+ """)
+ return
+
+
+@app.cell
+def _(entity, project, wandb):
+ # Enter the context of a W&B Run object, referenceable with the 'run' variable
+ with wandb.init(entity=entity, project=project, job_type='register-data') as _run:
+ data_source = wandb.Artifact('cola-raw', type='dataset')
+ data_source.add_reference('https://nyu-mll.github.io/CoLA/cola_public_1.1.zip', name='zipfile') # Construct a wandb.Artifact object
+ _run.log_artifact(data_source) # Store a reference to the download URL of the CoLA dataset # Log the artifact to W&B
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Tokenization 🪙
+
+ The cell below defines the function `tokenize_data`, which transforms a list of sentences and a list of labels into a tuple of `torch.tensor` objects which can be consumed by the transormer model we'll be using. The 3 tensors returned are the tokenized form of the sentences, the attention masks indicating which tokens in each sentence correspond to actual words, and a tensor containing the original labels.
+ """)
+ return
+
+
+@app.cell
+def _(torch, transformers):
+ def tokenize_data(sentences, labels):
+
+ # Tokenize all of the sentences and map the tokens to thier word IDs.
+ input_ids = []
+ attention_masks = []
+
+ # Get BertTokenizer from transformers
+ tokenizer = transformers.BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)
+
+ # For every sentence...
+ for sent in sentences:
+
+ # `encode_plus` will:
+ # (1) Tokenize the sentence.
+ # (2) Prepend the `[CLS]` token to the start.
+ # (3) Append the `[SEP]` token to the end.
+ # (4) Map tokens to their IDs.
+ # (5) Pad or truncate the sentence to `max_length`
+ # (6) Create attention masks for [PAD] tokens.
+ encoded_dict = tokenizer.encode_plus(
+ sent, # Sentence to encode.
+ add_special_tokens = True, # Add '[CLS]' and '[SEP]'
+ max_length = 64, # Pad & truncate all sentences.
+ padding='max_length',
+ return_attention_mask = True, # Construct attn. masks.
+ return_tensors = 'pt', # Return pytorch tensors.
+ )
+
+ # Add the encoded sentence to the list.
+ input_ids.append(encoded_dict['input_ids'])
+
+ # And its attention mask (simply differentiates padding from non-padding).
+ attention_masks.append(encoded_dict['attention_mask'])
+
+ # Convert the lists into tensors.
+ input_ids = torch.cat(input_ids, dim=0)
+ attention_masks = torch.cat(attention_masks, dim=0)
+ labels = torch.tensor(labels)
+ return input_ids, attention_masks, labels
+
+ return (tokenize_data,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The code below executes a run of type `preprocess-data`, which will
+ 1. Download the CoLA dataset using the reference artifact we logged previously
+ 2. Log the entire dataset to W&B as a Table
+ 3. Use the function `tokenize_data` to transform each sentence into a sequence of tokens and an attention mask
+ 4. Log the preprocessed data as an artifact to W&B.
+ """)
+ return
+
+
+@app.cell
+def _(entity, pd, project, subprocess, tokenize_data, torch, wandb):
+ with wandb.init(entity=entity, project=project, job_type='preprocess-data') as _run:
+ raw_data_artifact = _run.use_artifact('cola-raw:latest')
+ zip_path = raw_data_artifact.get_entry('zipfile').download() # Download the raw cola data from the 'zipfile' reference we added to the cola-raw artifact.
+ subprocess.call(['unzip', '-o', '$zip_path'])
+ df = pd.read_csv('./cola_public/raw/in_domain_train.tsv', delimiter='\t', header=None, names=['sentence_source', 'label', 'label_notes', 'sentence'])
+ _run.log({'raw-data': wandb.Table(dataframe=df)}) #! unzip -o $zip_path
+ input_ids, attention_masks, labels = tokenize_data(df.sentence.values, df.label.values)
+ preprocessed_data = torch.utils.data.TensorDataset(input_ids, attention_masks, labels) # jupyter hack to unzip data :P
+ data_artifact = wandb.Artifact('preprocessed-data', type='dataset')
+ with open('preprocessed-data.pt', 'wb') as f: # Read in the raw data, log it to W&B as a wandb.Table
+ torch.save(preprocessed_data, f)
+ data_artifact.add_file('preprocessed-data.pt', name='dataset')
+ _run.log_artifact(data_artifact) # Perform tokenization and store as a TensorDataset # 1. Create an artifact called preprocessed-data # 2. Save the dataset to a local fil called preprocessed-data.pt # 3. Add that file to the preprocessed-data artifact # 4. Log the artifact to W&B
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Splitting Our Data 🪓
+
+ For our training process, we want to split the data into a train and validation set. The train set is the data we will use to update the model parameters, while the validation set will be a smaller segment of data that we use to test whether our model is generalizing to examples that it hasn't been trained on.
+
+ The cell below executes a `wandb.Run` with `job_type="split-data"`. In the context of this run we will:
+
+ 1. Download the `preprocessed-data` artifact logged by our previous run
+ 2. Use the `random_split` function from `torch` to perform a randomn 90/10 test/valiation split on the preprocessed data
+ 3. Store the split datasets in a new artifact called `split-dataset`
+ """)
+ return
+
+
+@app.cell
+def _(entity, project, torch, wandb):
+ with wandb.init(entity=entity, project=project, job_type='split-data') as _run:
+ pp_data_artifact = _run.use_artifact('preprocessed-data:latest')
+ data_path = pp_data_artifact.get_entry('dataset').download() # Download the preprocessed data
+ dataset = torch.load(data_path, weights_only=False)
+ train_size = int(0.9 * len(dataset))
+ val_size = len(dataset) - train_size
+ train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
+ split_data_artifact = wandb.Artifact('split-dataset', type='dataset') # Calculate the number of samples to include in each set.
+ torch.save(train_dataset, 'train.pt')
+ torch.save(val_dataset, 'validation.pt')
+ split_data_artifact.add_file('train.pt', name='train-data')
+ split_data_artifact.add_file('validation.pt', name='validation-data') # Divide the dataset by randomly selecting samples.
+ _run.log_artifact(split_data_artifact) # Construct a new artifact # Save the dataset splits to disk # Add the data splits to the artifact # Log the artifact to W&B
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Defining Our Model ⚡
+
+ We define our model and the associated training + validation procedures in the `LightningModule` below. The model itself is a pre-trained `DistilBertForSequenceClassification` with two labels.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from torch.optim import AdamW
+
+ return (AdamW,)
+
+
+@app.cell
+def _(AdamW, pl, torch, transformers):
+ class SentenceClassifier(pl.LightningModule):
+
+ def __init__(self, learning_rate=5e-5):
+ super(SentenceClassifier, self).__init__()
+
+ # Load pretrained distilbert-base-uncased configured for classification with 2 labels
+ self.model = transformers.DistilBertForSequenceClassification.from_pretrained(
+ "distilbert-base-uncased",
+ num_labels = 2,
+ output_attentions = False, # Whether the model returns attentions weights.
+ output_hidden_states = False, # Whether the model returns all hidden-states.
+ )
+ self.learning_rate = learning_rate
+
+ def training_step(self, batch, batch_no):
+ """
+ This function overrides the pl.LightningModule class.
+
+ When trainer.fit is called, each batch from the provided data loader is fed
+ to this function successively.
+ """
+ ids, masks, labels = batch
+ outputs = self.model(ids, attention_mask=masks, labels=labels)
+ preds = torch.argmax(outputs["logits"], axis=1)
+ correct = sum(preds.flatten() == labels.flatten())
+ self.log("train/loss", outputs["loss"], on_step=True, on_epoch=True)
+ self.log("train/acc", correct/len(ids), on_step=True, on_epoch=True)
+ return outputs["loss"]
+
+ def validation_step(self, batch, batch_no):
+ """
+ """
+ ids, masks, labels = batch
+ outputs = self.model(ids, attention_mask=masks, labels=labels)
+ preds = torch.argmax(outputs["logits"], axis=1)
+ correct = sum(preds.flatten() == labels.flatten())
+ self.log("val/loss", outputs["loss"], on_step=False, on_epoch=True)
+ self.log("val/acc", correct/len(ids), on_step=False, on_epoch=True)
+
+ def configure_optimizers(self):
+ return AdamW(
+ self.model.parameters(),
+ lr=self.learning_rate,
+ eps=1e-8
+ )
+
+ return (SentenceClassifier,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Training & Tracking Our Model 📉
+
+ In the cell below, we define a function `train` which sets up and performs training in the context of a W&B run. The train function takes a configuration dictionary as input then passes it to `wandb.init` via the `config` keyword argument. We use the values saved in the `wandb.config` object associated with the run to set the parameters of our trainer and data loaders. This is a crucial best practice to ensure that the values logged in the `config` object (and displayed in the run table of the W&B app) represent the actual parameters of the experiment.
+ """)
+ return
+
+
+@app.cell
+def _(SentenceClassifier, entity, pl, project, torch, wandb):
+ def train(config={'learning_rate': 5e-05, 'batch_size': 16, 'epochs': 2}):
+ with wandb.init(project=project, entity=entity, job_type='train', config=config) as _run:
+ data = _run.use_artifact('split-dataset:latest')
+ train_dataset = torch.load(data.get_entry('train-data').download(), weights_only=False)
+ val_dataset = torch.load(data.get_entry('validation-data').download(), weights_only=False) # Load the datasets from the split-dataset artifact
+ config = _run.config
+ model = SentenceClassifier(learning_rate=config.learning_rate)
+ logger = pl.loggers.WandbLogger(experiment=_run, log_model=True)
+ gpus = -1 if torch.cuda.is_available() else 0
+ trainer = pl.Trainer(max_epochs=config.epochs, logger=logger) # Extract the config object associated with the run
+ train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config.batch_size)
+ val_data_loader = torch.utils.data.DataLoader(val_dataset, batch_size=config.batch_size)
+ trainer.fit(model, train_data_loader, val_data_loader) # Construct our LightningModule with the learning rate from the config object # This logger is used when we call self.log inside the LightningModule # Use as many GPUs as are available # Construct a Trainer object with the W&B logger we created and epoch set by the config object # Build data loaders for our datasets, using the batch_size from our config object # Execute training
+
+ return (train,)
+
+
+@app.cell
+def _(train):
+ train() # Run training with default parameters
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Running a Hyperparameter Sweep 🧹
+
+ W&B sweeps allow you to optimize your model hyperparameters with minimal effort. In general, the workflow of sweeps is:
+ 1. Construct a dictionary or YAML file that defines the hyperparameter space
+ 2. Call `wandb.sweep()` from the python library or `wandb sweep ` from the command line to initialize the sweep in W&B
+ 3. Run `wandb.agent()` (python lib) or `wandb agent ` (cli) to start a sweep agent to continuously:
+ - pull hyperparameter combinations from W&B
+ - run training with the given hyperparameters
+ - log training metrics back to W&B
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We implement the sweeps workflow laid out above by:
+ 1. Creating a `sweep_config` dictionary describing our hyperparameter space and objective
+ - The hyperparameters we will sweep over are `learning_rate`, `batch_size`, and `epochs`
+ - Our objective in this sweep is to maximize the `validation/epoch_acc` metric logged to W&B
+ - We will use the `random` strategy, which means we will sample uniformly from the parameter space indefinitely
+ 2. Calling `wandb.sweep(sweep_config)` to create the sweep in our W&B project
+ - `wandb.sweep` will return a unique id for the sweep, saved as `sweep_id`
+ 3. Calling `wandb.agent(sweep_id, function=train)` to create an agent that will execute training with different hyperparameter combinations
+ - The agent will repeatedly query W&B for hyperparameter combinations
+ - When `wandb.init` is called within an agent, the `config` dictionary of the returned `run` will be populated with the next hyperparameter combination in the sweep
+ """)
+ return
+
+
+@app.cell
+def _():
+ sweep_config = {
+ 'method': 'random', # Randomly sample the hyperparameter space (alternatives: grid, bayes)
+ 'metric': { # This is the metric we are interested in maximizing
+ 'name': 'validation/epoch_acc',
+ 'goal': 'maximize'
+ },
+ # Paramters and parameter values we are sweeping across
+ 'parameters': {
+ 'learning_rate': {
+ 'values': [5e-5, 3e-5, 2e-5]
+ },
+ 'batch_size': {
+ 'values': [16, 32]
+ },
+ 'epochs':{
+ 'values': [1, 2]
+ }
+ }
+ }
+ return (sweep_config,)
+
+
+@app.cell
+def _(entity, project, sweep_config, wandb):
+ # Create the sweep
+ sweep_id = wandb.sweep(sweep_config, project=project, entity=entity)
+ return (sweep_id,)
+
+
+@app.cell
+def _(sweep_id, train, wandb):
+ # Run an agent 🕵️ to try out 5 hyperparameter combinations
+ wandb.agent(sweep_id, function=train, count=5)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-image-classification-using-pytorch-lightning/pytorch_lightning_image_classification_using_pytorch_lightning.py b/marimo/convert/pytorch-lightning-image-classification-using-pytorch-lightning/pytorch_lightning_image_classification_using_pytorch_lightning.py
new file mode 100644
index 00000000..81ac7b3c
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-image-classification-using-pytorch-lightning/pytorch_lightning_image_classification_using_pytorch_lightning.py
@@ -0,0 +1,458 @@
+# /// script
+# dependencies = ["lightning", "torchvision", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Image Classification using PyTorch Lightning ⚡️
+
+ We will build an image classification pipeline using PyTorch Lightning. We will follow this [style guide](https://lightning.ai/docs/pytorch/stable/starter/style_guide.html) to increase the readability and reproducibility of our code. A cool explanation of this available [here](https://wandb.ai/wandb/wandb-lightning/reports/Image-Classification-using-PyTorch-Lightning--VmlldzoyODk1NzY).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setting up PyTorch Lightning and W&B
+
+ For this tutorial, we need PyTorch Lightning(ain't that obvious!) and Weights and Biases.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: lightning torchvision !pip install lightning torchvision -q
+ # install weights and biases
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You're gonna need these imports.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import lightning.pytorch as pl
+ # your favorite machine learning tracking tool
+ from lightning.pytorch.loggers import WandbLogger
+
+ import torch
+ from torch import nn
+ from torch.nn import functional as F
+ from torch.utils.data import random_split, DataLoader
+
+ from torchmetrics import Accuracy
+
+ from torchvision import transforms
+ from torchvision.datasets import CIFAR10
+
+ import wandb
+
+ return (
+ Accuracy,
+ CIFAR10,
+ DataLoader,
+ F,
+ WandbLogger,
+ nn,
+ pl,
+ random_split,
+ torch,
+ transforms,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now you'll need to login to you wandb account.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🔧 DataModule - The Data Pipeline we Deserve
+
+ DataModules are a way of decoupling data-related hooks from the LightningModule so you can develop dataset agnostic models.
+
+ It organizes the data pipeline into one shareable and reusable class. A datamodule encapsulates the five steps involved in data processing in PyTorch:
+ - Download / tokenize / process.
+ - Clean and (maybe) save to disk.
+ - Load inside Dataset.
+ - Apply transforms (rotate, tokenize, etc…).
+ - Wrap inside a DataLoader.
+
+ Learn more about datamodules [here](https://lightning.ai/docs/pytorch/stable/data/datamodule.html). Let's build a datamodule for the Cifar-10 dataset.
+ """)
+ return
+
+
+@app.cell
+def _(CIFAR10, DataLoader, pl, random_split, transforms):
+ class CIFAR10DataModule(pl.LightningDataModule):
+ def __init__(self, batch_size, data_dir: str = './'):
+ super().__init__()
+ self.data_dir = data_dir
+ self.batch_size = batch_size
+
+ self.transform = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+ ])
+
+ self.num_classes = 10
+
+ def prepare_data(self):
+ CIFAR10(self.data_dir, train=True, download=True)
+ CIFAR10(self.data_dir, train=False, download=True)
+
+ def setup(self, stage=None):
+ # Assign train/val datasets for use in dataloaders
+ if stage == 'fit' or stage is None:
+ cifar_full = CIFAR10(self.data_dir, train=True, transform=self.transform)
+ self.cifar_train, self.cifar_val = random_split(cifar_full, [45000, 5000])
+
+ # Assign test dataset for use in dataloader(s)
+ if stage == 'test' or stage is None:
+ self.cifar_test = CIFAR10(self.data_dir, train=False, transform=self.transform)
+
+ def train_dataloader(self):
+ return DataLoader(self.cifar_train, batch_size=self.batch_size, shuffle=True)
+
+ def val_dataloader(self):
+ return DataLoader(self.cifar_val, batch_size=self.batch_size)
+
+ def test_dataloader(self):
+ return DataLoader(self.cifar_test, batch_size=self.batch_size)
+
+ return (CIFAR10DataModule,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📱 Callbacks
+
+ A callback is a self-contained program that can be reused across projects. PyTorch Lightning comes with few [built-in callbacks](https://lightning.ai/docs/pytorch/latest/extensions/callbacks.html#built-in-callbacks) which are regularly used.
+ Learn more about callbacks in PyTorch Lightning [here](https://lightning.ai/docs/pytorch/latest/extensions/callbacks.html).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Built-in Callbacks
+
+ In this tutorial, we will use [Early Stopping](https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.EarlyStopping.html#lightning.callbacks.EarlyStopping) and [Model Checkpoint](https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.ModelCheckpoint.html#pytorch_lightning.callbacks.ModelCheckpoint) built-in callbacks. They can be passed to the `Trainer`.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Custom Callbacks
+ If you are familiar with Custom Keras callback, the ability to do the same in your PyTorch pipeline is just a cherry on the cake.
+
+ Since we are performing image classification, the ability to visualize the model's predictions on some samples of images can be helpful. This in the form of a callback can help debug the model at an early stage.
+ """)
+ return
+
+
+@app.cell
+def _(pl, torch, wandb):
+ class ImagePredictionLogger(pl.callbacks.Callback):
+ def __init__(self, val_samples, num_samples=32):
+ super().__init__()
+ self.num_samples = num_samples
+ self.val_imgs, self.val_labels = val_samples
+
+ def on_validation_epoch_end(self, trainer, pl_module):
+ # Bring the tensors to CPU
+ val_imgs = self.val_imgs.to(device=pl_module.device)
+ val_labels = self.val_labels.to(device=pl_module.device)
+ # Get model prediction
+ logits = pl_module(val_imgs)
+ preds = torch.argmax(logits, -1)
+ # Log the images as wandb Image
+ trainer.logger.experiment.log({
+ "examples":[wandb.Image(x, caption=f"Pred:{pred}, Label:{y}")
+ for x, pred, y in zip(val_imgs[:self.num_samples],
+ preds[:self.num_samples],
+ val_labels[:self.num_samples])]
+ })
+
+ return (ImagePredictionLogger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🎺 LightningModule - Define the System
+
+ The LightningModule defines a system and not a model. Here a system groups all the research code into a single class to make it self-contained. `LightningModule` organizes your PyTorch code into 5 sections:
+ - Computations (`__init__`).
+ - Train loop (`training_step`)
+ - Validation loop (`validation_step`)
+ - Test loop (`test_step`)
+ - Optimizers (`configure_optimizers`)
+
+ One can thus build a dataset agnostic model that can be easily shared. Let's build a system for Cifar-10 classification.
+ """)
+ return
+
+
+@app.cell
+def _(Accuracy, F, nn, pl, torch):
+ class LitModel(pl.LightningModule):
+ def __init__(self, input_shape, num_classes, learning_rate=2e-4):
+ super().__init__()
+
+ # log hyperparameters
+ self.save_hyperparameters()
+ self.learning_rate = learning_rate
+
+ self.conv1 = nn.Conv2d(3, 32, 3, 1)
+ self.conv2 = nn.Conv2d(32, 32, 3, 1)
+ self.conv3 = nn.Conv2d(32, 64, 3, 1)
+ self.conv4 = nn.Conv2d(64, 64, 3, 1)
+
+ self.pool1 = torch.nn.MaxPool2d(2)
+ self.pool2 = torch.nn.MaxPool2d(2)
+
+ n_sizes = self._get_conv_output(input_shape)
+
+ self.fc1 = nn.Linear(n_sizes, 512)
+ self.fc2 = nn.Linear(512, 128)
+ self.fc3 = nn.Linear(128, num_classes)
+
+ self.accuracy = Accuracy(task="multiclass", num_classes=num_classes)
+
+ # returns the size of the output tensor going into Linear layer from the conv block.
+ def _get_conv_output(self, shape):
+ batch_size = 1
+ input = torch.autograd.Variable(torch.rand(batch_size, *shape))
+
+ output_feat = self._forward_features(input)
+ n_size = output_feat.data.view(batch_size, -1).size(1)
+ return n_size
+
+ # returns the feature tensor from the conv block
+ def _forward_features(self, x):
+ x = F.relu(self.conv1(x))
+ x = self.pool1(F.relu(self.conv2(x)))
+ x = F.relu(self.conv3(x))
+ x = self.pool2(F.relu(self.conv4(x)))
+ return x
+
+ # will be used during inference
+ def forward(self, x):
+ x = self._forward_features(x)
+ x = x.view(x.size(0), -1)
+ x = F.relu(self.fc1(x))
+ x = F.relu(self.fc2(x))
+ x = F.log_softmax(self.fc3(x), dim=1)
+
+ return x
+
+ def training_step(self, batch, batch_idx):
+ x, y = batch
+ logits = self(x)
+ loss = F.nll_loss(logits, y)
+
+ # training metrics
+ preds = torch.argmax(logits, dim=1)
+ acc = self.accuracy(preds, y)
+ self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True)
+ self.log('train_acc', acc, on_step=True, on_epoch=True, logger=True)
+
+ return loss
+
+ def validation_step(self, batch, batch_idx):
+ x, y = batch
+ logits = self(x)
+ loss = F.nll_loss(logits, y)
+
+ # validation metrics
+ preds = torch.argmax(logits, dim=1)
+ acc = self.accuracy(preds, y)
+ self.log('val_loss', loss, prog_bar=True)
+ self.log('val_acc', acc, prog_bar=True)
+ return loss
+
+ def test_step(self, batch, batch_idx):
+ x, y = batch
+ logits = self(x)
+ loss = F.nll_loss(logits, y)
+
+ # validation metrics
+ preds = torch.argmax(logits, dim=1)
+ acc = self.accuracy(preds, y)
+ self.log('test_loss', loss, prog_bar=True)
+ self.log('test_acc', acc, prog_bar=True)
+ return loss
+
+ def configure_optimizers(self):
+ optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+ return optimizer
+
+ return (LitModel,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🚋 Train and Evaluate
+
+ Now that we have organized our data pipeline using `DataModule` and model architecture+training loop using `LightningModule`, the PyTorch Lightning `Trainer` automates everything else for us.
+
+ The Trainer automates:
+ - Epoch and batch iteration
+ - Calling of `optimizer.step()`, `backward`, `zero_grad()`
+ - Calling of `.eval()`, enabling/disabling grads
+ - Saving and loading weights
+ - Weights and Biases logging
+ - Multi-GPU training support
+ - TPU support
+ - 16-bit training support
+ """)
+ return
+
+
+@app.cell
+def _(CIFAR10DataModule):
+ dm = CIFAR10DataModule(batch_size=32)
+ # To access the x_dataloader we need to call prepare_data and setup.
+ dm.prepare_data()
+ dm.setup()
+
+ # Samples required by the custom ImagePredictionLogger callback to log image predictions.
+ val_samples = next(iter(dm.val_dataloader()))
+ val_imgs, val_labels = val_samples[0], val_samples[1]
+ val_imgs.shape, val_labels.shape
+ return dm, val_samples
+
+
+@app.cell
+def _(
+ ImagePredictionLogger,
+ LitModel,
+ WandbLogger,
+ dm,
+ pl,
+ val_samples,
+ wandb,
+):
+ model = LitModel((3, 32, 32), dm.num_classes)
+
+ # Initialize wandb logger
+ wandb_logger = WandbLogger(project='wandb-lightning', job_type='train')
+
+ # Initialize Callbacks
+ early_stop_callback = pl.callbacks.EarlyStopping(monitor="val_loss")
+ checkpoint_callback = pl.callbacks.ModelCheckpoint()
+
+ # Initialize a trainer
+ trainer = pl.Trainer(max_epochs=2,
+ logger=wandb_logger,
+ callbacks=[early_stop_callback,
+ ImagePredictionLogger(val_samples),
+ checkpoint_callback],
+ )
+
+ # Train the model ⚡🚅⚡
+ trainer.fit(model, dm)
+
+ # Evaluate the model on the held-out test set ⚡⚡
+ trainer.test(dataloaders=dm.test_dataloader())
+
+ # Close wandb run
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Final Thoughts
+ I come from the TensorFlow/Keras ecosystem and find PyTorch a bit overwhelming even though it's an elegant framework. Just my personal experience though. While exploring PyTorch Lightning, I realized that almost all of the reasons that kept me away from PyTorch is taken care of. Here's a quick summary of my excitement:
+ - Then: Conventional PyTorch model definition used to be all over the place. With the model in some `model.py` script and the training loop in the `train.py `file. It was a lot of looking back and forth to understand the pipeline.
+ - Now: The `LightningModule` acts as a system where the model is defined along with the `training_step`, `validation_step`, etc. Now it's modular and shareable.
+ - Then: The best part about TensorFlow/Keras is the input data pipeline. Their dataset catalog is rich and growing. PyTorch's data pipeline used to be the biggest pain point. In normal PyTorch code, the data download/cleaning/preparation is usually scattered across many files.
+ - Now: The DataModule organizes the data pipeline into one shareable and reusable class. It's simply a collection of a `train_dataloader`, `val_dataloader`(s), `test_dataloader`(s) along with the matching transforms and data processing/downloads steps required.
+ - Then: With Keras, one can call `model.fit` to train the model and `model.predict` to run inference on. `model.evaluate` offered a good old simple evaluation on the test data. This is not the case with PyTorch. One will usually find separate `train.py` and `test.py` files.
+ - Now: With the `LightningModule` in place, the `Trainer` automates everything. One needs to just call `trainer.fit` and `trainer.test` to train and evaluate the model.
+ - Then: TensorFlow loves TPU, PyTorch...well!
+ - Now: With PyTorch Lightning, it's so easy to train the same model with multiple GPUs and even on TPU. Wow!
+ - Then: I am a big fan of Callbacks and prefer writing custom callbacks. Something as trivial as Early Stopping used to be a point of discussion with conventional PyTorch.
+ - Now: With PyTorch Lightning using Early Stopping and Model Checkpointing is a piece of cake. I can even write custom callbacks.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🎨 Conclusion and Resources
+
+ I hope you find this report helpful. I will encourage to play with the code and train an image classifier with a dataset of your choice.
+
+ Here are some resources to learn more about PyTorch Lightning:
+ - [Step-by-step walk-through](https://lightning.ai/docs/pytorch/latest/starter/introduction.html) - This is one of the official tutorials. Their documentation is really well written and I highly encourage it as a good learning resource.
+ - [Use Pytorch Lightning with Weights & Biases](https://wandb.me/lightning) - This is a quick colab that you can run through to learn more about how to use W&B with PyTorch Lightning.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-optimize-pytorch-lightning-models-with-weights-biases/pytorch_lightning_optimize_pytorch_lightning_models_with_weights_biases.py b/marimo/convert/pytorch-lightning-optimize-pytorch-lightning-models-with-weights-biases/pytorch_lightning_optimize_pytorch_lightning_models_with_weights_biases.py
new file mode 100644
index 00000000..1f1784c1
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-optimize-pytorch-lightning-models-with-weights-biases/pytorch_lightning_optimize_pytorch_lightning_models_with_weights_biases.py
@@ -0,0 +1,444 @@
+# /// script
+# dependencies = ["lightning", "torchvision", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # ⚡ Pytorch Lightning models with Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Pytorch Lightning is a lightweight wrapper for organizing your PyTorch code and easily adding advanced features such as distributed training, 16-bit precision or gradient accumulation.
+
+ Coupled with the [Weights & Biases integration](https://docs.wandb.com/library/integrations/lightning), you can quickly train and monitor models for full traceability and reproducibility with only 2 extra lines of code:
+
+ ```python
+ from lightning.pytorch.loggers import WandbLogger
+ from lightning.pytorch import Trainer
+
+ wandb_logger = WandbLogger()
+ trainer = Trainer(logger=wandb_logger)
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ W&B integration with Pytorch-Lightning can automatically:
+ * log your configuration parameters
+ * log your losses and metrics
+ * log your model
+ * keep track of your code
+ * log your system metrics (GPU, CPU, memory, temperature, etc)
+
+ ### 📚 Docs
+ You can find the PyTorch Lightning WandbLogger docs [here](https://pytorch-lightning.readthedocs.io/en/latest/extensions/generated/pytorch_lightning.loggers.WandbLogger.html?highlight=wandblogger) and the Weights & Biases docs [here](https://docs.wandb.com/library/integrations/lightning)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🛠️ Installation and set-up
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: lightning wandb torchvision !pip install -q lightning wandb torchvision
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We make sure we're logged into W&B so that our experiments can be associated with our account.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📊 Setting up the dataloader
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For the context of this tutorial we use vanilla pytorch dataloaders on the MNIST dataset
+ """)
+ return
+
+
+@app.cell
+def _():
+ from torchvision.datasets import MNIST
+ from torchvision import transforms
+ from torch.utils.data import DataLoader, random_split
+
+ transform = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize((0.1307,), (0.3081,))])
+
+ dataset = MNIST(root="./MNIST", download=True, transform=transform)
+ training_set, validation_set = random_split(dataset, [55000, 5000])
+ return DataLoader, training_set, validation_set
+
+
+@app.cell
+def _(DataLoader, training_set, validation_set):
+ training_loader = DataLoader(training_set, batch_size=64, shuffle=True)
+ validation_loader = DataLoader(validation_set, batch_size=64)
+ return training_loader, validation_loader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🤓 Defining the Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Tips**:
+ * Call `self.save_hyperparameters()` in `__init__` to automatically log your hyperparameters to **W&B**
+ * Call self.log in `training_step` and `validation_step` to log the metrics
+ """)
+ return
+
+
+@app.cell
+def _():
+ import lightning.pytorch as pl
+
+ return (pl,)
+
+
+@app.cell
+def _(pl):
+ import torch
+ from torch.nn import Linear, CrossEntropyLoss, functional as F
+ from torch.optim import Adam
+ from torchmetrics.functional import accuracy
+
+ class MNIST_LitModule(pl.LightningModule):
+
+ def __init__(self, n_classes=10, n_layer_1=128, n_layer_2=256, lr=1e-3):
+ '''method used to define our model parameters'''
+ super().__init__()
+
+ # mnist images are (1, 28, 28) (channels, width, height)
+ self.layer_1 = Linear(28 * 28, n_layer_1)
+ self.layer_2 = Linear(n_layer_1, n_layer_2)
+ self.layer_3 = Linear(n_layer_2, n_classes)
+
+ # loss
+ self.loss = CrossEntropyLoss()
+
+ # optimizer parameters
+ self.lr = lr
+
+ # save hyper-parameters to self.hparams (auto-logged by W&B)
+ self.save_hyperparameters()
+
+ def forward(self, x):
+ '''method used for inference input -> output'''
+
+ batch_size, channels, width, height = x.size()
+
+ # (b, 1, 28, 28) -> (b, 1*28*28)
+ x = x.view(batch_size, -1)
+
+ # let's do 3 x (linear + relu)
+ x = self.layer_1(x)
+ x = F.relu(x)
+ x = self.layer_2(x)
+ x = F.relu(x)
+ x = self.layer_3(x)
+
+ return x
+
+ def training_step(self, batch, batch_idx):
+ '''needs to return a loss from a single batch'''
+ _, loss, acc = self._get_preds_loss_accuracy(batch)
+
+ # Log loss and metric
+ self.log('train_loss', loss)
+ self.log('train_accuracy', acc)
+
+ return loss
+
+ def validation_step(self, batch, batch_idx):
+ '''used for logging metrics'''
+ preds, loss, acc = self._get_preds_loss_accuracy(batch)
+
+ # Log loss and metric
+ self.log('val_loss', loss)
+ self.log('val_accuracy', acc)
+
+ # Let's return preds to use it in a custom callback
+ return preds
+
+ def test_step(self, batch, batch_idx):
+ '''used for logging metrics'''
+ _, loss, acc = self._get_preds_loss_accuracy(batch)
+
+ # Log loss and metric
+ self.log('test_loss', loss)
+ self.log('test_accuracy', acc)
+
+ def configure_optimizers(self):
+ '''defines model optimizer'''
+ return Adam(self.parameters(), lr=self.lr)
+
+ def _get_preds_loss_accuracy(self, batch):
+ '''convenience function since train/valid/test steps are similar'''
+ x, y = batch
+ logits = self(x)
+ preds = torch.argmax(logits, dim=1)
+ loss = self.loss(logits, y)
+ acc = accuracy(preds, y, 'multiclass', num_classes=10)
+ return preds, loss, acc
+
+ return (MNIST_LitModule,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The model is now ready!
+ """)
+ return
+
+
+@app.cell
+def _(MNIST_LitModule):
+ model = MNIST_LitModule(n_layer_1=128, n_layer_2=128)
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💾 Save Model Checkpoints
+
+ The `ModelCheckpoint` callback is required along with the `WandbLogger` argument to log model checkpoints to W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from lightning.pytorch.callbacks import ModelCheckpoint
+
+ checkpoint_callback = ModelCheckpoint(monitor='val_accuracy', mode='max')
+ return (checkpoint_callback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💡 Tracking Experiments with WandbLogger
+
+ PyTorch Lightning has a `WandbLogger` to easily log your experiments with Wights & Biases. Just pass it to your `Trainer` to log to W&B. See the [WandbLogger docs](https://lightning.ai/docs/pytorch/stable/extensions/generated/pytorch_lightning.loggers.WandbLogger.html#pytorch_lightning.loggers.WandbLogger) for all parameters. Note, to log the metrics to a specific W&B Team, pass your Team name to the `entity` argument in `WandbLogger`
+
+ #### `lightning.pytorch.loggers.WandbLogger()`
+
+ | Functionality | Argument/Function | PS |
+ | ------ | ------ | ------ |
+ | Logging models | `WandbLogger(... ,log_model='all')` or `WandbLogger(... ,log_model=True`) | Log all models if `log_model="all"` and at end of training if `log_model=True`
+ | Set custom run names | `WandbLogger(... ,name='my_run_name'`) | |
+ | Organize runs by project | `WandbLogger(... ,project='my_project')` | |
+ | Log histograms of gradients and parameters | `WandbLogger.watch(model)` | `WandbLogger.watch(model, log='all')` to log parameter histograms |
+ | Log hyperparameters | Call `self.save_hyperparameters()` within `LightningModule.__init__()` |
+ | Log custom objects (images, audio, video, molecules…) | Use `WandbLogger.log_text`, `WandbLogger.log_image` and `WandbLogger.log_table`, etc. |
+
+ See the [WandbLogger docs](https://lightning.ai/docs/pytorch/stable/extensions/generated/pytorch_lightning.loggers.WandbLogger.html#pytorch_lightning.loggers.WandbLogger) here for all parameters.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from lightning.pytorch.loggers import WandbLogger
+ from lightning.pytorch import Trainer
+
+ wandb_logger = WandbLogger(project='MNIST', # group runs in "MNIST" project
+ log_model='all') # log all new checkpoints during training
+ return Trainer, wandb_logger
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ⚙️ Using WandbLogger to log Images, Text and More
+ Pytorch Lightning is extensible through its callback system. We can create a custom callback to automatically log sample predictions during validation. `WandbLogger` provides convenient media logging functions:
+ * `WandbLogger.log_text` for text data
+ * `WandbLogger.log_image` for images
+ * `WandbLogger.log_table` for [W&B Tables](https://docs.wandb.ai/guides/data-vis).
+
+ An alternate to `self.log` in the Model class is directly using `wandb.log({dict})` or `trainer.logger.experiment.log({dict})`
+
+ In this case we log the first 20 images in the first batch of the validation dataset along with the predicted and ground truth labels.
+ """)
+ return
+
+
+@app.cell
+def _(wandb, wandb_logger):
+ from lightning.pytorch.callbacks import Callback
+
+ class LogPredictionsCallback(Callback):
+
+ def on_validation_batch_end(
+ self, trainer, pl_module, outputs, batch, batch_idx):
+ """Called when the validation batch ends."""
+
+ # `outputs` comes from `LightningModule.validation_step`
+ # which corresponds to our model predictions in this case
+
+ # Let's log 20 sample image predictions from first batch
+ if batch_idx == 0:
+ n = 20
+ x, y = batch
+ images = [img for img in x[:n]]
+ captions = [f'Ground Truth: {y_i} - Prediction: {y_pred}' for y_i, y_pred in zip(y[:n], outputs[:n])]
+
+ # Option 1: log images with `WandbLogger.log_image`
+ wandb_logger.log_image(key='sample_images', images=images, caption=captions)
+
+ # Option 2: log predictions as a Table
+ columns = ['image', 'ground truth', 'prediction']
+ data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs[:n]))]
+ wandb_logger.log_table(key='sample_table', columns=columns, data=data)
+
+ log_predictions_callback = LogPredictionsCallback()
+ return (log_predictions_callback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🏋️ Train Your Model
+ """)
+ return
+
+
+@app.cell
+def _(Trainer, checkpoint_callback, log_predictions_callback, wandb_logger):
+ trainer = Trainer(
+ logger=wandb_logger, # W&B integration
+ callbacks=[log_predictions_callback, # logging of sample predictions
+ checkpoint_callback], # our model checkpoint callback
+ accelerator="gpu", # use GPU
+ max_epochs=5) # number of epochs
+ return (trainer,)
+
+
+@app.cell
+def _(model, trainer, training_loader, validation_loader):
+ trainer.fit(model, training_loader, validation_loader)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ When we want to close our W&B run, we call `wandb.finish()` (mainly useful in notebooks, called automatically in scripts).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can monitor losses, metrics, gradients, parameters and sample predictions as the model trains.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📚 Resources
+
+ * [Pytorch Lightning and W&B integration documentation](https://docs.wandb.ai/integrations/lightning) contains a few tips for taking most advantage of W&B
+ * [Pytorch Lightning documentation](https://pytorch-lightning.readthedocs.io/en/stable/common/loggers.html#weights-and-biases) 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, please contact support@wandb.com
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-profile-pytorch-code/pytorch_lightning_profile_pytorch_code.py b/marimo/convert/pytorch-lightning-profile-pytorch-code/pytorch_lightning_profile_pytorch_code.py
new file mode 100644
index 00000000..43be24ed
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-profile-pytorch-code/pytorch_lightning_profile_pytorch_code.py
@@ -0,0 +1,471 @@
+# /// script
+# dependencies = ["lightning", "torch-tb-profiler", "torchvision", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Profiling PyTorch Code
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This notebook demonstrates how to incorporate [PyTorch Kineto](https://github.com/pytorch/kineto)'s
+ [Tensorboard plugin](https://github.com/pytorch/kineto/blob/master/tb_plugin/README.md)
+ for profiling PyTorch code
+ with [PyTorch Lightning](https://pytorch-lightning.readthedocs.io/)
+ as the high-level training API
+ and
+ [Weights & Biases](https://wandb.ai/site)
+ as the logging solution.
+
+ The final result looks something like what you see below:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The work done by processes, threads, and streams on the CPU and GPU
+ is displayed along with precise timing information
+ in an interactive viewer that can be incorporated into
+ Weights & Biases
+ [workspaces](https://docs.wandb.ai/ref/app/pages/workspaces)
+ and [Reports](https://docs.wandb.ai/guides/reports)
+ or exported to external viewers.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ That means you can incorporate tracing and profiling into
+ your model training and evaluation pipeline --
+ storing, visualizing, and communicating
+ performance results alongside other key metrics and metadata,
+ like [loss curves](https://docs.wandb.ai/guides/track/log),
+ [hard examples from datasets](https://docs.wandb.ai/guides/data-vis),
+ and [hyperparameter optimization results](https://docs.wandb.ai/guides/sweeps).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _NB:_ This tool is based on the
+ [Chrome Trace Viewer](https://www.chromium.org/developers/how-tos/trace-event-profiling-tool),
+ which works best with that browser.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb lightning torch_tb_profiler torchvision !pip install -q wandb lightning torch_tb_profiler torchvision
+ return
+
+
+@app.cell
+def _():
+ import glob
+
+ import lightning.pytorch as pl
+ import torch
+ import torch.nn as nn
+ import torch.nn.functional as F
+ import torch.optim as optim
+ import torchvision
+ from torchvision import datasets, transforms
+
+ from torch.profiler import tensorboard_trace_handler
+ import wandb
+
+ # 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 (
+ F,
+ datasets,
+ glob,
+ nn,
+ optim,
+ pl,
+ tensorboard_trace_handler,
+ torch,
+ transforms,
+ wandb,
+ )
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Set Up Profiled Training
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Network Module
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To profile neural network code,
+ we first need to write it.
+
+ For this demo,
+ we'll stick with a simple
+ [LeNet](http://yann.lecun.com/exdb/lenet/)-style DNN,
+ based on the
+ [PyTorch introductory tutorial](https://pytorch.org/tutorials/recipes/recipes/defining_a_neural_network.html).
+ """)
+ return
+
+
+@app.cell
+def _(F, nn, optim, pl, torch):
+ OPTIMIZERS = {
+ "Adadelta": optim.Adadelta,
+ "Adagrad" : optim.Adagrad,
+ "SGD": optim.SGD,
+ }
+
+ class Net(pl.LightningModule):
+ """Very simple LeNet-style DNN, plus DropOut."""
+
+ def __init__(self, optimizer="Adadelta"):
+ super(Net, self).__init__()
+ self.conv1 = nn.Conv2d(1, 32, 3, 1)
+ self.conv2 = nn.Conv2d(32, 64, 3, 1)
+ self.dropout1 = nn.Dropout(0.25)
+ self.dropout2 = nn.Dropout(0.5)
+ self.fc1 = nn.Linear(9216, 128)
+ self.fc2 = nn.Linear(128, 10)
+
+ self.optimizer = self.set_optimizer(optimizer)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = F.relu(x)
+ x = self.conv2(x)
+ x = F.relu(x)
+ x = F.max_pool2d(x, 2)
+ x = self.dropout1(x)
+ x = torch.flatten(x, 1)
+ x = self.fc1(x)
+ x = F.relu(x)
+ x = self.dropout2(x)
+ x = self.fc2(x)
+ output = F.log_softmax(x, dim=1)
+ return output
+
+ def set_optimizer(self, optimizer):
+ return OPTIMIZERS[optimizer]
+
+ return (Net,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To get this module to work with PyTorch Lightning,
+ we need to define two more methods,
+ which hook into the training loop.
+
+ Check out [this tutorial video and notebook](http://wandb.me/lit-video)
+ for more on using PyTorch Lightning and W&B.
+ """)
+ return
+
+
+@app.cell
+def _(F, Net):
+ def training_step(self, batch, idx):
+ inputs, labels = batch
+ outputs = self(inputs)
+ loss = F.nll_loss(outputs, labels)
+
+ return {"loss": loss}
+
+ def configure_optimizers(self):
+ return self.optimizer(self.parameters(), lr=0.1)
+
+ Net.training_step = training_step
+ Net.configure_optimizers = configure_optimizers
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Profiler Callback
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The profiler operates a bit like a PyTorch optimizer:
+ it has a `.step` method that we need to call
+ to demarcate the code we're interested in profiling.
+
+ A single training step (forward and backward prop)
+ is both the typical target of performance optimizations
+ and already rich enough to more than fill out a profiling trace,
+ so we want to call `.step` on each step.
+
+ The cell below defines a quick-and-dirty
+ method for doing so in PyTorch Lightning using the
+ [`Callback` system](https://pytorch-lightning.readthedocs.io/en/stable/extensions/callbacks.html).
+ """)
+ return
+
+
+@app.cell
+def _(pl):
+ class TorchTensorboardProfilerCallback(pl.Callback):
+ """Quick-and-dirty Callback for invoking TensorboardProfiler during training.
+
+ For greater robustness, extend the pl.profiler.profilers.BaseProfiler. See
+ https://pytorch-lightning.readthedocs.io/en/stable/advanced/profiler.html"""
+
+ def __init__(self, profiler):
+ super().__init__()
+ self.profiler = profiler
+
+ def on_train_batch_end(self, trainer, pl_module, outputs, *args, **kwargs):
+ self.profiler.step()
+ pl_module.log_dict(outputs) # also logging the loss, while we're here
+
+ return (TorchTensorboardProfilerCallback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Run Profiled Training
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We're now ready to go!
+
+ The cell below creates a `DataLoader`
+ based on the information in the `config`uration dictionary.
+ Choices made here have a substantial impact on performance
+ and show up very markedly in the trace.
+
+ After you've run with the default values,
+ check out the creation of the `trainloader`
+ and the `trainer`
+ for comments on what these arguments do
+ and then try a few different choices out, as suggested below.
+ """)
+ return
+
+
+@app.cell
+def _(
+ Net,
+ TorchTensorboardProfilerCallback,
+ datasets,
+ glob,
+ pl,
+ tensorboard_trace_handler,
+ torch,
+ transforms,
+ wandb,
+):
+ # initial values are defaults, for all except batch_size, which has no default
+ config = {"batch_size": 32, # try log-spaced values from 1 to 50,000
+ "num_workers": 0, # try 0, 1, and 2
+ "pin_memory": False, # try False and True
+ "precision": 32, # try 16 and 32
+ "optimizer": "Adadelta", # try optim.Adadelta and optim.SGD
+ }
+
+ with wandb.init(project="trace", config=config) as run:
+
+ # Set up MNIST data
+ transform=transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize((0.1307,), (0.3081,))
+ ])
+
+ dataset = datasets.MNIST("../data", train=True, download=True,
+ transform=transform)
+
+ ## Using a raw DataLoader, rather than LightningDataModule, for greater transparency
+ trainloader = torch.utils.data.DataLoader(
+ dataset,
+ # Key performance-relevant configuration parameters:
+ ## batch_size: how many datapoints are passed through the network at once?
+ batch_size=wandb.config.batch_size,
+ # larger batch sizes are more compute efficient, up to memory constraints
+
+ ## num_workers: how many side processes to launch for dataloading (should be >0)
+ num_workers=wandb.config.num_workers,
+ # needs to be tuned given model/batch size/compute
+
+ ## pin_memory: should a fixed "pinned" memory block be allocated on the CPU?
+ pin_memory=wandb.config.pin_memory,
+ # should nearly always be True for GPU models, see https://developer.nvidia.com/blog/how-optimize-data-transfers-cuda-cc/
+ )
+
+ # Set up model
+ model = Net(optimizer=wandb.config["optimizer"])
+
+ # Set up profiler
+ wait, warmup, active, repeat = 1, 1, 2, 1
+ total_steps = (wait + warmup + active) * (1 + repeat)
+ schedule = torch.profiler.schedule(
+ wait=wait, warmup=warmup, active=active, repeat=repeat)
+ profiler = torch.profiler.profile(
+ schedule=schedule, on_trace_ready=tensorboard_trace_handler("wandb/latest-run/tbprofile"), with_stack=False)
+
+ with profiler:
+ profiler_callback = TorchTensorboardProfilerCallback(profiler)
+
+ trainer = pl.Trainer(max_epochs=1, max_steps=total_steps,
+ logger=pl.loggers.WandbLogger(log_model=True, save_code=True),
+ callbacks=[profiler_callback], precision=wandb.config.precision)
+
+ trainer.fit(model, trainloader)
+
+ profile_art = wandb.Artifact(f"trace-{wandb.run.id}", type="profile")
+ profile_art.add_file(glob.glob("wandb/latest-run/tbprofile/*.pt.trace.json")[0], "trace.pt.trace.json")
+ run.log_artifact(profile_art)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Reading Profiling Results
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Head to the Artifacts tab
+ (identified by the
+ ["stacked pucks"](https://stackoverflow.com/questions/2822650/why-is-a-database-always-represented-with-a-cylinder)
+ database icon)
+ for your W&B [run page](https://docs.wandb.ai/ref/app/pages/run-page),
+ at the URL that appears in the output of the cell above,
+ then select the artifact named `trace-`.
+ In the Files tab, select `trace.pt.trace.json`
+ to pull up the Trace Viewer.
+
+ > You can also check out an example from an earlier run
+ [here](https://wandb.ai/wandb/trace/artifacts/profile/trace-224bfvza/56c5d50902233baa7710/files/trace.pt.trace.json).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The trace shows which operations were running and when
+ in each process/thread/stream
+ on the CPU and on the GPU.
+
+ In the main thread (the one in which the Profiler Steps appear),
+ locate the following steps:
+ 1. the loading of data (hint: look for `enumerate` on the CPU, nothing on the GPU)
+ 2. the forward pass to calculate the loss (hint: look for simultaneous activity on CPU+GPU,
+ with [`aten`](https://pytorch.org/cppdocs/#aten) in the operation names)
+ 3. the backward pass to calculate the gradient of the loss (hint: look for simultaneous activity on CPU+GPU, with [`backward`](https://pytorch.org/cppdocs/#autograd) in the operation names).
+
+ If you ran with the default settings
+ (in particular, `num_workers=0`),
+ you'll notice that these steps are all run sequentially,
+ meaning that between loading one batch
+ and loading the next,
+ the `DataLoader` is effectively idling,
+ and during the loading of a batch, the GPU is idling.
+
+ Change `num_workers` in the config to `1` or `2`
+ and then re-execute the cell above.
+ You should notice a difference,
+ in particular in the fraction of time the GPU is active.
+ (Note: the `DataLoader` may even be hard to find in this case!)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For more on how to read these results, check out
+ [this W&B Report](http://wandb.me/trace-report).
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-supercharge-your-training-with-pytorch-lightning-and-weights-and-biases/pytorch_lightning_supercharge_your_training_with_pytorch_lightning_and_weights_and_biases.py b/marimo/convert/pytorch-lightning-supercharge-your-training-with-pytorch-lightning-and-weights-and-biases/pytorch_lightning_supercharge_your_training_with_pytorch_lightning_and_weights_and_biases.py
new file mode 100644
index 00000000..78be2651
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-supercharge-your-training-with-pytorch-lightning-and-weights-and-biases/pytorch_lightning_supercharge_your_training_with_pytorch_lightning_and_weights_and_biases.py
@@ -0,0 +1,863 @@
+# /// script
+# dependencies = ["lightning", "onnx", "torchmetrics", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # ⚡ 💘 🏋️♀️ Supercharge your Training with PyTorch Lightning + Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ At Weights & Biases, we love anything
+ that makes training deep learning models easier.
+ That's why we worked with the folks at PyTorch Lightning to
+ [integrate our experiment tracking tool](https://docs.wandb.com/library/integrations/lightning)
+ directly into
+ [the Lightning library](https://pytorch-lightning.readthedocs.io/en/latest/common/loggers.html#weights-and-biases).
+
+ [PyTorch Lightning](https://pytorch-lightning.readthedocs.io/en/stable/) is a lightweight wrapper for organizing your PyTorch code and easily adding advanced features such as distributed training and 16-bit precision.
+ It retains all the flexibility of PyTorch,
+ in case you need it,
+ but adds some useful abstractions
+ and builds in some best practices.
+
+ ## What this notebook covers:
+
+ 1. Differences between PyTorch and PyTorch Lightning, including how to set up `LightningModules` and `LightningDataModules`
+ 2. How to get basic metric logging with the [`WandbLogger`](https://pytorch-lightning.readthedocs.io/en/latest/common/loggers.html#weights-and-biases)
+ 3. How to log media with W&B and fully customize logging with Lightning `Callbacks`
+
+ ## The interactive dashboard in W&B will look like this:
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Follow along with a [video tutorial](http://wandb.me/lit-video)!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🚀 Installing and importing
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ `wandb` and `pytorch-lightning` are both easily installable via [`pip`](https://pip.pypa.io/en/stable/).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb lightning torchmetrics onnx !pip install -qqq wandb lightning torchmetrics onnx
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ PyTorch Lightning is built on top of PyTorch,
+ so we still need to import vanilla PyTorch.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # numpy for non-GPU array math
+ import numpy as np
+
+ # 🍦 Vanilla PyTorch
+ import torch
+ from torch.nn import functional as F
+ from torch import nn
+ from torch.utils.data import DataLoader, random_split
+
+ # 👀 Torchvision for CV
+ from torchvision.datasets import MNIST
+ from torchvision import transforms
+
+ # remove slow mirror from list of MNIST mirrors
+ MNIST.mirrors = [mirror for mirror in MNIST.mirrors
+ if not mirror.startswith("http://yann.lecun.com")]
+ return DataLoader, F, MNIST, nn, np, random_split, torch, transforms
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Much of Lightning is built on the [Modules](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)
+ API from PyTorch,
+ but adds extra features
+ (like data loading and logging)
+ that are common to lots of PyTorch projects.
+
+ Let's bring those in,
+ plus W&B and the integration.
+
+ Lastly, we log in to the [Weights & Biases web service](https://wandb.ai).
+ If you've never used W&B,
+ you'll need to sign up first.
+ Accounts are free forever for academic and public projects.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # ⚡ PyTorch Lightning
+ import lightning.pytorch as pl
+ import torchmetrics
+ pl.seed_everything(hash("setting random seeds") % 2**32 - 1)
+
+ # 🏋️♀️ Weights & Biases
+ import wandb
+
+ # ⚡ 🤝 🏋️♀️
+ from lightning.pytorch.loggers import WandbLogger
+
+ return WandbLogger, pl, torchmetrics, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_: If you're executing your training in a terminal, rather than a notebook, you don't need to include `wandb.login()` in your script.
+ Instead, call `wandb login` in the terminal and we'll keep you logged in for future runs.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏗️ Building a Model with Lightning
+
+ In PyTorch Lightning, models are built with `LightningModule` ([docs here](https://pytorch-lightning.readthedocs.io/en/latest/lightning_module.html)), which has all the functionality of a vanilla `torch.nn.Module` (🍦) but with a few delicious cherries of added functionality on top (🍨).
+ These cherries are there to cut down on boilerplate and
+ help separate out the ML engineering code
+ from the actual machine learning.
+
+ For example, the mechanics of iterating over batches
+ as part of an epoch are extracted away,
+ so long as you define what happens on the `training_step`.
+
+ To make a working model out of a `LightningModule`,
+ we need to define a new `class` and add a few methods on top.
+
+ We'll demonstrate this process with `LitMLP`,
+ which applies a two-layer perceptron
+ (aka two fully-connected layers and
+ a fully-connected softmax readout layer)
+ to input `Tensors`.
+
+ > _Note_: It is common in the Lightning community to shorten "Lightning" to "[Lit](https://www.urbandictionary.com/define.php?term=it%27s%20lit)".
+ This sometimes it sound like
+ [your code was written by Travis Scott](https://www.youtube.com/watch?v=y3FCXV8oEZU).
+ We consider this a good thing.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🍦 `__init__` and `forward`
+
+ First, we need to add two methods that
+ are part of any vanilla PyTorch model.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Those methods are:
+ * `__init__` to do any setup, just like any Python class
+ * `forward` for inference, just like a PyTorch Module
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The `forward` pass method is standard,
+ and it'll be different for every project,
+ so we won't comment on it.
+
+ The `__init__` method,
+ which `init`ializes new instances of the class,
+ is a good place to log hyperparameter information to `wandb`.
+
+ This is done with the `save_hyperparameters` method,
+ which captures all of the arguments to the initializer
+ and adds them to a dictionary at `self.hparams` --
+ that all comes for free as part of the `LightningModule`.
+
+ > _Note_: `hparams` is logged to `wandb` as the `config`,
+ so you'll never lose track of the arguments you used to run a model again!
+ """)
+ return
+
+
+@app.cell
+def _(F, nn, np, pl, torchmetrics):
+ class LitMLP(pl.LightningModule):
+
+ def __init__(self, in_dims, n_classes=10,
+ n_layer_1=128, n_layer_2=256, lr=1e-4):
+ super().__init__()
+
+ # we flatten the input Tensors and pass them through an MLP
+ self.layer_1 = nn.Linear(np.prod(in_dims), n_layer_1)
+ self.layer_2 = nn.Linear(n_layer_1, n_layer_2)
+ self.layer_3 = nn.Linear(n_layer_2, n_classes)
+
+ # log hyperparameters
+ self.save_hyperparameters()
+
+ # compute the accuracy -- no need to roll your own!
+ self.train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=n_classes)
+ self.valid_acc = torchmetrics.Accuracy(task="multiclass", num_classes=n_classes)
+ self.test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=n_classes)
+
+ def forward(self, x):
+ """
+ Defines a forward pass using the Stem-Learner-Task
+ design pattern from Deep Learning Design Patterns:
+ https://www.manning.com/books/deep-learning-design-patterns
+ """
+ batch_size, *dims = x.size()
+
+ # stem: flatten
+ x = x.view(batch_size, -1)
+
+ # learner: two fully-connected layers
+ x = F.relu(self.layer_1(x))
+ x = F.relu(self.layer_2(x))
+
+ # task: compute class logits
+ x = self.layer_3(x)
+ x = F.log_softmax(x, dim=1)
+
+ return x
+
+ # convenient method to get the loss on a batch
+ def loss(self, xs, ys):
+ logits = self(xs) # this calls self.forward
+ loss = F.nll_loss(logits, ys)
+ return logits, loss
+
+ return (LitMLP,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_: for pedagogical purposes, we're splitting out
+ each stage of building the `LitMLP` into a different cell.
+ In a more typical workflow,
+ this would all happen in the `class` definition.
+
+ > _Note_: if you're familiar with PyTorch,
+ you might be surprised to see we aren't taking care with `.device`s:
+ no `to_cuda` etc. PyTorch Lightning handles all that for you! 😎
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🍨 `training_step` and `configure_optimizers`
+ Now, we add some special methods so that our `LitMLP` can be trained
+ using PyTorch Lightning's training API.
+
+ > _Note_: if you've used Keras, this might be familiar.
+ It's very similar to the `.fit` API in that library.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Those methods are
+
+ * `training_step`, which takes a batch and computes the loss; backprop goes through it
+ * `configure_optimizers`, which returns the `torch.optim.Optimizer` to apply after the `training_step`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_: `training_step` is part of a rich system of callbacks in PyTorch Lightning.
+ These callbacks are methods that get called
+ at specific points during training
+ (e.g. when a validation epoch ends),
+ and they are a major part of what makes
+ PyTorch Lightning both useful and extensible.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's where we add some more serious logging code.
+ `self.log` takes a name and value for a metric.
+ Under the hood, this will get passed to `wandb.log` if you're using W&B.
+
+ The logging behavior of PyTorch Lightning is both intelligent and configurable.
+ For example, by passing the `on_epoch`
+ keyword argument here,
+ we'll get `_epoch`-wise averages
+ of the metrics logged on each `_step`,
+ and those metrics will be named differently
+ in the W&B interface.
+ When training in a distributed setting,
+ these averages will be automatically computed across nodes.
+
+ Read more about the `log` method [in the docs](https://pytorch-lightning.readthedocs.io/en/latest/lightning_module.html#log).
+ """)
+ return
+
+
+@app.cell
+def _(LitMLP, torch):
+ def training_step(self, batch, batch_idx):
+ xs, ys = batch
+ logits, loss = self.loss(xs, ys)
+ preds = torch.argmax(logits, 1)
+
+ # logging metrics we calculated by hand
+ self.log('train/loss', loss, on_epoch=True)
+ # logging a pl.Metric
+ self.train_acc(preds, ys)
+ self.log('train/acc', self.train_acc, on_epoch=True)
+
+ return loss
+
+ def configure_optimizers(self):
+ return torch.optim.Adam(self.parameters(), lr=self.hparams["lr"])
+
+ LitMLP.training_step = training_step
+ LitMLP.configure_optimizers = configure_optimizers
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ➕ Optional methods for even better logging
+
+ The code above will log our model's performance,
+ system metrics, and more to W&B.
+
+ If we want to take our logging to the next level,
+ we need to make use of PyTorch Lightning's callback system.
+
+ > _Note_: thanks to the clean design of PyTorch Lightning,
+ the training code below will run with or without any
+ of this extra logging code. Nice!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The other callbacks we'll make use of fall into two categories:
+ * methods that trigger on each batch for a dataset: `validation_step` and `test_step`
+ * methods that trigger at the end of an epoch,
+ or a full pass over a given dataset: `{training, validation, test}_epoch_end`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 💾 `test`ing and saving the model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We use the test set to evaluate the performance of the final model,
+ so the `test` callbacks will be called at the end of the training pipeline.
+
+ For performance on the `test` and `validation` sets,
+ we're typically less concerned about how
+ we do on intermediate steps and more
+ with how we did overall.
+ That's why below, we pass in
+ `on_step=False` and `on_epoch=True`
+ so that we log only `epoch`-wise metrics.
+
+ > _Note_: That's actually the default behavior for `.log` when it's called inside of a `validation` or a `test` loop -- but not when it's called inside a `training` loop! Check out the table of default behaviors for `.log` [in the docs](https://pytorch-lightning.readthedocs.io/en/latest/lightning_module.html#log).
+ """)
+ return
+
+
+@app.cell
+def _(LitMLP, torch):
+ def test_step(self, batch, batch_idx):
+ xs, ys = batch
+ logits, loss = self.loss(xs, ys)
+ preds = torch.argmax(logits, 1)
+
+ self.test_acc(preds, ys)
+ self.log("test/loss_epoch", loss, on_step=False, on_epoch=True)
+ self.log("test/acc_epoch", self.test_acc, on_step=False, on_epoch=True)
+
+ LitMLP.test_step = test_step
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll also take the opportunity to save the model in the
+ [portable `ONNX` format](https://onnx.ai/).
+
+ Later,
+ we'll see that this allows us to use the
+ [Netron model viewer](https://github.com/lutzroeder/netron) in W&B.
+ """)
+ return
+
+
+@app.cell
+def _(LitMLP, torch, wandb):
+ def on_test_epoch_end(self): # args are defined as part of pl API
+ dummy_input = torch.zeros(self.hparams["in_dims"], device=self.device)
+ model_filename = "model_final.onnx"
+ self.to_onnx(model_filename, dummy_input, export_params=True)
+ artifact = wandb.Artifact(name="model.ckpt", type="model")
+ artifact.add_file(model_filename)
+ wandb.log_artifact(artifact)
+
+ LitMLP.on_test_epoch_end = on_test_epoch_end
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 📊 Logging `Histograms`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For the `validation_data`,
+ let's track not only the `acc`uracy and `loss`,
+ but also the `logits`:
+ the un-normalized class probabilities.
+ That way, we can track if our network
+ is becoming more or less confident over time.
+
+ There's a problem though:
+ `.log` wants to average,
+ but we'd rather look at a distribution.
+
+ So instead, on every `validation_step`,
+ we'll `return` the `logits`,
+ rather than `log`ging them.
+
+ Then, when we reach the `end`
+ of the `validation_epoch`,
+ the `logits` are available as the
+ `validation_step_outputs` -- a list.
+
+ So to log we'll take those `logits`,
+ concatenate them together,
+ and turn them into a histogram with [`wandb.Histogram`](https://docs.wandb.com/library/log#histograms).
+
+ Because we're no longer using Lightning's `.log` interface and are instead using `wandb`,
+ we need to drop down a level and use
+ `self.experiment.logger.log`.
+ """)
+ return
+
+
+@app.cell
+def _(LitMLP, torch, wandb):
+ def on_validation_epoch_start(self):
+ self.validation_step_outputs = []
+
+ def validation_step(self, batch, batch_idx):
+ xs, ys = batch
+ logits, loss = self.loss(xs, ys)
+ preds = torch.argmax(logits, 1)
+ self.valid_acc(preds, ys)
+
+ self.log("valid/loss_epoch", loss) # default on val/test is on_epoch only
+ self.log('valid/acc_epoch', self.valid_acc)
+
+ self.validation_step_outputs.append(logits)
+
+ return logits
+
+ def on_validation_epoch_end(self):
+
+ validation_step_outputs = self.validation_step_outputs
+
+ dummy_input = torch.zeros(self.hparams["in_dims"], device=self.device)
+ model_filename = f"model_{str(self.global_step).zfill(5)}.onnx"
+ torch.onnx.export(self, dummy_input, model_filename, opset_version=11)
+ artifact = wandb.Artifact(name="model.ckpt", type="model")
+ artifact.add_file(model_filename)
+ self.logger.experiment.log_artifact(artifact)
+
+ flattened_logits = torch.flatten(torch.cat(validation_step_outputs))
+ self.logger.experiment.log(
+ {"valid/logits": wandb.Histogram(flattened_logits.to("cpu")),
+ "global_step": self.global_step})
+
+ LitMLP.on_validation_epoch_start = on_validation_epoch_start
+ LitMLP.validation_step = validation_step
+ LitMLP.on_validation_epoch_end = on_validation_epoch_end
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Note that we're once again saving
+ the model in ONNX format.
+ That way, we can roll back our model to any given epoch --
+ useful in case the evaluation on the test set reveals we've overfit.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 📲 `Callback`s for extra-fancy logging
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ What we've done so far
+ will tell us how well our model
+ is using our system resources,
+ how well our model is training and generalizing,
+ and how confident it is.
+
+ But DNNs often fail in pernicious and silent ways.
+ Often, the only way to notice these failures
+ is to look at how the model is doing
+ on specific examples.
+
+ So let's additionally log some detailed information on some specific examples:
+ the inputs, outputs,
+ and `pred`ictions.
+
+ We'll do this by writing our own `Callback` --
+ one that, after every `validation_epoch` ends,
+ logs input images and output predictions
+ using W&B's `Image` logger.
+
+ > _Note_:
+ For more on the W&B media toolkit, read the [docs](https://docs.wandb.com/library/log#media)
+ or check out
+ [this Colab](http://wandb.me/media-colab)
+ to see everything it's capable of.
+ """)
+ return
+
+
+@app.cell
+def _(pl, torch, wandb):
+ class ImagePredictionLogger(pl.Callback):
+ def __init__(self, val_samples, num_samples=32):
+ super().__init__()
+ self.val_imgs, self.val_labels = val_samples
+ self.val_imgs = self.val_imgs[:num_samples]
+ self.val_labels = self.val_labels[:num_samples]
+
+ def on_validation_epoch_end(self, trainer, pl_module):
+ val_imgs = self.val_imgs.to(device=pl_module.device)
+
+ logits = pl_module(val_imgs)
+ preds = torch.argmax(logits, 1)
+
+ trainer.logger.experiment.log({
+ "examples": [wandb.Image(x, caption=f"Pred:{pred}, Label:{y}")
+ for x, pred, y in zip(val_imgs, preds, self.val_labels)],
+ "global_step": trainer.global_step
+ })
+
+ return (ImagePredictionLogger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🛒 Loading data
+
+ Data pipelines can be created with:
+ * 🍦 Vanilla Pytorch `DataLoaders`
+ * ⚡ Pytorch Lightning `DataModules`
+
+ `DataModules` are more structured definition, which allows for additional optimizations such as automated distribution of workload between CPU & GPU.
+ Using `DataModules` is recommended whenever possible!
+
+ A `DataModule` is also defined by an interface:
+ * `prepare_data` (optional) which is called only once and on 1 GPU -- typically something like the data download step we have below
+ * `setup`, which is called on each GPU separately and accepts `stage` to define if we are at `fit` or `test` step
+ * `train_dataloader`, `val_dataloader` and `test_dataloader` to load each dataset
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, MNIST, pl, random_split, transforms):
+ class MNISTDataModule(pl.LightningDataModule):
+
+ def __init__(self, data_dir='./', batch_size=128):
+ super().__init__()
+ self.data_dir = data_dir
+ self.batch_size = batch_size
+ self.transform = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize((0.1307,), (0.3081,))])
+
+ def prepare_data(self):
+ # download data, train then test
+ MNIST(self.data_dir, train=True, download=True)
+ MNIST(self.data_dir, train=False, download=True)
+
+ def setup(self, stage=None):
+
+ # we set up only relevant datasets when stage is specified
+ if stage == 'fit' or stage is None:
+ mnist = MNIST(self.data_dir, train=True, transform=self.transform)
+ self.mnist_train, self.mnist_val = random_split(mnist, [55000, 5000])
+ if stage == 'test' or stage is None:
+ self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)
+
+ # we define a separate DataLoader for each of train/val/test
+ def train_dataloader(self):
+ mnist_train = DataLoader(self.mnist_train, batch_size=self.batch_size)
+ return mnist_train
+
+ def val_dataloader(self):
+ mnist_val = DataLoader(self.mnist_val, batch_size=10 * self.batch_size)
+ return mnist_val
+
+ def test_dataloader(self):
+ mnist_test = DataLoader(self.mnist_test, batch_size=10 * self.batch_size)
+ return mnist_test
+
+ # setup data
+ mnist = MNISTDataModule()
+ mnist.prepare_data()
+ mnist.setup()
+
+ # grab samples to log predictions on
+ samples = next(iter(mnist.val_dataloader()))
+ return mnist, samples
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👟 Making a `Trainer`
+
+ The `DataLoader` and the `LightningModule`
+ are brought together by a `Trainer`,
+ which orchestrates data loading,
+ gradient calculation,
+ optimizer logic,
+ and logging.
+
+ Luckily, we don't need to sub-class the `Trainer`,
+ we just need to configure it with keyword arguments.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ And that is where we'll use the `pytorch_lightning.loggers.WandbLogger` to connect our logging to W&B.
+ """)
+ return
+
+
+@app.cell
+def _(WandbLogger):
+ wandb_logger = WandbLogger(project="lit-wandb")
+ return (wandb_logger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_: Check out [the documentation](https://docs.wandb.com/library/integrations/lightning) for customization options. I like `group`s and `tag`s!.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can then set up our `Trainer` and customize several options, such as gradient accumulation, half precision training and distributed computing.
+
+ We'll stick to the basics for this example,
+ but half-precision training and easy scaling to distributed settings are two of the major reasons why folks like PyTorch Lightning!
+ """)
+ return
+
+
+@app.cell
+def _(ImagePredictionLogger, pl, samples, wandb_logger):
+ trainer = pl.Trainer(
+ logger=wandb_logger, # W&B integration
+ log_every_n_steps=50, # set the logging frequency
+ max_epochs=5, # number of epochs
+ deterministic=True, # keep it deterministic
+ callbacks=[ImagePredictionLogger(samples)] # see Callbacks section
+ )
+ return (trainer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🏃♀️ Running our Model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, let's make it all happen:
+ """)
+ return
+
+
+@app.cell
+def _(LitMLP, mnist, trainer, wandb):
+ # setup model
+ model = LitMLP(in_dims=(1, 28, 28))
+
+ # fit the model
+ trainer.fit(model, mnist)
+
+ # evaluate the model on a test set
+ trainer.test(datamodule=mnist,
+ ckpt_path=None) # uses last-saved model
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_: In notebooks, we need to call `wandb.finish()` to indicate when we've finished our run. This isn't necessary in scripts.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Viewing the results on wandb.ai
+
+ Among the outputs from W&B,
+ you will have noticed a few URLs.
+ One of these is the
+ [run page](https://docs.wandb.ai/ref/app/pages/run-page),
+ which has a dashboard with all of the information logged in this run, complete with smart default charts
+ and more.
+ The run page is printed both at the start and end of training, and ends with `lit-wandb/runs/{run_id}`.
+
+ >_Note_: When visiting your run page, it is recommended to use `global_step` as x-axis to correctly superimpose metrics logged in different stages.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-track-pytorch-lightning-with-fabric-and-wandb/pytorch_lightning_track_pytorch_lightning_with_fabric_and_wandb.py b/marimo/convert/pytorch-lightning-track-pytorch-lightning-with-fabric-and-wandb/pytorch_lightning_track_pytorch_lightning_with_fabric_and_wandb.py
new file mode 100644
index 00000000..3c41e2f9
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-track-pytorch-lightning-with-fabric-and-wandb/pytorch_lightning_track_pytorch_lightning_with_fabric_and_wandb.py
@@ -0,0 +1,375 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # ⚡ Track PyTorch Lightning with Fabric and Wandb
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ At Weights & Biases, we love anything
+ that makes training deep learning models easier.
+ That's why we worked with the folks at PyTorch Lightning to
+ [integrate our experiment tracking tool](https://docs.wandb.com/library/integrations/lightning)
+ directly into the Fabric library of PyTorch Lightning
+
+ [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/) is a lightweight wrapper for organizing your PyTorch code and easily adding advanced features such as distributed training and 16-bit precision.
+ It retains all the flexibility of PyTorch,
+ in case you need it,
+ but adds some useful abstractions
+ and builds in some best practices.
+
+ [Pytorch Fabric](https://lightning.ai/docs/fabric/stable/) allows you to scale PyTorch models on
+ distributed machines while
+ maintaining full control of your
+ training loop.
+
+ ## What this notebook covers:
+
+ 1. How to get basic metric logging with the `WandbLogger`
+ 2. How to log media with W&B
+
+ ## The interactive dashboard in W&B will look like this:
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install wandb
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.environ["WANDB_API_KEY"]=""
+ return (os,)
+
+
+@app.cell
+def _():
+ import wandb
+ wandb.login()
+ return (wandb,)
+
+
+@app.cell
+def _():
+ import lightning as L
+ import torch; import torchvision as tv
+ from wandb.integration.lightning.fabric import WandbLogger
+
+ return L, WandbLogger, torch, tv
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💡 Tracking Experiments with WandbLogger
+
+ PyTorch Lightning has a `WandbLogger` to easily log your experiments with Wights & Biases. Just pass it to your `Trainer` to log to W&B. See the WandbLogger docs for all parameters. Note, to log the metrics to a specific W&B Team, pass your Team name to the `entity` argument in `WandbLogger`
+
+ #### `lightning.fabric.loggers.WandbLogger()`
+
+ | Functionality | Argument/Function | PS |
+ | ------ | ------ | ------ |
+ | Logging models | `WandbLogger(... ,log_model='all')` or `WandbLogger(... ,log_model=True`) | Log all models if `log_model="all"` and at end of training if `log_model=True`
+ | Set custom run names | `WandbLogger(... ,name='my_run_name'`) | |
+ | Organize runs by project | `WandbLogger(... ,project='my_project')` | |
+ | Log histograms of gradients and parameters | `WandbLogger.watch(model)` | `WandbLogger.watch(model, log='all')` to log parameter histograms |
+ | Log hyperparameters | Call `self.save_hyperparameters()` within `LightningModule.__init__()` |
+ | Log custom objects (images, audio, video, molecules…) | Use `WandbLogger.log_text`, `WandbLogger.log_image` and `WandbLogger.log_table`, etc. |
+ """)
+ return
+
+
+@app.cell
+def _(WandbLogger):
+ logger = WandbLogger(project="Cifar10_ptl_fabric")
+ return (logger,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log custom hyperparameters and configurations
+ """)
+ return
+
+
+@app.cell
+def _(logger):
+ lr = 0.001
+ batch_size = 16
+ num_epochs = 5
+ classes = ('plane', 'car', 'bird', 'cat',
+ 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
+ log_images_after_n_batches = 200
+
+ logger.log_hyperparams({
+ "lr": lr,
+ "batch_size": batch_size,
+ "num_epochs": num_epochs,
+ "classes": classes,
+ "log_images_after_n_batches": log_images_after_n_batches
+ })
+ return batch_size, classes, log_images_after_n_batches, lr, num_epochs
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Save Data to Weights and Biases Artifacts
+
+ This allows us to audit and create direct data lineages to our experiments
+ """)
+ return
+
+
+@app.cell
+def _():
+ root_folder = "data"
+ return (root_folder,)
+
+
+@app.cell
+def _(root_folder, tv):
+ train_dataset = tv.datasets.CIFAR10(root_folder, download=True,
+ train=True,
+ transform=tv.transforms.ToTensor())
+ test_dataset = tv.datasets.CIFAR10(root_folder, download=True,
+ train=False,
+ transform=tv.transforms.ToTensor())
+ return test_dataset, train_dataset
+
+
+@app.cell
+def _(train_dataset):
+ data_folder = train_dataset.base_folder # same as test_dataset.base_folder
+ return (data_folder,)
+
+
+@app.cell
+def _(data_folder, logger, os, root_folder, wandb):
+ data_art = wandb.Artifact(name="cifar10", type="dataset")
+ data_art.add_dir(os.path.join(root_folder, data_folder))
+ logger.experiment.log_artifact(data_art)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Configure our Model and Training
+ """)
+ return
+
+
+@app.cell
+def _(lr, torch, tv):
+ model = tv.models.resnet18()
+ optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+ return model, optimizer
+
+
+@app.cell
+def _(wandb):
+ class TableLoggingCallback:
+ def __init__(self, wandb_logger):
+ self.wandb_logger = wandb_logger
+ self.table = wandb.Table(columns=["image", "prediction", "ground_truth"])
+
+ def on_test_batch_end(self, images, predictions, ground_truths):
+ for image, prediction, ground_truth in zip(images, predictions, ground_truths):
+ self.table.add_data(wandb.Image(image), prediction, ground_truth)
+
+ def on_model_epoch_end(self):
+ prediction_table = self.table
+ print(self.table.data[0])
+ self.wandb_logger.experiment.log({"prediction_table": prediction_table}) # You can directly access the run object via `experiment`
+
+ # We could also use
+ # (1) wandb_logger.log_metrics()
+ # (2) wandb_logger.log_table()
+
+ self.table = wandb.Table(columns=["image", "prediction", "ground_truth"])
+
+ return (TableLoggingCallback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Load our model, datasources, and loggers into PyTorch Fabric
+ """)
+ return
+
+
+@app.cell
+def _(TableLoggingCallback, logger):
+ tlc = TableLoggingCallback(logger)
+ return (tlc,)
+
+
+@app.cell
+def _(L, logger, tlc):
+ fabric = L.Fabric(loggers=[logger], callbacks=[tlc])
+ fabric.launch()
+ return (fabric,)
+
+
+@app.cell
+def _(
+ batch_size,
+ fabric,
+ model,
+ optimizer,
+ test_dataset,
+ torch,
+ train_dataset,
+):
+ model_1, optimizer_1 = fabric.setup(model, optimizer)
+ train_dataloader = fabric.setup_dataloaders(torch.utils.data.DataLoader(train_dataset, batch_size=batch_size))
+ test_dataloader = fabric.setup_dataloaders(torch.utils.data.DataLoader(test_dataset, batch_size=batch_size))
+ return model_1, optimizer_1, test_dataloader, train_dataloader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Run training and log test predictions
+
+ For every epoch, run a training step and a test step. For each n test batches, we log the batch of test images caption by the prediction and label, and we create a wandb.Table() in which to store test predictions using our custom callback
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ No additional dependencies outside the Torch modeling you're used to!
+ """)
+ return
+
+
+@app.cell
+def _(logger, model_1):
+ logger.watch(model_1)
+ return
+
+
+@app.cell
+def _(
+ classes,
+ fabric,
+ log_images_after_n_batches,
+ logger,
+ model_1,
+ num_epochs,
+ optimizer_1,
+ test_dataloader,
+ torch,
+ train_dataloader,
+):
+ model_1.train()
+ for epoch in range(num_epochs):
+ fabric.print(f'Epoch: {epoch}')
+ cum_loss = 0 # Training Loop
+ for batch in train_dataloader:
+ inputs, labels = batch
+ optimizer_1.zero_grad()
+ outputs = model_1(inputs) # Batch by batch of data from training dataset
+ loss = torch.nn.functional.cross_entropy(outputs, labels)
+ cum_loss = cum_loss + loss.item()
+ fabric.backward(loss)
+ optimizer_1.step()
+ fabric.log_dict({'loss': loss.item()})
+ fabric.log_dict({'avg_loss': cum_loss / len(train_dataloader)})
+ correct = 0
+ total = 0
+ class_correct = list((0.0 for i in range(10)))
+ class_total = list((0.0 for i in range(10))) # Stream per batch training metrics
+ test_batch_ctr = 0
+ with torch.no_grad(): # Stream per epoch training metrics
+ for batch_ctr, batch in enumerate(test_dataloader):
+ images, labels = batch # Validation Loop
+ outputs = model_1(images)
+ _, predicted = torch.max(outputs, 1)
+ total = total + labels.size(0)
+ correct = correct + (predicted == labels).sum().item()
+ c = (predicted == labels).squeeze()
+ for i in range(batch[0].size(0)):
+ label = labels[i]
+ class_correct[label] = class_correct[label] + c[i].item()
+ class_total[label] = class_total[label] + 1
+ if batch_ctr % log_images_after_n_batches == 0: # Batch by batch of data from testing dataset
+ predictions = [classes[prediction] for prediction in predicted]
+ label_names = [classes[truth] for truth in labels]
+ loggable_images = [image for image in images]
+ captions = [f'pred: {pred}\nlabel: {truth}' for pred, truth in zip(predictions, label_names)]
+ logger.log_image(key='test_image_batch', images=loggable_images, step=None, caption=captions)
+ fabric.call('on_test_batch_end', images=loggable_images, predictions=predictions, ground_truths=label_names) # Overall Test Accuracy
+ test_acc = 100 * correct / total
+ class_acc = {f'{classes[i]}_acc': 100 * class_correct[i] / class_total[i] for i in range(10) if class_total[i] > 0}
+ loggable_dict = {'test_acc': test_acc}
+ loggable_dict.update(class_acc) # Per Class Accuracy
+ fabric.log_dict(loggable_dict)
+ fabric.call('on_model_epoch_end') # Test Images labeled with Class prediction for qualitative analysis # Automatically construct and log wandb.Images # Can also just directly log the below list via fabric.log_dict # [wandb.Image(image, caption=classes[predicted]) for image, predicted, label in zip(images, predicted, labels)]) # Populate per batch data within our table # Calculate cumulative test metrics # Stream per epoch validation metrics # Save epoch test data table to dashboard
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Finish our experiment!
+ """)
+ return
+
+
+@app.cell
+def _(logger):
+ logger.experiment.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-transfer-learning-using-pytorch-lightning/pytorch_lightning_transfer_learning_using_pytorch_lightning.py b/marimo/convert/pytorch-lightning-transfer-learning-using-pytorch-lightning/pytorch_lightning_transfer_learning_using_pytorch_lightning.py
new file mode 100644
index 00000000..3fb52d7c
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-transfer-learning-using-pytorch-lightning/pytorch_lightning_transfer_learning_using_pytorch_lightning.py
@@ -0,0 +1,398 @@
+# /// script
+# dependencies = ["lightning", "torchvision", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Transfer Learning Using PyTorch Lightning ⚡️
+
+ In this colab, we will extend the pipeline [here](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/pytorch-lightning/Image_Classification_using_PyTorch_Lightning.ipynb) to perform transfer learning with PyTorch Lightning.
+
+ Transfer Learning is a technique where the knowledge learned while training a model for "task" A and can be used for "task" B. Here A and B can be the same deep learning tasks but on a different dataset.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setting up PyTorch Lightning and W&B
+
+ For this tutorial, we need PyTorch Lightning and Weights and Biases.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb lightning torchvision !pip install wandb lightning torchvision -qqq
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You're gonna need these imports.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+
+ import lightning.pytorch as pl
+ # your favorite machine learning tracking tool
+ from lightning.pytorch.loggers import WandbLogger
+
+ import torch
+ from torch import nn
+ from torch.nn import functional as F
+ from torch.utils.data import random_split, DataLoader
+
+ from torchmetrics import Accuracy
+
+ from torchvision import transforms
+ from torchvision.datasets import StanfordCars
+ from torchvision.datasets.utils import download_url
+ import torchvision.models as models
+
+
+ import wandb
+
+ return (
+ Accuracy,
+ DataLoader,
+ StanfordCars,
+ WandbLogger,
+ models,
+ nn,
+ pl,
+ random_split,
+ torch,
+ transforms,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now you'll need to login to you wandb account.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## The Dataset 💿
+
+ We will be using the StanfordCars dataset to train our image classifier. It contains 16,185 images of 196 classes of cars. The data is split into 8,144 training images and 8,041 testing images, where each class has been split roughly in a 50-50 split. Classes are typically at the level of Make, Model, Year, e.g. 2012 Tesla Model S or 2012 BMW M3 coupe.
+ """)
+ return
+
+
+@app.cell
+def _(DataLoader, StanfordCars, pl, random_split, transforms):
+ class StanfordCarsDataModule(pl.LightningDataModule):
+ def __init__(self, batch_size, data_dir: str = './'):
+ super().__init__()
+ self.data_dir = data_dir
+ self.batch_size = batch_size
+
+ # Augmentation policy for training set
+ self.augmentation = transforms.Compose([
+ transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
+ transforms.RandomRotation(degrees=15),
+ transforms.RandomHorizontalFlip(),
+ transforms.CenterCrop(size=224),
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
+ ])
+ # Preprocessing steps applied to validation and test set.
+ self.transform = transforms.Compose([
+ transforms.Resize(size=256),
+ transforms.CenterCrop(size=224),
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
+ ])
+
+ self.num_classes = 196
+
+ def prepare_data(self):
+ pass
+
+ def setup(self, stage=None):
+ # build dataset
+ dataset = StanfordCars(root=self.data_dir, download=True, split="train")
+ # split dataset
+ self.train, self.val = random_split(dataset, [6500, 1644])
+
+ self.test = StanfordCars(root=self.data_dir, download=True, split="test")
+
+ self.test = random_split(self.test, [len(self.test)])[0]
+
+ self.train.dataset.transform = self.augmentation
+ self.val.dataset.transform = self.transform
+ self.test.dataset.transform = self.transform
+
+ def train_dataloader(self):
+ return DataLoader(self.train, batch_size=self.batch_size, shuffle=True, num_workers=2)
+
+ def val_dataloader(self):
+ return DataLoader(self.val, batch_size=self.batch_size, num_workers=2)
+
+ def test_dataloader(self):
+ return DataLoader(self.test, batch_size=self.batch_size, num_workers=2)
+
+ return (StanfordCarsDataModule,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## LightingModule - Define the System
+
+ Let us look at the model definition to see how transfer learning can be used with PyTorch Lightning.
+ In the `LitModel` class, we can use the pre-trained model provided by Torchvision as a feature extractor for our classification model. Here we are using ResNet-18. A list of pre-trained models provided by PyTorch Lightning can be found here.
+ - When `pretrained=True`, we use the pre-trained weights; otherwise, the weights are initialized randomly.
+ - If `.eval()` is used, then the layers are frozen.
+ - A single `Linear` layer is used as the output layer. We can have multiple layers stacked over the `feature_extractor`.
+
+ Setting the `transfer` argument to `True` will enable transfer learning.
+ """)
+ return
+
+
+@app.cell
+def _(Accuracy, models, nn, pl, torch):
+ class LitModel(pl.LightningModule):
+ def __init__(self, input_shape, num_classes, learning_rate=2e-4, transfer=False):
+ super().__init__()
+
+ # log hyperparameters
+ self.save_hyperparameters()
+ self.learning_rate = learning_rate
+ self.dim = input_shape
+ self.num_classes = num_classes
+
+ # transfer learning if pretrained=True
+ self.feature_extractor = models.resnet18(pretrained=transfer)
+
+ if transfer:
+ # layers are frozen by using eval()
+ self.feature_extractor.eval()
+ # freeze params
+ for param in self.feature_extractor.parameters():
+ param.requires_grad = False
+
+ n_sizes = self._get_conv_output(input_shape)
+
+ self.classifier = nn.Linear(n_sizes, num_classes)
+
+ self.criterion = nn.CrossEntropyLoss()
+ self.accuracy = Accuracy()
+
+ # returns the size of the output tensor going into the Linear layer from the conv block.
+ def _get_conv_output(self, shape):
+ batch_size = 1
+ tmp_input = torch.autograd.Variable(torch.rand(batch_size, *shape))
+
+ output_feat = self._forward_features(tmp_input)
+ n_size = output_feat.data.view(batch_size, -1).size(1)
+ return n_size
+
+ # returns the feature tensor from the conv block
+ def _forward_features(self, x):
+ x = self.feature_extractor(x)
+ return x
+
+ # will be used during inference
+ def forward(self, x):
+ x = self._forward_features(x)
+ x = x.view(x.size(0), -1)
+ x = self.classifier(x)
+
+ return x
+
+ def training_step(self, batch):
+ batch, gt = batch[0], batch[1]
+ out = self.forward(batch)
+ loss = self.criterion(out, gt)
+
+ acc = self.accuracy(out, gt)
+
+ self.log("train/loss", loss)
+ self.log("train/acc", acc)
+
+ return loss
+
+ def validation_step(self, batch, batch_idx):
+ batch, gt = batch[0], batch[1]
+ out = self.forward(batch)
+ loss = self.criterion(out, gt)
+
+ self.log("val/loss", loss)
+
+ acc = self.accuracy(out, gt)
+ self.log("val/acc", acc)
+
+ return loss
+
+ def test_step(self, batch, batch_idx):
+ batch, gt = batch[0], batch[1]
+ out = self.forward(batch)
+ loss = self.criterion(out, gt)
+
+ return {"loss": loss, "outputs": out, "gt": gt}
+
+ def test_epoch_end(self, outputs):
+ loss = torch.stack([x['loss'] for x in outputs]).mean()
+ output = torch.cat([x['outputs'] for x in outputs], dim=0)
+
+ gts = torch.cat([x['gt'] for x in outputs], dim=0)
+
+ self.log("test/loss", loss)
+ acc = self.accuracy(output, gts)
+ self.log("test/acc", acc)
+
+ self.test_gts = gts
+ self.test_output = output
+
+ def configure_optimizers(self):
+ return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+
+ return (LitModel,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train your Model 🏋️♂️
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To train the model, we instantiate the `StanfordCarsDataModule` and the `LitModel` along with the PyTorch Lightning Trainer. To the `Trainer`, we will pass the `WandbLogger` as the logger to use W&B to track the metrics during model training!
+ """)
+ return
+
+
+@app.cell
+def _(LitModel, StanfordCarsDataModule, WandbLogger, pl):
+ dm = StanfordCarsDataModule(batch_size=32)
+ model = LitModel((3, 300, 300), 196, transfer=True)
+ trainer = pl.Trainer(logger=WandbLogger(project="TransferLearning"), max_epochs=10, accelerator="gpu")
+ return dm, model, trainer
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We are good to go! Let's train our model!
+ """)
+ return
+
+
+@app.cell
+def _(dm, model, trainer):
+ trainer.fit(model, dm)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now that the model is trained, let's see how it performs on the test set
+ """)
+ return
+
+
+@app.cell
+def _(dm, model, trainer):
+ trainer.test(model, dm)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's close our W&B run, so we call `wandb.finish()`.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The workspace generated to compare training the model from scratch vs using transfer learning is available [here](https://wandb.ai/manan-goel/StanfordCars). The conclusions that can be drawn from this are explained in detail in [this report](https://wandb.ai/wandb/wandb-lightning/reports/Transfer-Learning-Using-PyTorch-Lightning--VmlldzoyMzMxMzk4/edit).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Conclusion
+
+ I will encourage you to play with the code and train an image classifier with a dataset of your choice from scratch and using transfer learning.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To learn more about transfer learning check out these resources:
+ - [Gotchas of transfer learning for image classification](https://docs.google.com/presentation/d/1s29WOQoQvBD5KoPUzE5TPcavjqno8ZgnZaSljHGGHVU/edit?usp=sharing) by Sayak Paul.
+ - [Transfer Learning with Keras and Deep Learning by PyImageSearch.](https://www.pyimagesearch.com/2019/05/20/transfer-learning-with-keras-and-deep-learning/)
+ - [Transfer Learning - Machine Learning's Next Frontier](https://ruder.io/transfer-learning/) by Sebastian Ruder.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-lightning-wandb-end-to-end-with-pytorch-lightning/pytorch_lightning_wandb_end_to_end_with_pytorch_lightning.py b/marimo/convert/pytorch-lightning-wandb-end-to-end-with-pytorch-lightning/pytorch_lightning_wandb_end_to_end_with_pytorch_lightning.py
new file mode 100644
index 00000000..48c381f5
--- /dev/null
+++ b/marimo/convert/pytorch-lightning-wandb-end-to-end-with-pytorch-lightning/pytorch_lightning_wandb_end_to_end_with_pytorch_lightning.py
@@ -0,0 +1,556 @@
+# /// script
+# dependencies = ["lightning", "torchvision", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # W&B Tutorial with Pytorch Lightning
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🛠️ Install `wandb` and `pytorch-lightning`
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: lightning wandb torchvision !pip install -q lightning wandb torchvision
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Login to W&B either through Python or CLI
+ If you are using the public W&B cloud, you don't need to specify the `WANDB_HOST`.
+
+ You can set environment variables `WANDB_API_KEY` and `WANDB_HOST` and pass them in as:
+ ```
+ import os
+ import wandb
+
+ wandb.login(host=os.getenv("WANDB_HOST"), key=os.getenv("WANDB_API_KEY"))
+ ```
+ You can also login via the CLI with:
+ ```
+ wandb login --host
+ ```
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ⚱ Logging the Raw Training Data as an Artifact
+ """)
+ return
+
+
+@app.cell
+def _():
+ #@title Enter your W&B project and entity
+
+ # FORM VARIABLES
+ PROJECT_NAME = "pytorch-lightning-e2e" #@param {type:"string"}
+ ENTITY = "wandb"#@param {type:"string"}
+
+ # set SIZE to "TINY", "SMALL", "MEDIUM", or "LARGE"
+ # to select one of these three datasets
+ # TINY dataset: 100 images, 30MB
+ # SMALL dataset: 1000 images, 312MB
+ # MEDIUM dataset: 5000 images, 1.5GB
+ # LARGE dataset: 12,000 images, 3.6GB
+
+ SIZE = "TINY"
+
+ if SIZE == "TINY":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_100.zip"
+ src_zip = "nature_100.zip"
+ DATA_SRC = "nature_100"
+ IMAGES_PER_LABEL = 10
+ BALANCED_SPLITS = {"train" : 8, "val" : 1, "test": 1}
+ elif SIZE == "SMALL":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_1K.zip"
+ src_zip = "nature_1K.zip"
+ DATA_SRC = "nature_1K"
+ IMAGES_PER_LABEL = 100
+ BALANCED_SPLITS = {"train" : 80, "val" : 10, "test": 10}
+ elif SIZE == "MEDIUM":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_12K.zip"
+ src_zip = "nature_12K.zip"
+ DATA_SRC = "inaturalist_12K/train" # (technically a subset of only 10K images)
+ IMAGES_PER_LABEL = 500
+ BALANCED_SPLITS = {"train" : 400, "val" : 50, "test": 50}
+ elif SIZE == "LARGE":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_12K.zip"
+ src_zip = "nature_12K.zip"
+ DATA_SRC = "inaturalist_12K/train" # (technically a subset of only 10K images)
+ IMAGES_PER_LABEL = 1000
+ BALANCED_SPLITS = {"train" : 800, "val" : 100, "test": 100}
+ return ENTITY, PROJECT_NAME
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !curl -SL $src_url > $src_zip
+ # !unzip $src_zip
+ return
+
+
+@app.cell
+def _(torch):
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ return
+
+
+@app.cell
+def _(ENTITY, PROJECT_NAME, wandb):
+ import pandas as pd
+ import os
+ with wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type='log_datasets') as run:
+ img_paths = []
+ for root, dirs, files in os.walk('nature_100', topdown=False):
+ for name in files:
+ img_path = os.path.join(root, name)
+ label = img_path.split('/')[1]
+ img_paths.append([img_path, label])
+ index_df = pd.DataFrame(columns=['image_path', 'label'], data=img_paths)
+ index_df.to_csv('index.csv', index=False)
+ train_art = wandb.Artifact(name='Nature_100', type='raw_images', description='nature image dataset with 10 classes, 10 images per class')
+ train_art.add_dir('nature_100')
+ train_art.add_file('index.csv')
+ wandb.log_artifact(train_art) # Also adding a csv indicating the labels of each image
+ return os, pd
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Using Artifacts in Pytorch Lightning `DataModule`'s and Pytorch `Dataset`'s
+ - Makes it easy to interopt your DataLoaders with new versions of datasets
+ - Just indicate the `name:alias` as an argument to your `Dataset` or `DataModule`
+ """)
+ return
+
+
+@app.cell
+def _(os, pd):
+ from torchvision import transforms
+ import lightning.pytorch as pl
+ import torch
+ from torch.utils.data import Dataset, DataLoader, random_split
+ from skimage import io, transform
+ from torchvision import utils, models
+ import math
+
+ class NatureDataset(Dataset):
+
+ def __init__(self, wandb_run, artifact_name_alias='Nature_100:latest', local_target_dir='Nature_100:latest', transform=None):
+ self.local_target_dir = local_target_dir
+ self.transform = transform
+ art = wandb_run.use_artifact(artifact_name_alias)
+ path_at = art.download(root=self.local_target_dir)
+ self.ref_df = pd.read_csv(os.path.join(self.local_target_dir, 'index.csv'))
+ self.class_names = self.ref_df.iloc[:, 1].unique().tolist()
+ self.idx_to_class = {k: v for k, v in enumerate(self.class_names)} # Pull down the artifact locally to load it into memory
+ self.class_to_idx = {v: k for k, v in enumerate(self.class_names)}
+
+ def __len__(self):
+ return len(self.ref_df)
+
+ def __getitem__(self, idx):
+ if torch.is_tensor(idx):
+ idx = idx.tolist()
+ img_path = self.ref_df.iloc[idx, 0]
+ image = io.imread(img_path)
+ label = self.ref_df.iloc[idx, 1]
+ label = torch.tensor(self.class_to_idx[label], dtype=torch.long)
+ if self.transform:
+ image = self.transform(image)
+ return (image, label)
+
+ class NatureDatasetModule(pl.LightningDataModule):
+
+ def __init__(self, wandb_run, artifact_name_alias: str='Nature_100:latest', local_target_dir: str='Nature_100:latest', batch_size: int=16, input_size: int=224, seed: int=42):
+ super().__init__()
+ self.wandb_run = wandb_run
+ self.artifact_name_alias = artifact_name_alias
+ self.local_target_dir = local_target_dir
+ self.batch_size = batch_size
+ self.input_size = input_size
+ self.seed = seed
+
+ def setup(self, stage=None):
+ self.nature_dataset = NatureDataset(wandb_run=self.wandb_run, artifact_name_alias=self.artifact_name_alias, local_target_dir=self.local_target_dir, transform=transforms.Compose([transforms.ToTensor(), transforms.CenterCrop(self.input_size), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]))
+ nature_length = len(self.nature_dataset)
+ train_size = math.floor(0.8 * nature_length)
+ val_size = math.floor(0.2 * nature_length)
+ self.nature_train, self.nature_val = random_split(self.nature_dataset, [train_size, val_size], generator=torch.Generator().manual_seed(self.seed))
+ return self
+
+ def train_dataloader(self):
+ return DataLoader(self.nature_train, batch_size=self.batch_size)
+
+ def val_dataloader(self):
+ return DataLoader(self.nature_val, batch_size=self.batch_size)
+
+ def predict_dataloader(self):
+ pass
+
+ def teardown(self, stage: str):
+ pass
+
+ return NatureDatasetModule, models, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ##How Logging in your Pytorch `LightningModule`works:
+ When you train the model using `Trainer`, ensure you have a `WandbLogger` instantiated and passed in as a `logger`.
+
+ ```
+ wandb_logger = WandbLogger(project="my_project", entity="machine-learning")
+ trainer = Trainer(logger=wandb_logger)
+ ```
+
+ You can always use `wandb.log` as normal throughout the module. When the `WandbLogger` is used, `self.log` will also log metrics to W&B.
+ - To access the current run from within the `LightningModule`, you can access `Trainer.logger.experiment`, which is a `wandb.Run` object
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Some helper functions
+ """)
+ return
+
+
+@app.cell
+def _(models, torch):
+ # Some helper functions
+ def set_parameter_requires_grad(model, feature_extracting):
+ if feature_extracting:
+ for param in _model.parameters():
+ param.requires_grad = False
+
+ def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
+ model_ft = None
+ input_size = 0 # Initialize these variables which will be set in this if statement. Each of these
+ if model_name == 'resnet': # variables is model specific.
+ ' Resnet18\n '
+ model_ft = models.resnet18(pretrained=use_pretrained)
+ set_parameter_requires_grad(model_ft, feature_extract)
+ num_ftrs = model_ft.fc.in_features
+ model_ft.fc = torch.nn.Linear(num_ftrs, num_classes)
+ input_size = 224
+ elif model_name == 'squeezenet':
+ ' Squeezenet\n '
+ model_ft = models.squeezenet1_0(pretrained=use_pretrained)
+ set_parameter_requires_grad(model_ft, feature_extract)
+ model_ft.classifier[1] = torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))
+ model_ft.num_classes = num_classes
+ input_size = 224
+ elif model_name == 'densenet':
+ ' Densenet\n '
+ model_ft = models.densenet121(pretrained=use_pretrained)
+ set_parameter_requires_grad(model_ft, feature_extract)
+ num_ftrs = model_ft.classifier.in_features
+ model_ft.classifier = torch.nn.Linear(num_ftrs, num_classes)
+ input_size = 224
+ else:
+ print('Invalid model name, exiting...')
+ exit()
+ return (model_ft, input_size)
+
+ return (initialize_model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Writing the `LightningModule`
+ """)
+ return
+
+
+@app.cell
+def _(initialize_model, torch, wandb):
+ from torch.nn import Linear, CrossEntropyLoss, functional as F
+ from torch.optim import Adam
+ from torchmetrics.functional import accuracy
+ from lightning.pytorch import LightningModule
+
+ class NatureLitModule(LightningModule):
+
+ def __init__(self, model_name, num_classes=10, feature_extract=True, lr=0.01):
+ """method used to define our model parameters"""
+ super().__init__()
+ self.model_name = model_name
+ self.num_classes = num_classes
+ self.feature_extract = feature_extract
+ self.model, self.input_size = initialize_model(model_name=self.model_name, num_classes=self.num_classes, feature_extract=True)
+ self.loss = CrossEntropyLoss()
+ self.lr = lr
+ self.save_hyperparameters()
+ wandb.watch(self.model)
+
+ def forward(self, x):
+ """method used for inference input -> output"""
+ x = self.model(x)
+ return x
+ # loss
+ def training_step(self, batch, batch_idx):
+ """needs to return a loss from a single batch"""
+ preds, y, loss, acc = self._get_preds_loss_accuracy(batch) # optimizer parameters
+ self.log('train/loss', loss)
+ self.log('train/accuracy', acc)
+ return loss # save hyper-parameters to self.hparams (auto-logged by W&B)
+
+ def validation_step(self, batch, batch_idx):
+ """used for logging metrics""" # Record the gradients of all the layers
+ preds, y, loss, acc = self._get_preds_loss_accuracy(batch)
+ self.log('validation/loss', loss)
+ self.log('validation/accuracy', acc)
+ return (preds, y)
+
+ def test_step(self, batch, batch_idx):
+ """used for logging metrics"""
+ preds, y, loss, acc = self._get_preds_loss_accuracy(batch)
+ self.log('test/loss', loss)
+ self.log('test/accuracy', acc)
+
+ def configure_optimizers(self):
+ """defines model optimizer""" # Log loss and metric
+ return Adam(self.parameters(), lr=self.lr)
+
+ def _get_preds_loss_accuracy(self, batch):
+ """convenience function since train/valid/test steps are similar"""
+ x, y = batch
+ logits = self(x)
+ preds = torch.argmax(logits, dim=1)
+ loss = self.loss(logits, y)
+ acc = accuracy(preds, y, task='multiclass', num_classes=10)
+ return (preds, y, loss, acc) # Log loss and metric # Let's return preds to use it in a custom callback # Log loss and metric
+
+ return (NatureLitModule,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Instrument Callbacks to log additional things at certain points in your code
+ """)
+ return
+
+
+@app.cell
+def _(pd, wandb):
+ from lightning.pytorch.callbacks import Callback
+
+ class LogPredictionsCallback(Callback):
+
+ def __init__(self):
+ super().__init__()
+
+ def on_validation_epoch_start(self, trainer, pl_module):
+ self.batch_dfs = []
+ self.image_list = []
+ self.val_table = wandb.Table(columns=['image', 'ground_truth', 'prediction'])
+
+ def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+ """Called when the validation batch ends."""
+ x, y = batch
+ preds, y = outputs
+ self.batch_dfs.append(pd.DataFrame({'Ground Truth': y.cpu().numpy(), 'Predictions': preds.cpu().numpy()}))
+ x = x.cpu().numpy().transpose(0, 2, 3, 1)
+ for x_i, y_i, y_pred in list(zip(x, y, preds)): # Append validation predictions and ground truth to log in confusion matrix
+ self.image_list.append(wandb.Image(x_i, caption=f'Ground Truth: {y_i} - Prediction: {y_pred}'))
+ self.val_table.add_data(wandb.Image(x_i), y_i, y_pred)
+
+ def on_validation_epoch_end(self, trainer, pl_module):
+ class_names = _trainer.datamodule.nature_dataset.class_names # Add wandb.Image to a table to log at the end of validation
+ val_df = pd.concat(self.batch_dfs)
+ wandb.log({'validation_table': self.val_table, 'images_over_time': self.image_list, 'validation_conf_matrix': wandb.plot.confusion_matrix(y_true=val_df['Ground Truth'].tolist(), preds=val_df['Predictions'].tolist(), class_names=class_names)}, step=_trainer.global_step)
+ del self.batch_dfs
+ del self.val_table # Collect statistics for whole validation set and log
+
+ return (LogPredictionsCallback,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🏋️ Main Training Loop
+ """)
+ return
+
+
+@app.cell
+def _(
+ ENTITY,
+ LogPredictionsCallback,
+ NatureDatasetModule,
+ NatureLitModule,
+ PROJECT_NAME,
+ wandb,
+):
+ from lightning.pytorch.callbacks import ModelCheckpoint
+ from lightning.pytorch.loggers import WandbLogger
+ from lightning.pytorch import Trainer
+ wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type='training', config={'model_name': 'squeezenet', 'batch_size': 16})
+ _wandb_logger = WandbLogger(log_model='all', checkpoint_name=f'nature-{wandb.run.id}')
+ _log_predictions_callback = LogPredictionsCallback()
+ _checkpoint_callback = ModelCheckpoint(every_n_epochs=1)
+ _model = NatureLitModule(model_name=wandb.config['model_name'])
+ _nature_module = NatureDatasetModule(wandb_run=_wandb_logger.experiment, artifact_name_alias='Nature_100:latest', local_target_dir='Nature_100:latest', batch_size=wandb.config['batch_size'], input_size=_model.input_size)
+ _nature_module.setup()
+ _trainer = Trainer(logger=_wandb_logger, callbacks=[_log_predictions_callback, _checkpoint_callback], max_epochs=5, log_every_n_steps=5)
+ _trainer.fit(_model, datamodule=_nature_module)
+ wandb.finish() # Access hyperparameters downstream to instantiate models/datasets # W&B integration
+ return ModelCheckpoint, Trainer, WandbLogger
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Syncing with W&B Offline
+ If for some reason, network communication is lost during the course of training, you can always sync progress with `wandb sync`
+
+ The W&B sdk caches all logged data in a local directory `wandb` and when you call `wandb sync`, this syncs the your local state with the web app.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Retrieve a model checkpoint artifact and resume training
+ - Artifacts make it easy to track state of your training remotely and then resume training from a checkpoint
+ """)
+ return
+
+
+@app.cell
+def _():
+ #@title Enter which checkpoint you want to resume training from:
+
+ # FORM VARIABLES
+ ARTIFACT_NAME_ALIAS = "nature-oyxk79m1:v4" #@param {type:"string"}
+ return (ARTIFACT_NAME_ALIAS,)
+
+
+@app.cell
+def _(
+ ARTIFACT_NAME_ALIAS,
+ ENTITY,
+ LogPredictionsCallback,
+ ModelCheckpoint,
+ NatureDatasetModule,
+ NatureLitModule,
+ PROJECT_NAME,
+ Trainer,
+ WandbLogger,
+ wandb,
+):
+ wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type='resume_training')
+ model_chkpt_art = wandb.use_artifact(f'{ENTITY}/{PROJECT_NAME}/{ARTIFACT_NAME_ALIAS}')
+ model_chkpt_art.download()
+ logging_run = model_chkpt_art.logged_by()
+ # Retrieve model checkpoint artifact and restore previous hyperparameters
+ wandb.config = logging_run.config
+ artifact_name = ARTIFACT_NAME_ALIAS.split(':')[0] # Can change download directory by adding `root`, defaults to "./artifacts"
+ _wandb_logger = WandbLogger(log_model='all', checkpoint_name=artifact_name)
+ _log_predictions_callback = LogPredictionsCallback()
+ _checkpoint_callback = ModelCheckpoint(every_n_epochs=1)
+ # Can create a new artifact name or continue logging to the old one
+ _model = NatureLitModule.load_from_checkpoint(f'./artifacts/{ARTIFACT_NAME_ALIAS}/model.ckpt')
+ _nature_module = NatureDatasetModule(wandb_run=_wandb_logger.experiment, artifact_name_alias='Nature_100:latest', local_target_dir='Nature_100:latest', batch_size=wandb.config['batch_size'], input_size=_model.input_size)
+ _nature_module.setup()
+ _trainer = Trainer(logger=_wandb_logger, callbacks=[_log_predictions_callback, _checkpoint_callback], max_epochs=10, log_every_n_steps=5)
+ _trainer.fit(_model, datamodule=_nature_module)
+ wandb.finish() # W&B integration
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Model Registry
+ After logging a bunch of checkpoints across multiple runs during experimentation, now comes time to hand-off the best checkpoint to the next stage of the workflow (e.g. testing, deployment).
+
+ The model registry offers a centralized place to house the best checkpoints for all your model tasks. Any `model` artifact you log can be "linked" to a Registered Model. Here are the steps to start using the model registry for more organized model management:
+ 1. Access your team's model registry by going the team page and selecting `Model Registry`
+ 
+
+ 2. Create a new Registered Model.
+ 
+
+ 3. Go to the artifacts tab of the project that holds all your model checkpoints
+ 
+
+ 4. Click "Link to Registry" for the model artifact version you want. (Alternatively you can [link a model via api](https://docs.wandb.ai/guides/models) with `wandb.run.link_artifact`)
+
+ **A note on linking:** The process of linking a model checkpoint is akin to "bookmarking" it. Each time you link a new model artifact to a Registered Model, this increments the version of the Registered Model. This helps delineate the model development side of the workflow from the model deployment/consumption side. The globally understood version/alias of a model should be unpolluted from all the experimental versions being generated in R&D and thus the versioning of a Registered Model increments according to new "bookmarked" models as opposed to model checkpoint logging.
+
+ ### Create a Centralized Hub for all your models
+ - Add a model card, tags, slack notifactions to your Registered Model
+ - Change aliases to reflect when models move through different phases
+ - Embed the model registry in reports for model documentation and regression reports. See this report as an [example](https://api.wandb.ai/links/wandb-smle/r82bj9at)
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/pytorch-simple-pytorch-integration/pytorch_simple_pytorch_integration.py b/marimo/convert/pytorch-simple-pytorch-integration/pytorch_simple_pytorch_integration.py
new file mode 100644
index 00000000..1fa7dfe7
--- /dev/null
+++ b/marimo/convert/pytorch-simple-pytorch-integration/pytorch_simple_pytorch_integration.py
@@ -0,0 +1,693 @@
+# /// script
+# dependencies = ["onnx", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥 = W&B ➕ PyTorch
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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):
+ mo.md(r"""
+ # Example Gallery
+
+ See examples of projects tracked and visualized with W&B in our gallery, [Fully Connected →](https://wandb.me/fc)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Basic Setup
+ 1. **Projects**: Log multiple runs to a project to compare them. `wandb.init(project="project-name")`
+ 2. **Groups**: For multiple processes or cross validation folds, log each process as a runs and group them together. `wandb.init(group='experiment-1')`
+ 3. **Tags**: Add tags to track your current baseline or production model.
+ 4. **Notes**: Type notes in the table to track the changes between runs.
+ 5. **Reports**: Take quick notes on progress to share with colleagues and make dashboards and snapshots of your ML projects.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Advanced Setup
+ 1. [Environment variables](https://docs.wandb.com/library/environment-variables): Set API keys in environment variables so you can run training on a managed cluster.
+ 2. [Offline mode](https://docs.wandb.com/library/technical-faq#can-i-run-wandb-offline): Use `dryrun` mode to train offline and sync results later.
+ 3. [On-prem](https://docs.wandb.com/self-hosted): 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.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/scikit-w-b-k-means-clustering/scikit_w_b_k_means_clustering.py b/marimo/convert/scikit-w-b-k-means-clustering/scikit_w_b_k_means_clustering.py
new file mode 100644
index 00000000..65b9fe64
--- /dev/null
+++ b/marimo/convert/scikit-w-b-k-means-clustering/scikit_w_b_k_means_clustering.py
@@ -0,0 +1,200 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # 🏋️♀️ W&B + 🧪 Scikit-learn
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+
+ ## What this notebook covers:
+ * Easy integration of Weights and Biases with Scikit.
+ * W&B Scikit plots for model interpretation and diagnostics for regression, classification, and clustering.
+
+ **Note**: Sections starting with _Step_ are all you need to integrate W&B to existing code.
+
+ ## The interactive W&B Dashboard will look like this:
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Author: [@SauravMaheshkar](https://twitter.com/MaheshkarSaurav)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Packages 📦 and Basic Setup
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install Packages
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Install the latest version of wandb client 🔥🔥
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell
+def _():
+ import numpy as np
+ from sklearn import datasets
+ from sklearn.cluster import KMeans
+
+ return KMeans, datasets, np
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Project Configuration using **`wandb.config`**
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ log to your weights and biases account
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(np, wandb):
+ # Initialize the run
+ run = wandb.init(project='simple-scikit')
+
+ # Feel free to change these and experiment !!
+ config = wandb.config
+ config.seed = 42
+ config.n_clusters = 3
+ config.dataset = 'iris'
+ config.labels=['Setosa', 'Versicolour', 'Virginica']
+
+ # Set random seed
+ np.random.seed(config.seed)
+
+ # Update the config
+ wandb.config.update(config)
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 💿 The Dataset
+ ---
+ """)
+ return
+
+
+@app.cell
+def _(datasets):
+ # Download the Iris dataset from sklearn
+ iris = datasets.load_iris()
+
+ # Get our data and target variables
+ X = iris.data
+ y = iris.target
+ return (X,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ✍️ Model Architecture & Training
+
+ ---
+ """)
+ return
+
+
+@app.cell
+def _(KMeans, X, config, wandb):
+ # Define the Estimator
+ est = KMeans(n_clusters = config.n_clusters, random_state = config.seed)
+
+ # Compute the Clusters
+ est.fit(X)
+
+ # Update our config with the cluster centers
+ wandb.config.update({'labels' : est.cluster_centers_})
+
+ # Plot the Clusters to W&B
+ wandb.sklearn.plot_clusterer(est, X, cluster_labels = est.fit_predict(X), labels=config.labels, model_name='KMeans')
+
+ # Finish the W&B Process
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/scikit-wandb-decision-tree/scikit_wandb_decision_tree.py b/marimo/convert/scikit-wandb-decision-tree/scikit_wandb_decision_tree.py
new file mode 100644
index 00000000..3d1d1e9f
--- /dev/null
+++ b/marimo/convert/scikit-wandb-decision-tree/scikit_wandb_decision_tree.py
@@ -0,0 +1,186 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Author: [@SauravMaheshkar](https://twitter.com/MaheshkarSaurav)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Packages 📦 and Basic Setup
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install Packages
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # ## Install Sklearn
+ # !pip install -U scikit-learn
+ # ## Install the latest version of wandb client 🔥🔥
+ # !pip install -q --upgrade wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Project Configuration using **`wandb.config`**
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ ## Importing Libraries
+ from sklearn.datasets import load_iris
+ from sklearn.model_selection import train_test_split
+ from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
+
+ return (
+ DecisionTreeClassifier,
+ DecisionTreeRegressor,
+ load_iris,
+ train_test_split,
+ wandb,
+ )
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(wandb):
+ # Initialize the run
+ run = wandb.init(project='simple-scikit')
+
+ # Feel free to change these and experiment !!
+ config = wandb.config
+ config.max_depth = 5
+ config.min_samples_split = 2
+ config.clf_criterion = "gini"
+ config.reg_criterion = "mse"
+ config.splitter = "best"
+ config.dataset = "iris"
+ config.test_size = 0.2
+ config.random_state = 42
+ config.labels =['setosa', 'versicolor', 'virginica']
+
+ # Update the config
+ wandb.config.update(config)
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 💿 Dataset
+ ---
+ """)
+ return
+
+
+@app.cell
+def _(load_iris):
+ ## Loading the Dataset
+ iris = load_iris(return_X_y = True, as_frame= True)
+ dataset = iris[0]
+ target = iris[1]
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ✍️ Model Architecture
+ ---
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Classification
+ """)
+ return
+
+
+@app.cell
+def _(DecisionTreeClassifier, config, load_iris, train_test_split, wandb):
+ _X, _y = load_iris(return_X_y=True)
+ _x_train, _x_test, _y_train, _y_test = train_test_split(_X, _y, test_size=config.test_size, random_state=config.random_state)
+ clf = DecisionTreeClassifier(max_depth=config.max_depth, min_samples_split=config.min_samples_split, criterion=config.clf_criterion, splitter=config.splitter)
+ clf = clf.fit(_x_train, _y_train)
+ y_pred = clf.predict(_x_test)
+ # Visualize Confustion Matrix
+ wandb.sklearn.plot_confusion_matrix(_y_test, y_pred, config.labels)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Regression
+ """)
+ return
+
+
+@app.cell
+def _(DecisionTreeRegressor, config, load_iris, train_test_split, wandb):
+ _X, _y = load_iris(return_X_y=True)
+ _x_train, _x_test, _y_train, _y_test = train_test_split(_X, _y, test_size=config.test_size, random_state=config.random_state)
+ reg = DecisionTreeRegressor(max_depth=config.max_depth, min_samples_split=config.min_samples_split, criterion=config.reg_criterion, splitter=config.splitter)
+ reg = reg.fit(_x_train, _y_train)
+ # All regression plots
+ wandb.sklearn.plot_regressor(reg, _x_train, _x_test, _y_train, _y_test, model_name='DecisionTreeRegressor')
+ return
+
+
+@app.cell
+def _(wandb):
+ # Finish the W&B Process
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/simpletransformers-simpletransformersqa/simpletransformers_simpletransformersqa.py b/marimo/convert/simpletransformers-simpletransformersqa/simpletransformers_simpletransformersqa.py
new file mode 100644
index 00000000..a17f31a7
--- /dev/null
+++ b/marimo/convert/simpletransformers-simpletransformersqa/simpletransformers_simpletransformersqa.py
@@ -0,0 +1,689 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B 💘 SimpleTransformers
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## What this notebook covers
+
+ In this notebook we show you how to integrate
+ [Weights & Biases](https://wandb.ai/site)
+ with your
+ [SimpleTransformers](https://github.com/ThilinaRajapakse/simpletransformers)
+ code to add experiment tracking to your pipeline. This includes:
+
+ 1. dataset and model versioning with W&B,
+ [Artifacts](https://docs.wandb.ai/guides/artifacts)
+ 2. storing configuration, hyperparameters, system metrics, and model metrics in an [interactive dashboard](https://docs.wandb.ai/guides/track/app), and
+ 3. examining evaluation outputs of your model using
+ [W&B Tables](https://docs.wandb.ai/guides/data-vis).
+
+ We'll add these features to a typical NLP pipeline:
+ training a question-answering model with a DistilBERT backbone on (a very small subset of) the Stanford QUestion Answering Dataset ([SQuAD](https://rajpurkar.github.io/SQuAD-explorer/)).
+ The extra code required is indicated in the notebook with headers that start with "Step".
+
+ We'll end up with an interactive dashboard ([link](https://wandb.ai/wandb/SimpleTransformers-QA?workspace=user-prashanthkurella))
+ for our QA experiments that looks like this:
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Step 0 : Install Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's get started by installing W&B.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -Uq wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Step 1: Import `wandb`, log in, and set the project name
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now that we have installed W&B, let's import it and log in.
+ If you don't have a W&B account, you'll be prompted to create one.
+ W&B is free for open projects, just like GitHub.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ wandb_project = "SimpleTransformers-QA"
+ return wandb, wandb_project
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Install SimpleTransformers
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ SimpleTransformers is a library for quickly spinning up Transformer models
+ for natural language processing tasks.
+ It has an easy to use interface and requires only a minimal amount of code.
+
+ You can look at [the docs](https://simpletransformers.ai/) for more about what SimpleTransformers can do.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install simpletransformers
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download the SQuAD Dataset
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We'll use a subset of the SQuAD Dataset to train a question-answering model.
+
+ The SQuAD Dataset consists of a context and set of questions relevant to the context.
+ Our model's task is to understand the context and use it to answer the questions.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # data_dir = "data"
+ # raw_data_dir = data_dir + "/" + "raw"
+ #
+ # !mkdir -p {raw_data_dir}
+ # !curl https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json --output {raw_data_dir}/train.json
+ # !curl https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v2.0.json --output {raw_data_dir}/eval.json
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Step 2: Log the raw dataset
+
+ In order to make our work more reproducible and portable,
+ we'll store, version, and distribute our dataset with
+ [W&B Artifacts](https://docs.wandb.ai/guides/artifacts).
+
+ To log an Artifact, we need to start a `wandb.Run`
+ with [`wandb.init`](https://docs.wandb.ai/guides/track/launch).
+ We'll give it the `upload-raw-dataset` job type.
+ """)
+ return
+
+
+@app.cell
+def _(raw_data_dir, wandb, wandb_project):
+ # initialize a run to log the datasets
+ _run = wandb.init(project=wandb_project, job_type='upload-raw-dataset')
+ raw_data_artifact = wandb.Artifact('raw-data', 'dataset')
+ raw_data_artifact.add_dir(raw_data_dir)
+ _run.log_artifact(raw_data_artifact)
+ # log the raw data
+ # finish the run
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Perform the train-test split on the raw dataset
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Datasets are typically transformed in an ML pipeline.
+
+ For example, when training models, we usually want to split out some data to hold for validation.
+
+ With W&B Artifacts we can track those splits
+ and then reuse them when debugging models, or reproducing and extending results.
+ """)
+ return
+
+
+@app.cell
+def _(data_dir, raw_data_dir, subprocess):
+ import json
+ import random
+
+ # make the data dirs for splits
+ split_data_dir = data_dir + "/" + "split"
+ #! mkdir -p {split_data_dir}
+ subprocess.call(['mkdir', '-p', str(split_data_dir)])
+
+ # shuffle and subset train data
+ with open(f"{raw_data_dir}/train.json", "r") as f:
+ train_data = json.load(f)
+ train_data = [item for topic in train_data["data"] for item in topic["paragraphs"] ]
+ random.shuffle(train_data)
+ train_data = train_data[:int(len(train_data) * 0.01)]
+
+ # shuffle and subset eval data
+ with open(f"{raw_data_dir}/eval.json", "r") as f:
+ eval_data = json.load(f)
+ eval_data = [item for topic in eval_data["data"] for item in topic["paragraphs"] ]
+ random.shuffle(eval_data)
+ eval_data = eval_data[:int(len(eval_data) * 0.01)]
+
+ # write the subsets to disk
+ with open(f"{split_data_dir}/train.json", "w") as f:
+ json.dump(train_data, f)
+ with open(f"{split_data_dir}/eval.json", "w") as f:
+ json.dump(eval_data, f)
+ return eval_data, split_data_dir
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Step 3: Log the train-test split we'll be using
+ """)
+ return
+
+
+@app.cell
+def _(split_data_dir, wandb, wandb_project):
+ # initialize a run to save the datasets
+ _run = wandb.init(project=wandb_project, job_type='split-dataset')
+ _run.use_artifact('raw-data:latest')
+ train_artifact = wandb.Artifact('train-data', 'dataset')
+ train_artifact.add_file(f'{split_data_dir}/train.json')
+ _run.log_artifact(train_artifact)
+ # use the raw data artifact
+ eval_artifact = wandb.Artifact('eval-data', 'dataset')
+ eval_artifact.add_file(f'{split_data_dir}/eval.json')
+ # log the training data as an artifact
+ _run.log_artifact(eval_artifact)
+ # log the evaluation data as an artifact
+ # finish logging the data logging run
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ By using Artifacts, we can track which runs produced and used particular resources,
+ like datasets, models, and analyses.
+ The relationships are tracked as a
+ [graph](https://docs.wandb.ai/ref/app/pages/project-page#graph-view-panel).
+
+ You can view and interact with a complete artifact graph
+ for this project in the browser
+ [here](https://wandb.ai/wandb/SimpleTransformers-QA/artifacts/run_table/run-3n08kirq-evalresults/ce84b13b2961e7b30e13/graph).
+
+ You'll see square nodes, representing runs,
+ and circular nodes, representing generated artifacts.
+ Arrows connect runs to the artifacts they generated
+ and artifacts to the runs that use them.
+
+ In the screenshot below,
+ see if you can find the runs used to upload and split the dataset
+ and the dataset artifacts that those runs generated.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Configure model training
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ SimpleTransformers makes it easy to run and configure your transformer model training:
+ `train_args` are ["all you need"](https://arxiv.org/abs/1706.03762)
+ to train your model.
+ """)
+ return
+
+
+@app.cell
+def _():
+ train_args = {
+ "learning_rate": 3e-5, # learning rate of our model
+ "num_train_epochs": 2, # number of epochs
+ "max_seq_length": 384, # maximum sequence length in tokens
+ "doc_stride": 128, # stride when processing sentences
+ "overwrite_output_dir": True, # overwrite the output directory
+ "reprocess_input_data": False, # reprocess the input data
+ "train_batch_size": 16, # training batch size
+ "gradient_accumulation_steps": 1, # steps before applying gradients
+ "evaluate_during_training": True, # run evaluation during training
+ "evaluate_during_training_steps": 40, # steps in training before eval
+ "save_eval_checkpoints": False, # save evaluation checkpoints
+ "eval_batch_size": 16, # evaluation batch size
+ }
+ return (train_args,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Step 4: Include `wandb_project` in `train_args` to use W&B for logging our training progress
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ SimpleTransformers comes with W&B logging built in -- no extra code required.
+ To enable it you just need to pass the `wandb_project` argument.
+
+ You can also customize what's passed to the `wandb.init` function
+ used to launch your training run
+ with the `wandb_kwargs` argument.
+ Refer to the docs
+ [here](https://docs.wandb.ai/guides/track/launch)
+ for more info.
+ """)
+ return
+
+
+@app.cell
+def _(train_args, wandb_project):
+ train_args.update(
+ {
+ "logging_steps": 1, # number of steps before logging
+ "wandb_project": wandb_project, # wandb project name
+ "wandb_kwargs": {"job_type": "training"} # additional args for wandb init
+ }
+ )
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Initialize the model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Another killer feature of SimpleTransformers is that it comes with a bunch of
+ implementations of widely-used transformer architectures, like BERT, ALBERT, and others.
+
+ It also includes utilities for downloading their pretrained versions and adapting
+ them to specific tasks.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture --no-display
+ # from simpletransformers.question_answering import QuestionAnsweringModel
+ #
+ # # initialize the model with a distilbert backbone
+ # model = QuestionAnsweringModel("distilbert", "distilbert-base-cased", args=train_args)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train the model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Every time you call the `train_model` function,
+ you launch a new experiment.
+
+ W&B prints out the links to
+ [project-level](https://docs.wandb.ai/ref/app/pages/project-page)
+ and
+ [run-level](https://docs.wandb.ai/ref/app/pages/run-page)
+ dashboards.
+
+ Click on those links to view the training progress
+ and compare to other experiments.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture --no-display
+ # model.train_model(train_data, eval_data=eval_data)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can reorganize your workspace into an interactive dashboard
+ to share with team members or put in your portfolio.
+
+ Below is a screenshot of a dashboard made for this project.
+
+ You can view and interact with it in your browser
+ [here](https://wandb.ai/wandb/SimpleTransformers-QA?workspace=user-prashanthkurella).
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Custom Logging for SimpleTransformers
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ SimpleTransformers automatically logs important metrics to W&B.
+
+ You can also customize what you log using two methods:
+
+ 1. [Resuming](https://docs.wandb.ai/guides/track/advanced/resuming) the run,
+ "restarting" the experiment so that you can log additional stuff, including more training.
+ 2. Using the [`wandb.api`](https://docs.wandb.ai/guides/track/public-api-guide) to update existing runs with additional metadata.
+
+ We show both below.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Use resuming to add model checkpoints
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Runs that have finished can be resumed
+ so that additional information can be added to an experiment.
+ For example, you might be using
+ [pre-emptible compute](https://www.parkmycloud.com/blog/google-preemptible-vms/)
+ where training runs can be stopped prematurely.
+
+ Here we use it to log the model checkpoints to the training run,
+ since it was responsible for creating them.
+ """)
+ return
+
+
+@app.cell
+def _(model, wandb, wandb_project):
+ import os
+ with wandb.init(id=model.wandb_run_id, resume='allow', project=wandb_project) as _training_run:
+ for dir in sorted(os.listdir('outputs')):
+ if 'checkpoint' in dir:
+ artifact = wandb.Artifact('model-checkpoints', type='checkpoints')
+ artifact.add_dir('outputs' + '/' + dir)
+ _training_run.log_artifact(artifact)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Use resuming to add evaluation results as a `Table`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To evaluate models and their performance,
+ it's important to be able to visualize and analyze model predictions.
+ W&B supports this workflow with
+ [Tables](https://docs.wandb.ai/guides/data-vis).
+
+ Here, we'll grab our model's predictions on the evaluation data,
+ convert them into a pandas `DataFrame`,
+ and then log them to W&B as a `Table` attached to the resumed run.
+
+ For more on using Tables for NLP, check out our
+ [video guide](https://www.youtube.com/watch?v=756JcKiDvqo)
+ on applying Tables to the
+ [GoEmotions dataset](https://arxiv.org/abs/2005.00547).
+ """)
+ return
+
+
+@app.cell
+def _(eval_data, model):
+ _, outputs = model.eval_model(eval_data)
+ return (outputs,)
+
+
+@app.cell
+def _(eval_data, model, outputs, wandb, wandb_project):
+ import pandas as pd
+ eval_data_df = pd.DataFrame(columns=['id', 'question', 'context'])
+ # create an empty dataframe
+ for context in eval_data:
+ for qas in context['qas']:
+ eval_data_df = eval_data_df.append([{'id': qas['id'], 'context': context['context'], 'question': qas['question']}])
+ eval_data_df = eval_data_df.reset_index(drop=True)
+ results = pd.DataFrame(columns=['id', 'predicted_answer', 'actual_answer', 'category'])
+ for entry in outputs['correct_text']:
+ results = results.append([{'id': entry, 'predicted_answer': outputs['correct_text'][entry], 'actual_answer': outputs['correct_text'][entry], 'category': 'correct'}])
+ for entry in outputs['similar_text']:
+ # load the eval data into the dataframe
+ results = results.append([{'id': entry, 'predicted_answer': outputs['similar_text'][entry]['predicted'], 'actual_answer': outputs['similar_text'][entry]['truth'], 'category': 'similar'}])
+ for entry in outputs['incorrect_text']:
+ results = results.append([{'id': entry, 'predicted_answer': outputs['incorrect_text'][entry]['predicted'], 'actual_answer': outputs['incorrect_text'][entry]['truth'], 'category': 'incorrect'}])
+ results = results.reset_index(drop=True)
+ results = eval_data_df.set_index('id').join(results.set_index('id'))
+ results = results.drop_duplicates()
+ with wandb.init(resume=model.wandb_run_id, project=wandb_project) as _training_run:
+ # reset index for clear indexing
+ # create an empty results data frame
+ # load all the correctly predicted answers
+ # load all the similar answers
+ # load all the incorrect answers
+ # join the evaluation data with the predictions
+ # resume the training run and log the table
+ _training_run.log({'eval-results': wandb.Table(dataframe=results)})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Inside the W&B web app,
+ you can interact with logged table data
+ to perform post-hoc analyses,
+ including filtering, grouping, and computing derived metrics.
+
+ Below is a screenshot of a table that compares the model's outputs
+ to the actual ground truth across multiple runs.
+
+ You can view and interact with it in your browser
+ [here](https://wandb.ai/wandb/SimpleTransformers-QA/reports/Shared-panel-21-09-10-12-09-00--VmlldzoxMDExNDQw).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Use the API to attach the train-test splits to the training run
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Logged information from experiments and workflows
+ often needs to be programmatically accessed or updated.
+ For those tasks, we provide
+ [a public API](https://docs.wandb.ai/guides/track/public-api-guide).
+
+ Here we'll use it to update the training run with
+ the dataset artifacts that we uploaded earlier.
+ """)
+ return
+
+
+@app.cell
+def _(model, wandb, wandb_project):
+ # initialize the wandb api object
+ api = wandb.Api()
+ _training_run = api.run(wandb_project + '/' + model.wandb_run_id)
+ # retrieve our training run
+ train_data_artifact = api.artifact(wandb_project + '/' + 'train-data:latest')
+ eval_data_artifact = api.artifact(wandb_project + '/' + 'eval-data:latest')
+ # retrieve the artifacts we'll be using
+ _training_run.use_artifact(train_data_artifact)
+ # mark the training run as using the training and eval data artifacts
+ _training_run.use_artifact(eval_data_artifact)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/spacy-spacy-v3-and-w-b-sweeps/spacy_spacy_v3_and_w_b_sweeps.py b/marimo/convert/spacy-spacy-v3-and-w-b-sweeps/spacy_spacy_v3_and_w_b_sweeps.py
new file mode 100644
index 00000000..480c409d
--- /dev/null
+++ b/marimo/convert/spacy-spacy-v3-and-w-b-sweeps/spacy_spacy_v3_and_w_b_sweeps.py
@@ -0,0 +1,191 @@
+# /// script
+# dependencies = ["spacy", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [Weights & Biases](https://wandb.ai/site) makes running collaborative machine learning projects a breeze. You can focus on what you're trying to experiment with, and W&B will take on the burden of keeping track of everything. If you want to review a loss plot, download the latest model for production, or just see which configurations produced a certain model, W&B is your friend. There's also a bunch of features to help you and your team collaborate, like having a shared dashboard and sharing interactive reports.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [spaCy](https://spacy.io/) can serve a lot of your Natural Language Processing (NLP) needs out-of-the-box. This includes Named Entity Recognition (NER), Part of Speech tagging, text classification and more. Even better, these components are all customizable, extendable and composable.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Hyperparameter Search with W&B + spaCy
+
+ This notebook accompanies and implements a
+ [blog post](https://wandb.ai/wandb/wandb_spacy_sweeps/reports/Hyperparameter-Search-with-spaCy-and-Weights-Biases--Vmlldzo5NDA2MjE)
+ on automating hyperparameter search for spaCy models using W&B sweeps.
+
+ Run the cells below to start your own hyperparameter search.
+
+ The code is organized as a [spaCy `project`](https://spacy.io/usage/projects/), available [here](https://github.com/explosion/projects/tree/v3/integrations/wandb).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: spacy >= 3.0.6 !pip install -qq "spacy >= 3.0.6"
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ✍️ Login to W&B
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Clone the project with `spacy project clone`
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python -m spacy project clone integrations/wandb
+ subprocess.call(['python', '-m', 'spacy', 'project', 'clone', 'integrations/wandb'])
+ import os
+ os.chdir('wandb/')
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install the project-specific dependencies
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python -m spacy project run install
+ subprocess.call(['python', '-m', 'spacy', 'project', 'run', 'install'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download the project-specific assets
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python -m spacy project assets
+ subprocess.call(['python', '-m', 'spacy', 'project', 'assets'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Run the hyperparameter search
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Running the cell below will produce a Sweeps page 🧹 link which you can follow to see the all the metrics of your runs.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python -m spacy project run parameter-search
+ subprocess.call(['python', '-m', 'spacy', 'project', 'run', 'parameter-search'])
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/spacy-spacy-v3-and-w-b/spacy_spacy_v3_and_w_b.py b/marimo/convert/spacy-spacy-v3-and-w-b/spacy_spacy_v3_and_w_b.py
new file mode 100644
index 00000000..8fcb58ab
--- /dev/null
+++ b/marimo/convert/spacy-spacy-v3-and-w-b/spacy_spacy_v3_and_w_b.py
@@ -0,0 +1,371 @@
+# /// script
+# dependencies = ["spacy", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Text Classification with spaCy v3 and W&B
+
+ In this notebook, we'll be training a multi-label CNN text classifier using spaCy v3 on Google's GoEmotions dataset. We'll be tracking our models' progress and saving its outputs using Weights and Biases (W&B).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **TextCat GoEmotions Project**
+
+ The notebook is based on the [TextCat GoEmotions](https://github.com/explosion/projects/tree/v3/tutorials/textcat_goemotions) spaCy Project. In spaCy's full TextCat GoEmotions project on github you can do also pre-processing, evaluation and even package the project and visualize the outputs!
+
+ **SpaCy Projects**
+
+ [SpaCy Projects](https://spacy.io/usage/projects) let you manage and share end-to-end spaCy workflows for different use cases and domains, and orchestrate training, packaging and serving your custom pipelines.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: spacy >= 3.0.6 !pip install -qq "spacy >= 3.0.6"
+ # packages added via marimo's package management: wandb !pip install -qU wandb
+ return
+
+
+@app.cell
+def _():
+ import os
+ from pathlib import Path
+
+ return (Path,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Google's GoEmotions Dataset
+
+ Here we'll be doing multilabel text classification on the [GoEmotions dataset](https://github.com/google-research/google-research/tree/master/goemotions), a corpus of 58k curated comments extracted from Reddit, with human annotations of 28 different emotions. Here is an example from the [GoEmotions paper](https://arxiv.org/pdf/2005.00547v2.pdf):
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### ✍️ Login to wandb ✍️
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### ✍️ Download from W&B Artifacts ✍️
+
+ To keep things focused on the training, we will download a pre-processed version of the GoEmotions dataset using Weights and Biases Artifacts. With Artifacts you get **100GB free storage** to use for your data and model versioning or however you like!
+
+ #### Corpora
+ The train, dev and test corpora are stored in the `.spacy` [binary format](https://spacy.io/api/data-formats#training) that spaCy's training function expects. To see how Googles `.txt` files were converted to `.spacy` files have a look at this [gist file here](https://gist.github.com/morganmcg1/a43842b847e2ff7dc78d2c3e5990bb96)
+
+ #### Training Config
+
+ [Training in spaCy v3](https://spacy.io/usage/training) revolves around using a configuration file to set up your model and any hyperparameters you might need. In this notebook we will use a modified version of the [`cnn.cfg`](https://github.com/explosion/projects/blob/v3/tutorials/textcat_goemotions/configs/cnn.cfg) from the spaCy project repo
+
+ #### Start Download
+
+ Lets download the directory which contains our 3 corpora and 1 configuration file
+ """)
+ return
+
+
+@app.cell
+def _(Path):
+ # W&B Artifact naming convention: `wandb_entity/wandb_project/artifact_name:version`
+ spacy_artifact = 'wandb/spacy/spacy_demo:v3'
+
+ # Our output directory name
+ spacy_dir = Path("my_spacy_demo")
+ return spacy_artifact, spacy_dir
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ By creating a wandb run context we can easily download artifacts
+ """)
+ return
+
+
+@app.cell
+def _(spacy_artifact, spacy_dir, wandb):
+ with wandb.init(project='spacy_demo') as run: # "config" is optional here
+ artifact = run.use_artifact(spacy_artifact)
+ _ = artifact.download(spacy_dir)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training a Classifier
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### spaCy Config
+
+ **✍️ WandbLogger ✍️**
+
+ The only change made to the configuration file from the original spaCy project was to change the training logger. Instead of using the default console logger, we can use the [spaCy W&B logger](https://spacy.io/api/top-level#WandbLogger) to get a richer display of our training metrics. This can be done by simply adding the following code to our config file:
+
+ ```
+ [training.logger]
+ @loggers = "spacy.WandbLogger.v2"
+
+ # Our W&B Project name
+ project_name = "spacy_demo"
+
+ # Any config data you do not want logged to W&B
+ remove_config_values = ["paths.train", "paths.dev", "corpora.train.path", "corpora.dev.path"]
+
+ # Optional, log this dataset folder to W&B Artifacts
+ log_dataset_dir = "./my_spacy_demo/corpus"
+
+ # Optional, log the model every N steps to W&B Artifacts
+ model_log_interval = 1200
+ ```
+
+ **Training Params**
+
+ In our training we can define our training regime. This example will train until the eval score hasn't improved for 400 steps or until it has trained for 10,000 steps. We can also choose to accumulate gradients if we like, in this case we will not
+
+ ```
+ [training]
+ ...
+ max_epochs = 0
+ patience = 400
+ max_steps = 6000
+ eval_frequency = 200
+ accumulate_gradient = 1
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Model Architecture
+
+ Our config also defines with model architecture to use, in this case we are using spaCy's [`TextCatCNN`](https://spacy.io/api/architectures#TextCatCNN) architecture which is:
+
+ > a neural network model where token vectors are calculated using a CNN. The vectors are mean pooled and used as features in a feed-forward network.
+
+ Our embedding layer is defined [`MultiHashEmbed`](https://spacy.io/api/architectures#MultiHashEmbed) which:
+ > lets the model take into account some subword information, without construction a fully character-based representation.
+
+ You can read more about spaCy model architectures [here](https://spacy.io/usage/layers-architectures#_title)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Tokenizer
+
+ This config uses spaCy's [default tokenizer](https://spacy.io/api/tokenizer#_title)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 🏋️ Train
+
+ Once we have our configuration setup, we can call `spacy train` to begin training our models. If we like we can override the data paths in the config, as well as define what GPU to run the training on in the case that we have multiple GPUs
+
+ Once training starts we will be able to track our runs by clicking on the 🚀 run link that is output
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python -m spacy train my_spacy_demo/configs/cnn.cfg --paths.train my_spacy_demo/corpus/train.spacy --paths.dev my_spacy_demo/corpus/dev.spacy -o my_spacy_demo/training/cnn --gpu-id 0
+ subprocess.call(['python', '-m', 'spacy', 'train', 'my_spacy_demo/configs/cnn.cfg', '--paths.train', 'my_spacy_demo/corpus/train.spacy', '--paths.dev', 'my_spacy_demo/corpus/dev.spacy', '-o', 'my_spacy_demo/training/cnn', '--gpu-id', '0'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using our Trained Model
+
+ We can now quickly use our trained classifier model to retrieve predictions on the emotion of different texts. First we get write texts
+ """)
+ return
+
+
+@app.cell
+def _():
+ texts = ["This is a fabulous idea, this made be so happy and excited, I want to jump for joy",
+ "This movie was terrifying, I jumped out of my seat I was so scared, I never want to watch this again"]
+ return (texts,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We load our spaCy model from disk
+ """)
+ return
+
+
+@app.cell
+def _():
+ import spacy
+ nlp = spacy.load("my_spacy_demo/training/cnn/model-best")
+ return (nlp,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ And now we can easily get the model's score for each of the emotions by using spaCys [pipeline method](https://spacy.io/usage/processing-pipelines#pipelines), `nlp.pipe`
+ """)
+ return
+
+
+@app.cell
+def _(nlp, texts):
+ category_scores = [doc.cats for doc in nlp.pipe(texts)]
+ category_scores[0]
+ return (category_scores,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ While training our multi-label classifier, we defined our score threshold as 0.5. Using that we can now get the top emotions predicted for our 2 samples texts
+ """)
+ return
+
+
+@app.cell
+def _(category_scores):
+ thresh = 0.5
+ for d in category_scores:
+ print(dict((k, v) for k, v in d.items() if v >= thresh))
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Not bad!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # ✅ Done
+
+ In this notebook we have downloaded our dataset from, trained a model using spacy v3 and saved our datast and model checkpoints to W&B Artifacts.
+
+ ### W&B Integrations
+
+ Weights and Biases has been integrated with all of the major machine learning libraries, including Hugging Face, Keras, Pytorch Lightning, Fastai and more. You can see the full list in our [integrations docs here](https://docs.wandb.ai/guides/integrations)
+
+ ### More to come!
+ This is just a quick intro to the power of spaCy v3 and how W&B can support your ML workflow - **stay tuned for more**! In the meantime, come check out [Fully Connected](https://wandb.ai/fully-connected), a place for the ML community to share their work, collaborate and learn from each other.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/stable-baselines3-stable-baselines3-wandb-experiment-tracking/stable_baselines3_stable_baselines3_wandb_experiment_tracking.py b/marimo/convert/stable-baselines3-stable-baselines3-wandb-experiment-tracking/stable_baselines3_stable_baselines3_wandb_experiment_tracking.py
new file mode 100644
index 00000000..45c7b3a1
--- /dev/null
+++ b/marimo/convert/stable-baselines3-stable-baselines3-wandb-experiment-tracking/stable_baselines3_stable_baselines3_wandb_experiment_tracking.py
@@ -0,0 +1,147 @@
+# /// script
+# dependencies = ["pyvirtualdisplay", "stable-baselines3", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Stable Baselines 3 - Track Experiments with Weights and Biases
+
+
+
+ Github repo: https://github.com/araffin/rl-tutorial-jnrr19
+
+ Stable-Baselines3: https://github.com/DLR-RM/stable-baselines3
+
+ Documentation: https://stable-baselines.readthedocs.io/en/master/
+
+ RL Baselines3 zoo: https://github.com/DLR-RM/rl-baselines3-zoo
+
+ Weights & Biases: https://wandb.ai/site
+
+ Weights & Biases Docs: https://docs.wandb.ai/
+
+ ## Introduction
+
+ [Weights & Biases (W&B)](https://wandb.ai/site) is a tool for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+
+
+
+
+
+ In this notebook, you will learn how to track reinforcement learning experiments using W&B. In particular, W&B helps track your experiment configs, metrics, and videos of the agents playing the game. At the end, you should see a run page like https://wandb.ai/wandb/cartpole_test/runs/37ppqzxc
+
+ ## Install Dependencies and Set up Virtual Displays for Video Recordings
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! apt install python-opengl xvfb
+ subprocess.call(['apt', 'install', 'python-opengl', 'xvfb'])
+ # packages added via marimo's package management: pyvirtualdisplay stable_baselines3[extra] wandb !pip install pyvirtualdisplay stable_baselines3[extra] wandb
+ from pyvirtualdisplay import Display
+ virtual_display = Display(visible=0, size=(1400, 900))
+ virtual_display.start()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Track experiments with W&B
+
+ Here is a clean end-to-end example to run. It will prompt you to login in to W&B if you haven't.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import gym
+ from stable_baselines3 import PPO
+ from stable_baselines3.common.monitor import Monitor
+ from stable_baselines3.common.vec_env import DummyVecEnv, VecVideoRecorder
+ import wandb
+ from wandb.integration.sb3 import WandbCallback
+
+
+ config = {
+ "policy_type": "MlpPolicy",
+ "total_timesteps": 25000,
+ "env_name": "CartPole-v1",
+ }
+ run = wandb.init(
+ project="sb3",
+ config=config,
+ sync_tensorboard=True, # auto-upload sb3's tensorboard metrics
+ monitor_gym=True, # auto-upload the videos of agents playing the game
+ save_code=True, # optional
+ )
+
+
+ def make_env():
+ env = gym.make(config["env_name"])
+ env = Monitor(env) # record stats such as returns
+ return env
+
+
+ env = DummyVecEnv([make_env])
+ env = VecVideoRecorder(env, f"videos/{run.id}", record_video_trigger=lambda x: x % 2000 == 0, video_length=200)
+ model = PPO(config["policy_type"], env, verbose=1, tensorboard_log=f"runs/{run.id}")
+ model.learn(
+ total_timesteps=config["total_timesteps"],
+ callback=WandbCallback(
+ gradient_save_freq=100,
+ model_save_path=f"models/{run.id}",
+ verbose=2,
+ ),
+ )
+ run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ After finishing the cell above you should see a dashbaord similar to the gif below:
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/stylegan-nada-stylegan-nada/stylegan_nada_stylegan_nada.py b/marimo/convert/stylegan-nada-stylegan-nada/stylegan_nada_stylegan_nada.py
new file mode 100644
index 00000000..c9dcb477
--- /dev/null
+++ b/marimo/convert/stylegan-nada-stylegan-nada/stylegan_nada_stylegan_nada.py
@@ -0,0 +1,625 @@
+# /// script
+# dependencies = ["CLIP @ git+https://github.com/openai/CLIP.git", "ftfy", "regex", "tqdm", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 StyelGAN-NADA + WandB Playground 🪄🐝
+
+
+
+ **Original Implementation:** https://github.com/rinongal/StyleGAN-nada
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 1: Setup required libraries and models.
+ This may take a few minutes.
+
+ You may optionally enable downloads with pydrive in order to authenticate and avoid drive download limits when fetching pre-trained ReStyle and StyleGAN2 models.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %tensorflow_version 1.x
+
+ import os
+
+ restyle_dir = os.path.join("/content", "restyle")
+ stylegan_ada_dir = os.path.join("/content", "stylegan_ada")
+ stylegan_nada_dir = os.path.join("/content", "stylegan_nada")
+
+ output_dir = os.path.join("/content", "output")
+
+ output_model_dir = os.path.join(output_dir, "models")
+ output_image_dir = os.path.join(output_dir, "images")
+ return os, output_dir, restyle_dir, stylegan_nada_dir
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Installing Requirements
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! git clone --depth 1 https://github.com/yuval-alaluf/restyle-encoder.git $restyle_dir
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/yuval-alaluf/restyle-encoder.git', '$restyle_dir'])
+
+ #! wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip
+ subprocess.call(['wget', 'https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip'])
+ #! sudo unzip ninja-linux.zip -d /usr/local/bin/
+ subprocess.call(['sudo', 'unzip', 'ninja-linux.zip', '-d', '/usr/local/bin/'])
+ #! sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force
+ subprocess.call(['sudo', 'update-alternatives', '--install', '/usr/bin/ninja', 'ninja', '/usr/local/bin/ninja', '1', '--force'])
+
+ # packages added via marimo's package management: ftfy regex tqdm wandb !pip install ftfy regex tqdm wandb
+ # packages added via marimo's package management: git+https://github.com/openai/CLIP.git !pip install git+https://github.com/openai/CLIP.git
+
+ #! git clone --depth 1 https://github.com/NVlabs/stylegan2-ada/ $stylegan_ada_dir
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/NVlabs/stylegan2-ada/', '$stylegan_ada_dir'])
+ #! git clone --depth 1 https://github.com/rinongal/stylegan-nada.git $stylegan_nada_dir
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/rinongal/stylegan-nada.git', '$stylegan_nada_dir'])
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(os, restyle_dir, stylegan_nada_dir):
+ from argparse import Namespace
+
+ import sys
+
+ import numpy as np
+ from PIL import Image
+ from glob import glob
+
+ import torch
+ import torchvision.transforms as transforms
+
+ sys.path.append(restyle_dir)
+ sys.path.append(stylegan_nada_dir)
+ sys.path.append(os.path.join(stylegan_nada_dir, "ZSSGAN"))
+
+ device = 'cuda'
+
+ # magic command not supported in marimo; please file an issue to add support
+ # %load_ext autoreload
+ # '%autoreload 2' command supported automatically in marimo
+ return Image, Namespace, device, glob, np, torch, transforms
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 2: Choose a model type.
+ Model will be downloaded and converted to a pytorch compatible version.
+
+ Re-runs of the cell with the same model will re-use the previously downloaded version.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Defining the Configs for Selecting the Model
+ """)
+ return
+
+
+@app.cell
+def _():
+ project = "stylegan-nada" #@param {"type": "string"}
+ source_model_type = 'ffhq' #@param['ffhq', 'cat', 'dog', 'church', 'horse', 'car']
+
+ artifact_adressed = {
+ "car": "geekyrakshit/stylegan-nada/car:v0",
+ "horse": "geekyrakshit/stylegan-nada/horse:v0",
+ "church": "geekyrakshit/stylegan-nada/church:v0",
+ "dog": "geekyrakshit/stylegan-nada/dog:v0",
+ "cat": "geekyrakshit/stylegan-nada/cat:v0",
+ "ffhq": "geekyrakshit/stylegan-nada/ffhq:v0"
+ }
+
+ model_names = {
+ "ffhq": "ffhq.pt",
+ "cat": "afhqcat.pkl",
+ "dog": "afhqdog.pkl",
+ "church": "stylegan2-church-config-f.pkl",
+ "car": "stylegan2-car-config-f.pkl",
+ "horse": "stylegan2-horse-config-f.pkl"
+ }
+
+ dataset_sizes = {
+ "ffhq": 1024,
+ "cat": 512,
+ "dog": 512,
+ "church": 256,
+ "horse": 256,
+ "car": 512,
+ }
+ return (
+ artifact_adressed,
+ dataset_sizes,
+ model_names,
+ project,
+ source_model_type,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Initializing WandB
+ """)
+ return
+
+
+@app.cell
+def _(project, source_model_type, wandb):
+ wandb.init(project=project, job_type="train")
+ config = wandb.config
+ config.source_model_type = source_model_type
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Fetching Models from [WandB Artifacts](https://docs.wandb.ai/guides/artifacts/artifacts-core-concepts)
+ """)
+ return
+
+
+@app.cell
+def _(artifact_adressed, model_names, source_model_type, wandb):
+ _artifact = wandb.use_artifact(artifact_adressed[source_model_type])
+ pretrained_model_dir = _artifact.download()
+ pt_file_name = model_names[source_model_type].split('.')[0] + '.pt'
+ return pretrained_model_dir, pt_file_name
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 3: Train the model.
+ Describe your source and target class. These describe the direction of change you're trying to apply (e.g. "photo" to "sketch", "dog" to "the joker" or "dog" to "avocado dog").
+
+ Alternatively, upload a directory with a small (~3) set of target style images (there is no need to preprocess them in any way) and set `style_image_dir` to point at them. This will use the images as a target rather than the source/class texts.
+
+ We reccomend leaving the 'improve shape' button unticked at first, as it will lead to an increase in running times and is often not needed.
+ For more drastic changes, turn it on and increase the number of iterations.
+
+ As a rule of thumb:
+ - Style and minor domain changes ('photo' -> 'sketch') require ~200-400 iterations.
+ - Identity changes ('person' -> 'taylor swift') require ~150-200 iterations.
+ - Simple in-domain changes ('face' -> 'smiling face') may require as few as 50.
+ - The `style_image_dir` option often requires ~400-600 iterations.
+ """)
+ return
+
+
+@app.cell
+def _(config, dataset_sizes, np, source_model_type):
+ from tqdm import notebook
+ from ZSSGAN.model.ZSSGAN import ZSSGAN
+ from ZSSGAN.utils.file_utils import save_images, get_dir_img_list
+ from ZSSGAN.utils.training_utils import mixing_noise
+ from IPython.display import display
+ source_class = 'Human'
+ config.source_class = source_class
+ target_class = 'The Joker'
+ config.target_class = target_class
+ style_image_dir = ''
+ config.style_image_dir = style_image_dir
+ seed = 3 #@param {"type": "string"}
+ config.seed = seed
+ target_img_list = get_dir_img_list(style_image_dir) if style_image_dir else None
+ improve_shape = False #@param {"type": "string"}
+ config.improve_shape = improve_shape
+ model_choice = ['ViT-B/32', 'ViT-B/16']
+ model_weights = [1.0, 0.0] #@param {'type': 'string'}
+ if improve_shape or style_image_dir:
+ model_weights[1] = 1.0
+ mixing = 0.9 if improve_shape else 0.0 #@param {"type": "integer"}
+ auto_layers_k = int(2 * (2 * np.log2(dataset_sizes[source_model_type]) - 2) / 3) if improve_shape else 0
+ auto_layer_iters = 1 if improve_shape else 0
+ training_iterations = 251
+ config.training_iterations = training_iterations
+ output_interval = 10 #@param{type:"boolean"}
+ config.output_interval = output_interval
+ save_interval = 10
+ config.save_interval = save_interval #@param {type: "integer"}
+ return (
+ ZSSGAN,
+ auto_layer_iters,
+ auto_layers_k,
+ mixing,
+ mixing_noise,
+ model_choice,
+ model_weights,
+ notebook,
+ output_interval,
+ save_interval,
+ seed,
+ source_class,
+ target_class,
+ target_img_list,
+ training_iterations,
+ )
+
+
+@app.cell
+def _(
+ auto_layer_iters,
+ auto_layers_k,
+ config,
+ dataset_sizes,
+ mixing,
+ model_choice,
+ model_weights,
+ os,
+ output_dir,
+ pretrained_model_dir,
+ pt_file_name,
+ save_interval,
+ source_class,
+ source_model_type,
+ target_class,
+ target_img_list,
+ training_iterations,
+):
+ training_args = {
+ "size": dataset_sizes[source_model_type],
+ "batch": 2,
+ "n_sample": 4,
+ "output_dir": output_dir,
+ "lr": 0.002,
+ "frozen_gen_ckpt": os.path.join(pretrained_model_dir, pt_file_name),
+ "train_gen_ckpt": os.path.join(pretrained_model_dir, pt_file_name),
+ "iter": training_iterations,
+ "source_class": source_class,
+ "target_class": target_class,
+ "lambda_direction": 1.0,
+ "lambda_patch": 0.0,
+ "lambda_global": 0.0,
+ "lambda_texture": 0.0,
+ "lambda_manifold": 0.0,
+ "auto_layer_k": auto_layers_k,
+ "auto_layer_iters": auto_layer_iters,
+ "auto_layer_batch": 8,
+ "output_interval": 50,
+ "clip_models": model_choice,
+ "clip_model_weights": model_weights,
+ "mixing": mixing,
+ "phase": None,
+ "sample_truncation": 0.7,
+ "save_interval": save_interval,
+ "target_img_list": target_img_list,
+ "img2img_batch": 16,
+ "channel_multiplier": 2,
+ "sg3": False,
+ "sgxl": False,
+ }
+ config.training_args = training_args
+ return (training_args,)
+
+
+@app.cell
+def _(
+ Namespace,
+ ZSSGAN,
+ config,
+ glob,
+ np,
+ os,
+ seed,
+ torch,
+ training_args,
+ wandb,
+):
+ args = Namespace(**training_args)
+ resume_training_from_artifact = False
+ config.resume_training_from_artifact = resume_training_from_artifact #@param{type:"boolean"}
+ checkpoint_artifact_address = 'geekyrakshit/stylegan-nada/model-winter-frost-8:v14'
+ config.checkpoint_artifact_address = checkpoint_artifact_address
+ print('Loading base models...') #@param {'type': 'string'}
+ net = ZSSGAN(args)
+ print('Done')
+ g_reg_ratio = 4 / 5
+ g_optim = torch.optim.Adam(net.generator_trainable.parameters(), lr=args.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio))
+ if resume_training_from_artifact and checkpoint_artifact_address is not None:
+ _artifact = wandb.use_artifact(checkpoint_artifact_address)
+ _artifact_dir = _artifact.download()
+ _checkpoint = torch.load(glob(os.path.join(_artifact_dir, '*.pt'))[0])
+ net.generator_trainable.generator.load_state_dict(_checkpoint['g_ema'])
+ g_optim.load_state_dict(_checkpoint['g_optim'])
+ sample_dir = os.path.join(args.output_dir, 'sample')
+ config.sample_dir = sample_dir
+ ckpt_dir = os.path.join(args.output_dir, 'checkpoint')
+ config.ckpt_dir = ckpt_dir
+ os.makedirs(sample_dir, exist_ok=True)
+ os.makedirs(ckpt_dir, exist_ok=True)
+ torch.manual_seed(seed)
+ # Set up output directories.
+ np.random.seed(seed)
+ return args, ckpt_dir, g_optim, net
+
+
+@app.cell
+def _(
+ args,
+ ckpt_dir,
+ device,
+ g_optim,
+ mixing_noise,
+ net,
+ notebook,
+ output_interval,
+ source_model_type,
+ torch,
+ wandb,
+):
+ fixed_z = torch.randn(args.n_sample, 512, device=device)
+ for i in notebook.tqdm(range(args.iter)):
+ net.train()
+ _sample_z = mixing_noise(args.batch, 512, args.mixing, device)
+ [_sampled_src, _sampled_dst], clip_loss = net(_sample_z)
+ wandb.log({'CLIP-Loss': clip_loss.item()}, step=i)
+ net.zero_grad()
+ clip_loss.backward()
+ g_optim.step()
+ if i % output_interval == 0:
+ net.eval()
+ with torch.no_grad():
+ [_sampled_src, _sampled_dst], _loss = net([fixed_z], truncation=args.sample_truncation)
+ if source_model_type == 'car':
+ _sampled_dst = _sampled_dst[:, :, 64:448, :]
+ _sampled_dst = torch.permute(_sampled_dst, (0, 2, 3, 1)).cpu()
+ _sampled_dst = [wandb.Image(dst.numpy()) for dst in _sampled_dst]
+ wandb.log({'Samples': _sampled_dst}, step=i)
+ if args.save_interval > 0 and i > 0 and (i % args.save_interval == 0):
+ model_file = f'{ckpt_dir}/{str(i).zfill(6)}.pt'
+ torch.save({'g_ema': net.generator_trainable.generator.state_dict(), 'g_optim': g_optim.state_dict()}, model_file)
+ _artifact = wandb.Artifact(f'model-{wandb.run.name}', type='model')
+ _artifact.add_file(model_file)
+ wandb.log_artifact(_artifact, aliases=['latest', f'step_{i}'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 4: Generate samples with the new model
+ """)
+ return
+
+
+@app.cell
+def _(ZSSGAN, args, config, device, glob, os, source_model_type, torch, wandb):
+ truncation = 0.7
+ config.truncation = truncation
+ samples = 9
+ config.samples = samples
+ _artifact = wandb.use_artifact(f'model-{wandb.run.name}:latest')
+ _artifact_dir = _artifact.download()
+ _checkpoint = torch.load(glob(os.path.join(_artifact_dir, '*.pt'))[0])
+ print('Loading models from checkpoint artifact...')
+ net_1 = ZSSGAN(args)
+ net_1.generator_trainable.generator.load_state_dict(_checkpoint['g_ema'])
+ print('Done')
+ with torch.no_grad():
+ net_1.eval()
+ _sample_z = torch.randn(samples, 512, device=device)
+ [_sampled_src, _sampled_dst], _loss = net_1([_sample_z], truncation=truncation)
+ if source_model_type == 'car':
+ _sampled_dst = _sampled_dst[:, :, 64:448, :]
+ _sampled_dst = torch.permute(_sampled_dst, (0, 2, 3, 1)).cpu()
+ _sampled_dst = [wandb.Image(dst.numpy()) for dst in _sampled_dst]
+ wandb.log({'Generated Samples': _sampled_dst})
+ return (net_1,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Editing a real image with Re-Style inversion (currently only FFHQ inversion is supported):
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 1: Fetch ReStyle Models from [WandB Artifacts](https://docs.wandb.ai/guides/artifacts/artifacts-core-concepts)
+
+ This may take a few minutes
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ from restyle.utils.common import tensor2im
+ from restyle.models.psp import pSp
+ from restyle.models.e4e import e4e
+ _artifact = wandb.use_artifact('geekyrakshit/stylegan-nada/restyle:v0')
+ pretrained_model_dir_1 = _artifact.download()
+ return e4e, pSp, pretrained_model_dir_1
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 2: Choose a re-style model
+
+ We reccomend choosing the e4e model as it performs better under domain translations. Choose pSp for better reconstructions on minor domain changes (typically those that require less than 150 training steps).
+ """)
+ return
+
+
+@app.cell
+def _(Namespace, e4e, os, pSp, pretrained_model_dir_1, torch, transforms):
+ encoder_type = 'e4e' #@param['psp', 'e4e']
+ restyle_experiment_args = {'model_path': os.path.join(pretrained_model_dir_1, f'restyle_{encoder_type}_ffhq_encode.pt'), 'transform': transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])}
+ model_path = restyle_experiment_args['model_path']
+ ckpt = torch.load(model_path, map_location='cpu')
+ opts = ckpt['opts']
+ opts['checkpoint_path'] = model_path
+ opts = Namespace(**opts)
+ restyle_net = (pSp if encoder_type == 'psp' else e4e)(opts)
+ restyle_net.eval()
+ restyle_net.cuda()
+ print('Model successfully loaded!')
+ return opts, restyle_experiment_args, restyle_net
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 3: Align and invert an image
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ def run_alignment(image_path):
+ import dlib
+ from scripts.align_faces_parallel import align_face
+ if not os.path.exists("shape_predictor_68_face_landmarks.dat"):
+ print('Downloading files for aligning face image...')
+ os.system('wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2')
+ os.system('bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2')
+ print('Done.')
+ predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+ aligned_image = align_face(filepath=image_path, predictor=predictor)
+ print("Aligned image has shape: {}".format(aligned_image.size))
+ return aligned_image
+
+ return (run_alignment,)
+
+
+@app.cell
+def _(
+ Image,
+ opts,
+ os,
+ restyle_experiment_args,
+ restyle_net,
+ run_alignment,
+ subprocess,
+ torch,
+):
+ image_url = "https://engineering.nyu.edu/sites/default/files/styles/square_large_default_2x/public/2018-06/yann-lecun.jpg" #@param {"type": "string"}
+ file_name = "yann-lecun.jpg" #@param {"type": "string"}
+
+ if not os.path.isfile(file_name):
+ #! wget {image_url}
+ subprocess.call(['wget', str(image_url)])
+
+ image_path = os.path.join("/content", file_name)
+ original_image = Image.open(image_path).convert("RGB")
+
+ input_image = run_alignment(image_path)
+
+ img_transforms = restyle_experiment_args['transform']
+ transformed_image = img_transforms(input_image)
+
+ def get_avg_image(net):
+ avg_image = net(
+ net.latent_avg.unsqueeze(0),
+ input_code=True,
+ randomize_noise=False,
+ return_latents=False,
+ average_code=True
+ )[0]
+ avg_image = avg_image.to('cuda').float().detach()
+ return avg_image
+
+ opts.n_iters_per_batch = 5
+ opts.resize_outputs = False # generate outputs at full resolution
+
+ from restyle.utils.inference_utils import run_on_batch
+
+ with torch.no_grad():
+ avg_image = get_avg_image(restyle_net)
+ result_batch, result_latents = run_on_batch(
+ transformed_image.unsqueeze(0).cuda(), restyle_net, opts, avg_image
+ )
+ return (result_latents,)
+
+
+@app.cell
+def _(net_1, result_latents, source_class, target_class, torch, wandb):
+ inverted_latent = torch.Tensor(result_latents[0][4]).cuda().unsqueeze(0).unsqueeze(1)
+ with torch.no_grad():
+ net_1.eval()
+ [_sampled_src, _sampled_dst] = net_1(inverted_latent, input_is_latent=True)[0]
+ _sampled_src = torch.permute(_sampled_src, (0, 2, 3, 1)).cpu().numpy()[0]
+ _sampled_dst = torch.permute(_sampled_dst, (0, 2, 3, 1)).cpu().numpy()[0]
+ table = wandb.Table(columns=['Source-Class-Text', 'Source-Image', 'Target-Class-Text', 'Translated-Image'], data=[[source_class, wandb.Image(_sampled_src), target_class, wandb.Image(_sampled_dst)]])
+ wandb.log({'Restyle': table})
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/super-gradients-yolo-nas-data-analysis/super_gradients_yolo_nas_data_analysis.py b/marimo/convert/super-gradients-yolo-nas-data-analysis/super_gradients_yolo_nas_data_analysis.py
new file mode 100644
index 00000000..e62dc0e5
--- /dev/null
+++ b/marimo/convert/super-gradients-yolo-nas-data-analysis/super_gradients_yolo_nas_data_analysis.py
@@ -0,0 +1,314 @@
+# /// script
+# dependencies = ["pycairo", "roboflow", "super-gradients", "sweeps", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! sudo apt install libcairo2-dev pkg-config python3-dev -qq
+ subprocess.call(['sudo', 'apt', 'install', 'libcairo2-dev', 'pkg-config', 'python3-dev', '-qq'])
+ # packages added via marimo's package management: roboflow pycairo wandb sweeps !pip install roboflow pycairo wandb sweeps -qqq
+ # packages added via marimo's package management: super_gradients !pip install super_gradients
+ return
+
+
+@app.cell
+def _():
+ import os
+ import glob
+ import torch
+ import wandb
+ import warnings
+ import numpy as np
+ import pandas as pd
+ import matplotlib.pyplot as plt
+
+ from matplotlib import patches
+ from google.colab import userdata
+ from torchvision.io import read_image
+ from torch.utils.data import DataLoader
+ from super_gradients.training import models, dataloaders
+ from super_gradients.training.dataloaders.dataloaders import (
+ coco_detection_yolo_format_train, coco_detection_yolo_format_val
+ )
+
+ warnings.filterwarnings("ignore")
+
+ os.environ["WANDB_API_KEY"] = userdata.get('wandb')
+ os.environ["ROBOFLOW_API_KEY"] = userdata.get('roboflow')
+ return (
+ coco_detection_yolo_format_train,
+ coco_detection_yolo_format_val,
+ np,
+ os,
+ wandb,
+ )
+
+
+@app.cell
+def _(os):
+ from roboflow import Roboflow
+ rf = Roboflow(api_key=os.getenv("ROBOFLOW_API_KEY"))
+ project = rf.workspace("easyhyeon").project("trash-sea")
+ dataset = project.version(10).download("yolov5")
+ return
+
+
+@app.cell
+def _():
+ DATASET_PATH = "/content/trash-sea-10"
+ WANDB_PROJECT_NAME = "fconn-yolo-nas"
+ ENTITY = "ml-colabs"
+ return DATASET_PATH, ENTITY, WANDB_PROJECT_NAME
+
+
+@app.cell
+def _(DATASET_PATH):
+ dataset_params = {
+ 'data_dir':DATASET_PATH,
+ 'train_images_dir':'train/images',
+ 'train_labels_dir':'train/labels',
+ 'val_images_dir':'valid/images',
+ 'val_labels_dir':'valid/labels',
+ 'test_images_dir':'test/images',
+ 'test_labels_dir':'test/labels',
+ 'classes': ["Buoy", "Can", "Paper", "Plastic Bag", "Plastic Bottle"]
+ }
+ return (dataset_params,)
+
+
+@app.cell
+def _(
+ coco_detection_yolo_format_train,
+ coco_detection_yolo_format_val,
+ dataset_params,
+):
+ from IPython.display import clear_output
+
+ train_data = coco_detection_yolo_format_train(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['train_images_dir'],
+ 'labels_dir': dataset_params['train_labels_dir'],
+ 'classes': dataset_params['classes'],
+ },
+ dataloader_params={
+ 'batch_size':16,
+ 'num_workers':2
+ }
+ )
+
+ val_data = coco_detection_yolo_format_val(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['val_images_dir'],
+ 'labels_dir': dataset_params['val_labels_dir'],
+ 'classes': dataset_params['classes'],
+ },
+ dataloader_params={
+ 'batch_size':16,
+ 'num_workers':2
+ }
+ )
+
+ test_data = coco_detection_yolo_format_val(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['test_images_dir'],
+ 'labels_dir': dataset_params['test_labels_dir'],
+ 'classes': dataset_params['classes'],
+ },
+ dataloader_params={
+ 'batch_size':16,
+ 'num_workers':2
+ }
+ )
+
+ train_data.dataset.transforms = train_data.dataset.transforms[5:]
+ return (train_data,)
+
+
+@app.cell
+def _():
+ colors = {
+ 0: 'red',
+ 1: 'green',
+ 2: 'blue',
+ 3: 'yellow',
+ 4: 'black'
+ }
+ classes = {
+ 0:"Buoy",
+ 1:"Can",
+ 2:"Paper",
+ 3:"Plastic Bag",
+ 4:"Plastic Bottle"
+ }
+ return (classes,)
+
+
+@app.cell
+def _(ENTITY, WANDB_PROJECT_NAME, classes, train_data, wandb):
+ def _process_bounding_boxes_list(annots):
+ result = []
+ class_count = {i: 0 for i in range(0, 5)}
+ for annot_idx, annotation in enumerate(annots):
+ class_count[int(annotation[1])] += 1
+ result.append({'position': {'middle': [float(annotation[2]), float(annotation[3])], 'width': float(annotation[4]), 'height': float(annotation[5])}, 'domain': 'pixel', 'class_id': int(annotation[1]), 'box_caption': classes[int(annotation[1])]})
+ return (result, class_count)
+
+ def populate_wandb_image_samples(train_data):
+ wandb.init(project=WANDB_PROJECT_NAME, entity=ENTITY, id='add-image-samples', job_type='add-tables', resume='allow')
+ class_set = wandb.Classes([{'name': 'Buoy', 'id': 0}, {'name': 'Can', 'id': 1}, {'name': 'Paper', 'id': 2}, {'name': 'Plastic Bag', 'id': 3}, {'name': 'Plastic Bottle', 'id': 4}])
+ table = wandb.Table(columns=['Annotated-Image', 'Number-of-objects', 'Number-Buoy', 'Number-Can', 'Number-Paper', 'Number-Plastic-Bag', 'Number-Plastic-Bottle'])
+ img_count = 0
+ for batch_idx, batch_sample in enumerate(train_data):
+ batch_images = batch_sample[0]
+ batch_annotations = batch_sample[1]
+ annots_dict = {i: [] for i in range(0, batch_images.shape[0])}
+ for annot in batch_annotations:
+ annots_dict[int(annot[0])].append(annot)
+ for idx, image in enumerate(batch_images):
+ bbox, class_count = _process_bounding_boxes_list(annots_dict[idx])
+ image = image.flip(0)
+ img = wandb.Image(image, boxes={'ground_truth': {'box_data': bbox, 'class_labels': classes}}, classes=class_set)
+ table.add_data(img, len(bbox), class_count[0], class_count[1], class_count[2], class_count[3], class_count[4])
+ img_count += 1
+ print(f'{img_count}/{len(train_data) * 16} completed')
+ wandb.log({'ground_truth_dataset': table})
+ wandb.finish()
+ populate_wandb_image_samples(train_data)
+ return
+
+
+@app.cell
+def _(ENTITY, WANDB_PROJECT_NAME, classes, train_data, wandb):
+ def _process_bounding_boxes_list(annots):
+ result = []
+ class_count = {i: 0 for i in range(0, 5)}
+ for annot_idx, annotation in enumerate(annots):
+ class_count[int(annotation[1])] += 1
+ result.append({'position': {'middle': [float(annotation[2]), float(annotation[3])], 'width': float(annotation[4]), 'height': float(annotation[5])}, 'domain': 'pixel', 'class_id': int(annotation[1]), 'box_caption': classes[int(annotation[1])]})
+ return (result, class_count)
+
+ def populate_wandb_bbox(train_data):
+ wandb.init(project=WANDB_PROJECT_NAME, entity=ENTITY, id='add-bbox-data', job_type='add-tables', resume='allow')
+ class_set = wandb.Classes([{'name': 'Buoy', 'id': 0}, {'name': 'Can', 'id': 1}, {'name': 'Paper', 'id': 2}, {'name': 'Plastic Bag', 'id': 3}, {'name': 'Plastic Bottle', 'id': 4}])
+ table = wandb.Table(columns=['Image-Id', 'BBox-Height', 'BBox-Width', 'Class-Id'])
+ img_count = 0
+ for batch_idx, batch_sample in enumerate(train_data):
+ batch_images = batch_sample[0]
+ batch_annotations = batch_sample[1]
+ annots_dict = {i: [] for i in range(0, batch_images.shape[0])}
+ for annot in batch_annotations:
+ annots_dict[int(annot[0])].append(annot)
+ for idx, image in enumerate(batch_images):
+ result, class_count = _process_bounding_boxes_list(annots_dict[idx])
+ for bbox in result:
+ height = bbox['position']['height']
+ width = bbox['position']['width']
+ class_id = bbox['class_id']
+ table.add_data(img_count, height, width, classes[class_id])
+ img_count += 1
+ print(f'{img_count}/{len(train_data) * 16} completed')
+ wandb.log({'bounding_box_information': table})
+ wandb.finish()
+ populate_wandb_bbox(train_data)
+ return
+
+
+@app.cell
+def _(ENTITY, WANDB_PROJECT_NAME, classes, np, train_data, wandb):
+ def populate_wandb_spatial_heatmaps(train_data):
+ wandb.init(
+ project=WANDB_PROJECT_NAME,
+ entity=ENTITY,
+ id='add-heatmap',
+ job_type="add-tables",
+ resume='allow'
+ )
+
+ class_set = wandb.Classes(
+ [
+ {"name": "Buoy", "id": 0},
+ {"name": "Can", "id": 1},
+ {"name": "Paper", "id": 2},
+ {"name": "Plastic Bag", "id": 3},
+ {"name": "Plastic Bottle", "id": 4},
+ ]
+ )
+ heatmaps = [np.zeros((224, 224, 1), dtype=np.float32) for _ in classes]
+ annotation_counts = {i:0 for i in range(len(classes))}
+
+ table = wandb.Table(columns=["Class-Id", "Class-Name", "Spatial-Heatmap",
+ "Num-Total-Objects"])
+
+ for batch_idx, batch_sample in enumerate(train_data):
+ batch_images = batch_sample[0]
+ batch_annotations = batch_sample[1]
+
+ annots_dict = {i:[] for i in range(0, batch_images.shape[0])}
+
+ for annot in batch_annotations:
+ class_idx = int(annot[1])
+
+ midpoint_x = int(annot[2])
+ midpoint_y = int(annot[3])
+ width = int(annot[4])
+ height = int(annot[5])
+
+ x_min = midpoint_x - (width//2)
+ x_max = midpoint_x + (width//2)
+
+ y_min = midpoint_y - (height//2)
+ y_max = midpoint_y + (height//2)
+
+ heatmaps[class_idx][y_min:y_max, x_min:x_max] += 1
+
+ annotation_counts[class_idx] += 1
+
+ print(f"{batch_idx+1}/{len(train_data)} batches completed")
+
+ for class_idx in range(len(classes)):
+ heatmap = wandb.Image(
+ heatmaps[class_idx],
+ caption=classes[class_idx]
+ )
+ table.add_data(class_idx, classes[class_idx], heatmap, annotation_counts[class_idx])
+
+ wandb.log({"spatial_heatmap_information": table})
+ wandb.finish()
+
+ populate_wandb_spatial_heatmaps(train_data)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/super-gradients-yolo-nas-sweep-run/super_gradients_yolo_nas_sweep_run.py b/marimo/convert/super-gradients-yolo-nas-sweep-run/super_gradients_yolo_nas_sweep_run.py
new file mode 100644
index 00000000..57b1af82
--- /dev/null
+++ b/marimo/convert/super-gradients-yolo-nas-sweep-run/super_gradients_yolo_nas_sweep_run.py
@@ -0,0 +1,342 @@
+# /// script
+# dependencies = ["pycairo", "roboflow", "super-gradients", "sweeps", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installation and Imports
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! sudo apt install libcairo2-dev pkg-config python3-dev -qq
+ subprocess.call(['sudo', 'apt', 'install', 'libcairo2-dev', 'pkg-config', 'python3-dev', '-qq'])
+ # packages added via marimo's package management: roboflow pycairo wandb sweeps !pip install roboflow pycairo wandb sweeps -qqq
+ # packages added via marimo's package management: super_gradients !pip install super_gradients
+ return
+
+
+@app.cell
+def _():
+ import os
+ import glob
+ import torch
+ import wandb
+ import warnings
+ import pandas as pd
+
+ from google.colab import userdata
+ from torchvision.io import read_image
+ from torch.utils.data import DataLoader
+ from IPython.display import clear_output
+
+ from super_gradients.training import models, Trainer, dataloaders
+ from super_gradients.training.losses import PPYoloELoss
+ from super_gradients.training.metrics import DetectionMetrics_050
+ from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
+ from super_gradients.training.dataloaders.dataloaders import coco_detection_yolo_format_train, coco_detection_yolo_format_val
+
+ os.environ["WANDB_API_KEY"] = userdata.get('wandb')
+ os.environ["ROBOFLOW_API_KEY"] = userdata.get('roboflow')
+ return (
+ DetectionMetrics_050,
+ PPYoloELoss,
+ PPYoloEPostPredictionCallback,
+ Trainer,
+ coco_detection_yolo_format_train,
+ coco_detection_yolo_format_val,
+ models,
+ os,
+ torch,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Definitions
+ """)
+ return
+
+
+@app.cell
+def _(torch):
+ seed = 42
+ torch.manual_seed(seed)
+
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download and Register dataset
+ """)
+ return
+
+
+@app.cell
+def _(os):
+ from roboflow import Roboflow
+ rf = Roboflow(api_key=os.getenv("ROBOFLOW_API_KEY"))
+ project = rf.workspace("easyhyeon").project("trash-sea")
+ dataset = project.version(10).download("yolov5")
+ return
+
+
+@app.cell
+def _():
+ ENTITY = "ml-colabs"
+ SWEEP_NUM_RUNS = 100
+ WANDB_PROJECT_NAME = "fconn-yolo-nas"
+ DATASET_PATH = "/content/trash-sea-10"
+ return DATASET_PATH, ENTITY, SWEEP_NUM_RUNS, WANDB_PROJECT_NAME
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Define Sweep Configuration and functions
+ """)
+ return
+
+
+@app.cell
+def _(WANDB_EXP_NAME):
+ sweep_configuration = {
+ "name": WANDB_EXP_NAME,
+ "metric": {"name": "Valid_mAP@0.50", "goal": "maximize"},
+ "method": "bayes",
+ "parameters": {
+ "batch_size": {"values": [16, 24, 32]},
+ "optimizer": {"values": ["Adam", "SGD", "RMSProp", "AdamW"]},
+ "ema_decay": {"min":0.5, "max":0.9},
+ "ema_decay_type": {"values": ["constant", "threshold"]},
+ "cosine_lr_ratio": {"min": 0.01, "max": 0.4},
+ "iou_loss_weight": {"min": 0.25, "max": 2.0},
+ "dfl_loss_weight": {"min": 0.25, "max": 2.0},
+ "classification_loss_weight": {"min": 0.25, "max": 2.0},
+ "model_flavor": {"values": ["yolo_nas_s", "yolo_nas_m", "yolo_nas_l"]},
+ "weight_decay": {"min": 0.0001, "max": 0.01},
+ },
+ }
+ return (sweep_configuration,)
+
+
+@app.cell
+def _(
+ DATASET_PATH,
+ DetectionMetrics_050,
+ ENTITY,
+ PPYoloELoss,
+ PPYoloEPostPredictionCallback,
+ Trainer,
+ WANDB_EXP_NAME,
+ WANDB_PROJECT_NAME,
+ coco_detection_yolo_format_train,
+ coco_detection_yolo_format_val,
+ models,
+ wandb,
+):
+ def main_call():
+
+ CHECKPOINT_DIR = 'checkpoints'
+
+ wandb.init(
+ project=WANDB_PROJECT_NAME,
+ entity=ENTITY,
+ resume="allow",
+ save_code=True,
+ id=WANDB_EXP_NAME
+ )
+
+ config = wandb.config
+
+ dataset_params = {
+ 'data_dir':DATASET_PATH,
+ 'train_images_dir':'train/images',
+ 'train_labels_dir':'train/labels',
+ 'val_images_dir':'valid/images',
+ 'val_labels_dir':'valid/labels',
+ 'test_images_dir':'test/images',
+ 'test_labels_dir':'test/labels',
+ 'classes': ["Buoy", "Can", "Paper", "Plastic Bag", "Plastic Bottle"]
+ }
+
+ train_data = coco_detection_yolo_format_train(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['train_images_dir'],
+ 'labels_dir': dataset_params['train_labels_dir'],
+ 'classes': dataset_params['classes'],
+ },
+ dataloader_params={
+ 'batch_size':config["batch_size"],
+ 'num_workers':4
+ }
+ )
+
+ val_data = coco_detection_yolo_format_val(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['val_images_dir'],
+ 'labels_dir': dataset_params['val_labels_dir'],
+ 'classes': dataset_params['classes'],
+ },
+ dataloader_params={
+ 'batch_size':config["batch_size"],
+ 'num_workers':4
+ }
+ )
+
+ test_data = coco_detection_yolo_format_val(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['test_images_dir'],
+ 'labels_dir': dataset_params['test_labels_dir'],
+ 'classes': dataset_params['classes'],
+ },
+ dataloader_params={
+ 'batch_size':config["batch_size"],
+ 'num_workers':4
+ }
+ )
+
+ train_data.dataset.transforms = train_data.dataset.transforms[1:]
+
+ model = models.get(
+ config["model_flavor"],
+ num_classes=len(dataset_params['classes']),
+ pretrained_weights="coco"
+ )
+
+ train_params = {
+ 'silent_mode': False,
+ "average_best_models":True,
+ "warmup_mode": "linear_epoch_step",
+ "warmup_initial_lr": 1e-6,
+ "lr_warmup_epochs": 3,
+ "initial_lr": 1e-3,
+ "lr_mode": "cosine",
+ "cosine_final_lr_ratio": config["cosine_lr_ratio"],
+ "optimizer": config["optimizer"],
+ "optimizer_params": {
+ "weight_decay": config["weight_decay"]
+ },
+ "zero_weight_decay_on_bias_and_bn": True,
+ "ema": True,
+ "ema_params": {
+ "decay": config["ema_decay"],
+ "decay_type": config["ema_decay_type"]
+ },
+ "max_epochs": 5,
+ "mixed_precision": False,
+ "loss": PPYoloELoss(
+ use_static_assigner=False,
+ num_classes=len(dataset_params['classes']),
+ reg_max=16,
+ iou_loss_weight=config["iou_loss_weight"],
+ dfl_loss_weight=config["dfl_loss_weight"],
+ classification_loss_weight=config["classification_loss_weight"]
+ ),
+ "valid_metrics_list": [
+ DetectionMetrics_050(
+ score_thres=0.1,
+ top_k_predictions=300,
+ num_cls=len(dataset_params['classes']),
+ normalize_targets=True,
+ post_prediction_callback=PPYoloEPostPredictionCallback(
+ score_threshold=0.01,
+ nms_top_k=1000,
+ max_predictions=300,
+ nms_threshold=0.7
+ )
+ )
+ ],
+ "metric_to_watch": 'mAP@0.50',
+ "sg_logger": "wandb_sg_logger",
+ "sg_logger_params": {
+ "project_name": WANDB_PROJECT_NAME,
+ "save_checkpoints_remote": True,
+ "save_tensorboard_remote": True,
+ "save_logs_remote": True,
+ "entity": ENTITY
+ }
+ }
+
+ trainer = Trainer(
+ experiment_name=WANDB_EXP_NAME,
+ ckpt_root_dir=CHECKPOINT_DIR
+ )
+
+ trainer.train(
+ model=model,
+ training_params=train_params,
+ train_loader=train_data,
+ valid_loader=val_data
+ )
+
+ wandb.finish()
+
+ return (main_call,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Execute Sweep
+ """)
+ return
+
+
+@app.cell
+def _(SWEEP_NUM_RUNS, main_call, sweep_configuration, wandb):
+ sweep_id = wandb.sweep(
+ sweep=sweep_configuration,
+ project="yolo-nas-sweep"
+ )
+
+ wandb.agent(sweep_id, function=main_call, count=SWEEP_NUM_RUNS)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/super-gradients-yolo-nas/super_gradients_yolo_nas.py b/marimo/convert/super-gradients-yolo-nas/super_gradients_yolo_nas.py
new file mode 100644
index 00000000..e3ef224f
--- /dev/null
+++ b/marimo/convert/super-gradients-yolo-nas/super_gradients_yolo_nas.py
@@ -0,0 +1,370 @@
+# /// script
+# dependencies = ["super-gradients", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥 Fine-tuning YOLO-NAS using Super-Gradients and Weights & Biases 🐝
+
+ This notebook demonstrates how to fine-tune YOLO-NAS on a custom dataset using the [Super-Graidents](https://github.com/Deci-AI/super-gradients) library and performing experiment-tracking, logging and versioning model checkpoints, and viusalizing your detection datasets and prediction results during training.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Installing Dependencies
+
+ We install [Super-Graidents](https://github.com/Deci-AI/super-gradients) and [Weights & Biases](wandb.ai/site).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: super_gradients wandb !pip install -qq super_gradients wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Experiment Setup
+
+ First, we initialize a W&B [run](https://docs.wandb.ai/guides/runs) using `wandb.init`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+
+ wandb.init(project="yolo-nas-integration-2")
+ return (wandb,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we will initialize our trainer which will be in charge of the whole workflow, including the likes of training, evaluation, saving checkpoints, visualization of results, etc.
+ """)
+ return
+
+
+@app.cell
+def _():
+ from super_gradients.training import Trainer
+
+
+ trainer = Trainer(
+ experiment_name='transfer_learning_object_detection_yolo_nas', ckpt_root_dir="./checkpoints/"
+ )
+ return (trainer,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setting up the Dataset
+
+ Next, we fetch a subset of the [BDD100K dataset](https://www.vis.xyz/bdd100k/) hosted on Weights & Biases as a [dataset artifact](https://docs.wandb.ai/guides/artifacts). Hosting the dataset as an artifact not only enables us to maintain different versions of our datasets, but also enables us to keep track of the runs that produced it or the runs that are using it via the [lineage panel](https://docs.wandb.ai/guides/app/pages/project-page#lineage-panel).
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ artifact = wandb.use_artifact('geekyrakshit/yolo-nas-integration/bdd100k-subset-yolo:v2', type='dataset')
+ artifact_dir = artifact.download()
+ return (artifact_dir,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we set up the dataset parameters, and create the dataloaders for training, validation and testing using the `coco_detection_yolo_format_train` and the `coco_detection_yolo_format_val` functions from [Super-Graidents](https://github.com/Deci-AI/super-gradients), that would automatically create the dataset loading, pre-processing and augmentation pipelines.
+ """)
+ return
+
+
+@app.cell
+def _(artifact_dir):
+ import os
+ from super_gradients.training.dataloaders.dataloaders import (
+ coco_detection_yolo_format_train, coco_detection_yolo_format_val
+ )
+
+
+ with open(os.path.join(artifact_dir, "labels.txt"), "r") as f:
+ labels = f.read().split("\n")
+
+ dataset_params = {
+ 'data_dir': artifact_dir,
+ 'train_images_dir':'train/images',
+ 'train_labels_dir':'train/labels',
+ 'val_images_dir':'val/images',
+ 'val_labels_dir':'val/labels',
+ 'test_images_dir':'test/images',
+ 'test_labels_dir':'test/labels',
+ 'classes': labels
+ }
+
+ train_data = coco_detection_yolo_format_train(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['train_images_dir'],
+ 'labels_dir': dataset_params['train_labels_dir'],
+ 'classes': dataset_params['classes']
+ },
+ dataloader_params={
+ 'batch_size':16,
+ 'num_workers':2
+ }
+ )
+
+ val_data = coco_detection_yolo_format_val(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['val_images_dir'],
+ 'labels_dir': dataset_params['val_labels_dir'],
+ 'classes': dataset_params['classes']
+ },
+ dataloader_params={
+ 'batch_size':16,
+ 'num_workers':2
+ }
+ )
+
+ test_data = coco_detection_yolo_format_val(
+ dataset_params={
+ 'data_dir': dataset_params['data_dir'],
+ 'images_dir': dataset_params['test_images_dir'],
+ 'labels_dir': dataset_params['test_labels_dir'],
+ 'classes': dataset_params['classes']
+ },
+ dataloader_params={
+ 'batch_size':16,
+ 'num_workers':2
+ }
+ )
+ return dataset_params, labels, test_data, train_data, val_data
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We can visualize our datasets using the `plot_detection_dataset_on_wandb` function on our Weights & Biases dashboard. The datasets are logged into a [Weights & Biases Table](https://docs.wandb.ai/guides/tables), in which the images can be visualized overlayed with an [interactive overlays for computer vision tasks](https://docs.wandb.ai/guides/track/log/media#image-overlays).
+ """)
+ return
+
+
+@app.cell
+def _(test_data, train_data, val_data):
+ from super_gradients.common.plugins.wandb import plot_detection_dataset_on_wandb
+
+
+ plot_detection_dataset_on_wandb(train_data.dataset, max_examples=20, dataset_name="Train-Dataset")
+ plot_detection_dataset_on_wandb(val_data.dataset, max_examples=20, dataset_name="Validation-Dataset")
+ plot_detection_dataset_on_wandb(test_data.dataset, max_examples=20, dataset_name="Test-Dataset")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's how the datasets look on Weights & Biases 👇
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Performing Transfer Learning
+
+ For performing transfer learning, we first define the `YOLO_NAS_S` model as the pre-trained backbone from [Super-Graidents](https://github.com/Deci-AI/super-gradients). You can check the [Super-Gradients Model Zoo](https://docs.deci.ai/super-gradients/documentation/source/model_zoo.html#computer-vision-models-pretrained-checkpoints) for all the available pre-trained models for object detection.
+ """)
+ return
+
+
+@app.cell
+def _(dataset_params):
+ from super_gradients.training import models
+ from super_gradients.common.object_names import Models
+
+
+ net = models.get(Models.YOLO_NAS_S, pretrained_weights="coco", num_classes=len(dataset_params["classes"]))
+ return (net,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Experiment Tracking with Weights & Biases
+
+ Now, we define the training parameters for the experiment. In order to perform experiment tracking with Weights & Biases, in the training parameters, we set the value of `sg_logger` to `wandb_sg_logger`. You can also set the value of additional parameters for the `wandb_sg_logger` by defining them under `sg_logger_params` like
+
+ ```python
+ train_params = {
+ ...
+ "sg_logger": "wandb_sg_logger",
+ "sg_logger_params": {
+ "save_checkpoints_remote": True,
+ "save_tensorboard_remote": True,
+ "save_logs_remote": True,
+ "save_checkpoint_as_artifact": True,
+ },
+ }
+ ```
+
+ Here's a complete list of params for the `wandb_sg_logger`:
+
+ |Parameter|Description|Default Value|
+ |---|---|---|
+ |experiment_name|Name used for logging and loading purposes| |
+ |storage_location|If set to 's3' (i.e. s3://my-bucket) saves the Checkpoints in AWS S3 otherwise saves the Checkpoints Locally| |
+ |resumed|If true, then old tensorboard files will **NOT** be deleted when tb_files_user_prompt=True| |
+ |training_params|training_params for the experiment| |
+ |checkpoints_dir_path|Local root directory path where all experiment logging directories will reside.| |
+ |tb_files_user_prompt|Asks user for Tensorboard deletion prompt.| |
+ |launch_tensorboard|Whether to launch a TensorBoard process.|False|
+ |tensorboard_port|Specific port number for the tensorboard to use when launched (when set to None, some free port number will be used)|False|
+ |save_checkpoints_remote|Saves checkpoints in s3.|True|
+ |save_tensorboard_remote|Saves tensorboard in s3.|True|
+ |save_logs_remote|Saves log files in s3.|True|
+ |save_code|Save current code to wandb|False|
+ |save_checkpoint_as_artifact|Save model checkpoint using Weights & Biases Artifact. Note that setting this option to True would save model checkpoints every epoch as a versioned artifact, which will result in use of increased storage usage on Weights & Biases.|False|
+
+ The Weights & Biases integration for Super Gradients also come with the callback `WandBDetectionValidationPredictionLoggerCallback` for logging object detection predictions to Weights & Biases during training. In order to log object detection predictions to Weights & Biases, we can include this callback in the training parameters:
+
+ ```python
+ train_params = {
+ ...
+ "sg_logger": "wandb_sg_logger",
+ "sg_logger_params": {
+ "save_checkpoints_remote": True,
+ "save_tensorboard_remote": True,
+ "save_logs_remote": True,
+ "save_checkpoint_as_artifact": True,
+ },
+ "phase_callbacks": [
+ WandBDetectionValidationPredictionLoggerCallback(class_names=labels),
+ ]
+ }
+ ```
+ """)
+ return
+
+
+@app.cell
+def _(dataset_params, labels):
+ from super_gradients.training.losses import PPYoloELoss
+ from super_gradients.training.metrics import DetectionMetrics_050
+ from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
+ from super_gradients.common.plugins.wandb.validation_logger import WandBDetectionValidationPredictionLoggerCallback
+
+
+ train_params = {
+ 'silent_mode': True,
+ "average_best_models":True,
+ "warmup_mode": "linear_epoch_step",
+ "warmup_initial_lr": 1e-6,
+ "lr_warmup_epochs": 3,
+ "initial_lr": 5e-4,
+ "lr_mode": "cosine",
+ "cosine_final_lr_ratio": 0.1,
+ "optimizer": "Adam",
+ "optimizer_params": {"weight_decay": 0.0001},
+ "zero_weight_decay_on_bias_and_bn": True,
+ "ema": True,
+ "ema_params": {"decay": 0.9, "decay_type": "threshold"},
+ "max_epochs": 10,
+ "mixed_precision": True,
+ "loss": PPYoloELoss(
+ use_static_assigner=False,
+ num_classes=len(dataset_params['classes']),
+ reg_max=16
+ ),
+ "valid_metrics_list": [
+ DetectionMetrics_050(
+ score_thres=0.1,
+ top_k_predictions=300,
+ # NOTE: num_classes needs to be defined here
+ num_cls=len(dataset_params['classes']),
+ normalize_targets=True,
+ post_prediction_callback=PPYoloEPostPredictionCallback(
+ score_threshold=0.01,
+ nms_top_k=1000,
+ max_predictions=300,
+ nms_threshold=0.7
+ )
+ )
+ ],
+ "metric_to_watch": 'mAP@0.50',
+ "sg_logger": "wandb_sg_logger",
+ "sg_logger_params": {
+ "save_checkpoints_remote": True,
+ "save_tensorboard_remote": True,
+ "save_logs_remote": True,
+ "save_checkpoint_as_artifact": True,
+ },
+ "phase_callbacks": [
+ WandBDetectionValidationPredictionLoggerCallback(class_names=labels),
+ ]
+ }
+ return (train_params,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, we perform the training.
+ """)
+ return
+
+
+@app.cell
+def _(net, train_data, train_params, trainer, val_data):
+ trainer.train(model=net, training_params=train_params, train_loader=train_data, valid_loader=val_data)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's how the validation predictions look at the end of the training 👇
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/tables-log-tables-incrementally/tables_log_tables_incrementally.py b/marimo/convert/tables-log-tables-incrementally/tables_log_tables_incrementally.py
new file mode 100644
index 00000000..07f280d9
--- /dev/null
+++ b/marimo/convert/tables-log-tables-incrementally/tables_log_tables_incrementally.py
@@ -0,0 +1,100 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Example: Log Tables Incrementally
+
+ As you log more and more data to a Table over time, log incrementally to the same table using the pattern below. You will see each version logged as a new artifact, but data is deduplicated and shared between versions, meaning that we minimize redundant upload and storage. The final artifact version will contain the ultimate table, with all the data concatenated together.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(wandb):
+ PROJECT_NAME = "incremental_table_demo"
+ TABLE_NAME = "incremental_table"
+ TABLE_COLUMNS = ["col_1", "col_2"]
+
+ # Let's build up a table over 5 runs
+ for i in range(5):
+ # Create a new run
+ run = wandb.init(project=PROJECT_NAME)
+
+ # Create an artifact to hold the partitioned table
+ # Setting `incremental=True` allows you to append to the last version
+ # of the artifact without downloading everything locally
+ art = wandb.Artifact(TABLE_NAME, "example", incremental=True)
+
+ # Create a Partitioned Table pointing to a directory in the artifact (only
+ # need to do this once)
+ if i == 0:
+ parts_dir = "{}_parts".format(TABLE_NAME)
+ tab = wandb.data_types.PartitionedTable(parts_dir)
+ art.add(tab, TABLE_NAME)
+
+ # Create the table, and add it to the artifact
+ sub_tab = wandb.Table(data=[["a", i], ["b", i**2]], columns=TABLE_COLUMNS)
+ tab_path = "{}/tab_{}".format(parts_dir, i)
+ art.add(sub_tab, tab_path)
+
+ # Log the artifact
+ run.log_artifact(art)
+
+ # Finish the run
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/tables-w-b-tables-quickstart/tables_w_b_tables_quickstart.py b/marimo/convert/tables-w-b-tables-quickstart/tables_w_b_tables_quickstart.py
new file mode 100644
index 00000000..2091e007
--- /dev/null
+++ b/marimo/convert/tables-w-b-tables-quickstart/tables_w_b_tables_quickstart.py
@@ -0,0 +1,118 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Tables for Data Visualization
+
+ Try logging tabular data to visualize and query in the Weights & Biases interactive dashboard.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install Weights & Biases logging library
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ # Import libraries
+ import numpy as np
+ import pandas as pd
+ from sklearn.datasets import load_iris
+
+ return load_iris, np, pd
+
+
+@app.cell
+def _(load_iris, np, pd):
+ # Download a simple dataset
+ iris = load_iris()
+ # Load it into a dataframe
+ iris_dataframe = pd.DataFrame(data=np.c_[iris['data'], iris['target']],
+ columns=iris['feature_names'] + ['target'])
+ return (iris_dataframe,)
+
+
+@app.cell
+def _(iris_dataframe, wandb):
+ # Start a W&B run to log data
+ wandb.init(project="Tables-Quickstart")
+ # Log the dataframe to visualize
+ wandb.log({"iris": iris_dataframe})
+ # Finish the run (useful in notebooks)
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once you execute the code cells above, look for the blue [run page](https://docs.wandb.com/ref/app/pages/run-page) link in the console, and click to view the dashboard in the Weights & Biases app.
+
+ [Here's an example dashboard](https://wandb.ai/wandb/Tables%20Quickstart?workspace=user-carey) from a previous execution of this notebook.
+
+ You can also run the cell below to [render the dashboard inside the notebook](http://wandb/me/jupyter-interact-colab).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %wandb charlesfrye/Tables-Quickstart -h 1024
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/tensorboard-tensorboard-and-weights-and-biases/tensorboard_tensorboard_and_weights_and_biases.py b/marimo/convert/tensorboard-tensorboard-and-weights-and-biases/tensorboard_tensorboard_and_weights_and_biases.py
new file mode 100644
index 00000000..792fa846
--- /dev/null
+++ b/marimo/convert/tensorboard-tensorboard-and-weights-and-biases/tensorboard_tensorboard_and_weights_and_biases.py
@@ -0,0 +1,166 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ By the end of this colab you will have a TensorBoard server running in Weights & Biases, just like this:
+
+ 
+
+ This code is modified from the offical TensorBoard [getting started](https://www.tensorflow.org/tensorboard/get_started) code
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🪴 Start a Weights & Biases run
+ When using Weights & Biases for the first time you will need to:
+
+ 1️⃣. Sign-up for a free W&B [account here](https://wandb.ai/site)
+
+ 2️⃣. Create a new W&B [API key at your settings page](https://wandb.ai/settings) and store it securely. API keys can only be viewed once when created.
+
+ 3️⃣. Initialise a W&B run with wandb.init and you will be prompted to enter your API key to log in
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -qqq wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Initialising a Weights & Biases run with `sync_tensorboard=True` will enable wandb to pick up
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ run = wandb.init(project="my-wonderful-project", sync_tensorboard=True)
+ return (run,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🌿 Create Your Dataset and Model
+ """)
+ return
+
+
+@app.cell
+def _():
+ import tensorflow as tf
+
+ mnist = tf.keras.datasets.mnist
+
+ (x_train, y_train),(x_test, y_test) = mnist.load_data()
+ x_train, x_test = x_train / 255.0, x_test / 255.0
+
+ def create_model():
+ return tf.keras.models.Sequential([
+ tf.keras.layers.Flatten(input_shape=(28, 28)),
+ tf.keras.layers.Dense(512, activation='relu'),
+ tf.keras.layers.Dropout(0.2),
+ tf.keras.layers.Dense(10, activation='softmax')
+ ])
+
+ model = create_model()
+ model.compile(optimizer='adam',
+ loss='sparse_categorical_crossentropy',
+ metrics=['accuracy'])
+ return model, tf, x_test, x_train, y_test, y_train
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🌲 Train Your Model and Log to TensorBoard AND Weights & Biases
+
+ The tensorboard logs will be automatically picked up by Weights & Biases and logged
+ """)
+ return
+
+
+@app.cell
+def _(model, tf, x_test, x_train, y_test, y_train):
+ tensorboard_callback = tf.keras.callbacks.TensorBoard(histogram_freq=1)
+
+ model.fit(x=x_train,
+ y=y_train,
+ epochs=5,
+ validation_data=(x_test, y_test),
+ callbacks=[tensorboard_callback])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## (Notebook only) Finish the Weights & Biases Run
+ """)
+ return
+
+
+@app.cell
+def _(run):
+ run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Documentation
+
+ You can find additional documentation of how to use [Weights & Biases with Tensorboard here](https://docs.wandb.ai/guides/integrations/tensorboard)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/tensorflow-convert-imagenette-tfrecord/tensorflow_convert_imagenette_tfrecord.py b/marimo/convert/tensorflow-convert-imagenette-tfrecord/tensorflow_convert_imagenette_tfrecord.py
new file mode 100644
index 00000000..98356ff4
--- /dev/null
+++ b/marimo/convert/tensorflow-convert-imagenette-tfrecord/tensorflow_convert_imagenette_tfrecord.py
@@ -0,0 +1,259 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ #! wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz
+ subprocess.call(['wget', 'https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz'])
+ #! tar -xzf imagenette2-320.tgz
+ subprocess.call(['tar', '-xzf', 'imagenette2-320.tgz'])
+ #! rm imagenette2-320.tgz
+ subprocess.call(['rm', 'imagenette2-320.tgz'])
+ return
+
+
+@app.cell
+def _():
+ import os
+ import cv2
+ import math
+ import wandb
+ import random
+ import numpy as np
+ from glob import glob
+ from PIL import Image
+ import tensorflow as tf
+ from tqdm.auto import tqdm
+ import matplotlib.pyplot as plt
+
+ return Image, glob, math, os, plt, random, tf, tqdm, wandb
+
+
+@app.cell
+def _():
+ LABEL_DICT = {
+ "n01440764": ["tench", 0],
+ "n02102040": ["english_springer", 1],
+ "n02979186": ["cassette_player", 2],
+ "n03000684": ["chain_saw", 3],
+ "n03028079": ["church", 4],
+ "n03394916": ["french_horn", 5],
+ "n03417042": ["grabage_truck", 6],
+ "n03425413": ["gas_pump", 7],
+ "n03445777": ["golf_ball", 8],
+ "n03888257": ["parachute", 9]
+ }
+ return (LABEL_DICT,)
+
+
+@app.cell
+def _(wandb):
+ wandb.init(
+ project="simple-training-loop",
+ entity="jax-series",
+ job_type="tfrecord"
+ )
+ return
+
+
+@app.cell
+def _(LABEL_DICT, tf):
+ def create_example(image_file, label):
+ feature = {
+ "image": tf.train.Feature(
+ bytes_list=tf.train.BytesList(
+ value=[tf.io.read_file(image_file).numpy()]
+ )
+ ),
+ "label": tf.train.Feature(
+ int64_list=tf.train.Int64List(value=[LABEL_DICT[label][1]])
+ ),
+ "label_name": tf.train.Feature(
+ bytes_list=tf.train.BytesList(
+ value=[LABEL_DICT[label][0].encode('utf8')]
+ )
+ )
+ }
+ return tf.train.Example(
+ features=tf.train.Features(feature=feature)
+ )
+
+ return (create_example,)
+
+
+@app.cell
+def _(glob, os, random):
+ train_images = glob(os.path.join("imagenette2-320", "train/*/*.JPEG"))
+ random.shuffle(train_images)
+ train_labels = [img.split("/")[-2] for img in train_images]
+
+ val_images = glob(os.path.join("imagenette2-320", "val/*/*.JPEG"))
+ random.shuffle(val_images)
+ val_labels = [img.split("/")[-2] for img in val_images]
+ return train_images, train_labels, val_images, val_labels
+
+
+@app.cell
+def _(Image, LABEL_DICT, create_example, math, os, tf, tqdm, wandb):
+ def chunkify(input_list, chunk_size):
+ chunk_size = max(1, chunk_size)
+ return [
+ input_list[i: i + chunk_size]
+ for i in range(0, len(input_list), chunk_size)
+ ]
+
+
+ def create_tfrecords(images, labels, max_chunk_size: int, dump_dir: str):
+ os.makedirs(dump_dir)
+ num_chunks = math.ceil(len(images) / max_chunk_size)
+ print("Total number of image-label pairs:", len(images))
+ print("Total number of image-label pair chunks:", num_chunks)
+ image_chunks = chunkify(images, max_chunk_size)
+ label_chunks = chunkify(labels, max_chunk_size)
+ table = wandb.Table(columns=[
+ "Image", "Label-Name", "Label-ID", "Split-Name", "Chunk-ID"
+ ])
+ for idx in range(num_chunks):
+ image_chunk = image_chunks[idx]
+ label_chunk = label_chunks[idx]
+ current_chunk_size = len(image_chunk)
+ file_name = "%.2i-%.3i.tfrec" % (idx + 1, current_chunk_size)
+ tfrecord_file = os.path.join(dump_dir, file_name)
+ writer = tf.io.TFRecordWriter(tfrecord_file)
+ progress_bar = tqdm(
+ range(current_chunk_size),
+ desc=f"Writing {file_name}"
+ )
+ for chunk_idx in progress_bar:
+ image = Image.open(image_chunk[chunk_idx])
+ table.add_data(
+ wandb.Image(image),
+ LABEL_DICT[label_chunk[chunk_idx]][0],
+ LABEL_DICT[label_chunk[chunk_idx]][1],
+ dump_dir.split("/")[-1],
+ idx
+ )
+ example = create_example(
+ image_chunk[chunk_idx], label_chunk[chunk_idx]
+ )
+ writer.write(example.SerializeToString())
+ writer.close()
+ return table
+
+ return (create_tfrecords,)
+
+
+@app.cell
+def _(create_tfrecords, train_images, train_labels, val_images, val_labels):
+ print("Creating TFRecords for train data...")
+ train_table = create_tfrecords(
+ train_images,
+ train_labels,
+ max_chunk_size=512,
+ dump_dir="tfrecords/train"
+ )
+
+ print("Creating TFRecords for validation data...")
+ val_table = create_tfrecords(
+ val_images,
+ val_labels,
+ max_chunk_size=512,
+ dump_dir="tfrecords/val"
+ )
+ return train_table, val_table
+
+
+@app.cell
+def _(train_table, val_table, wandb):
+ wandb.log({"Train-Data": train_table})
+ wandb.log({"Validation-Data": val_table})
+ return
+
+
+@app.cell
+def _(glob, tf):
+ def parse_tfrecord(example):
+ example = tf.io.parse_single_example(
+ example, {
+ "image": tf.io.FixedLenFeature([], tf.string),
+ "label": tf.io.VarLenFeature(tf.int64),
+ "label_name": tf.io.VarLenFeature(tf.string)
+ }
+ )
+ example["image"] = tf.io.decode_jpeg(example["image"], channels=3)
+ example["label"] = tf.sparse.to_dense(example["label"])
+ example["label_name"] = tf.sparse.to_dense(example["label_name"])
+ return example
+
+
+ raw_dataset = tf.data.TFRecordDataset(glob("./tfrecords/train/*"))
+ parsed_dataset = raw_dataset.map(parse_tfrecord)
+ return (parsed_dataset,)
+
+
+@app.cell
+def _(parsed_dataset, plt):
+ for features in parsed_dataset.take(1):
+ plt.imshow(features["image"].numpy())
+ label = features["label"].numpy()
+ label_name = features["label_name"].numpy()
+ plt.title(f"{label}-{label_name}")
+ plt.show()
+ return
+
+
+@app.cell
+def _(wandb):
+ artifact = wandb.Artifact(
+ 'imagenette-tfrecords',
+ type='dataset',
+ metadata={
+ "author": "Jeremy Howard",
+ "title": "imagenette",
+ "url": "https://github.com/fastai/imagenette/",
+ "source": "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz"
+ }
+ )
+ artifact.add_dir('tfrecords')
+ wandb.log_artifact(artifact, aliases=["320px"])
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/tensorflow-hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/tensorflow_hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py b/marimo/convert/tensorflow-hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/tensorflow_hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py
new file mode 100644
index 00000000..39c4fe51
--- /dev/null
+++ b/marimo/convert/tensorflow-hyperparameter-optimization-in-tensorflow-using-w-b-sweeps/tensorflow_hyperparameter_optimization_in_tensorflow_using_w_b_sweeps.py
@@ -0,0 +1,477 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🧹 Weights & Biases Sweep + 🌊 TensorFlow 2.x
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use Weights & Biases for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use Weights & Biases Sweeps to automate hyperparameter optimization and explore the space of possible models, complete with interactive dashboards like this:
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🤔 Why Should I Use Sweeps?
+
+ * **Quick setup**: With just a few lines of code you can run W&B sweeps.
+ * **Transparent**: We cite all the algorithms we're using, and [our code is open source](https://github.com/wandb/client/tree/master/wandb/sweeps).
+ * **Powerful**: Our sweeps are completely customizable and configurable. You can launch a sweep across dozens of machines, and it's just as easy as starting a sweep on your laptop.
+
+ **[Check out the official documentation $\rightarrow$](https://docs.wandb.com/sweeps)**
+
+ ## What this notebook covers
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ * Simple steps to get started with W&B Sweep with custom training loop in TensorFlow.
+ * We will find best hyperparameters for our image classification task.
+
+ **Note**: Sections starting with _Step_ are all you need to perform hyperparameter sweep in existing code.
+ The rest of the code is there to set up a simple example.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🚀 Install, Import, and Log in
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 0️⃣: Install W&B
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 1️⃣: Import W&B and Login
+ """)
+ return
+
+
+@app.cell
+def _():
+ import tqdm
+ import tensorflow as tf
+ from tensorflow import keras
+ from tensorflow.keras.datasets import cifar10
+
+ import os
+ import numpy as np
+ import pandas as pd
+ import matplotlib.pyplot as plt
+
+ return keras, np, tf, tqdm
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.integration.keras import WandbCallback
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > Side note: If this is your first time using W&B or you are not logged in, the link that appears after running `wandb.login()` will take you to sign-up/login page. Signing up is as easy as a few clicks.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👩🍳 Prepare Dataset
+ """)
+ return
+
+
+@app.cell
+def _(keras, np):
+ # Prepare the training dataset
+ (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
+
+ x_train = x_train/255.
+ x_test = x_test/255.
+ x_train = np.reshape(x_train, (-1, 784))
+ x_test = np.reshape(x_test, (-1, 784))
+ return x_test, x_train, y_test, y_train
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🧠 Define the Model and Training Loop
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🏗️ Build a Simple Classifier MLP
+ """)
+ return
+
+
+@app.cell
+def _(keras, tf):
+ def Model():
+ inputs = keras.Input(shape=(784,), name="digits")
+ x1 = keras.layers.Dense(64, activation="relu")(inputs)
+ x2 = keras.layers.Dense(64, activation="relu")(x1)
+ outputs = keras.layers.Dense(10, name="predictions")(x2)
+
+ return keras.Model(inputs=inputs, outputs=outputs)
+
+
+ def train_step(x, y, model, optimizer, loss_fn, train_acc_metric):
+ with tf.GradientTape() as tape:
+ logits = model(x, training=True)
+ loss_value = loss_fn(y, logits)
+
+ grads = tape.gradient(loss_value, model.trainable_weights)
+ optimizer.apply_gradients(zip(grads, model.trainable_weights))
+
+ train_acc_metric.update_state(y, logits)
+
+ return loss_value
+
+
+ def test_step(x, y, model, loss_fn, val_acc_metric):
+ val_logits = model(x, training=False)
+ loss_value = loss_fn(y, val_logits)
+ val_acc_metric.update_state(y, val_logits)
+
+ return loss_value
+
+ return Model, test_step, train_step
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🔁 Write a Training Loop
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Step 3️⃣: Log metrics with `wandb.log`
+ """)
+ return
+
+
+@app.cell
+def _(np, test_step, tqdm, train_step, wandb):
+ def train(train_dataset,
+ val_dataset,
+ model,
+ optimizer,
+ loss_fn,
+ train_acc_metric,
+ val_acc_metric,
+ epochs=10,
+ log_step=200,
+ val_log_step=50):
+
+ for epoch in range(epochs):
+ print("\nStart of epoch %d" % (epoch,))
+
+ train_loss = []
+ val_loss = []
+
+ # Iterate over the batches of the dataset
+ for step, (x_batch_train, y_batch_train) in tqdm.tqdm(enumerate(train_dataset), total=len(train_dataset)):
+ loss_value = train_step(x_batch_train, y_batch_train,
+ model, optimizer,
+ loss_fn, train_acc_metric)
+ train_loss.append(float(loss_value))
+
+ # Run a validation loop at the end of each epoch
+ for step, (x_batch_val, y_batch_val) in enumerate(val_dataset):
+ val_loss_value = test_step(x_batch_val, y_batch_val,
+ model, loss_fn,
+ val_acc_metric)
+ val_loss.append(float(val_loss_value))
+
+ # Display metrics at the end of each epoch
+ train_acc = train_acc_metric.result()
+ print("Training acc over epoch: %.4f" % (float(train_acc),))
+
+ val_acc = val_acc_metric.result()
+ print("Validation acc: %.4f" % (float(val_acc),))
+
+ # Reset metrics at the end of each epoch
+ train_acc_metric.reset_states()
+ val_acc_metric.reset_states()
+
+ # 3️⃣ log metrics using wandb.log
+ wandb.log({'epochs': epoch,
+ 'loss': np.mean(train_loss),
+ 'acc': float(train_acc),
+ 'val_loss': np.mean(val_loss),
+ 'val_acc':float(val_acc)})
+
+ return (train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 4️⃣: Configure the Sweep
+
+ This is where you will:
+ * Define the hyperparameters you're sweeping over
+ * Provide your hyperparameter optimization method. We have `random`, `grid` and `bayes` methods.
+ * Provide an objective and a `metric` if using `bayes`, for example to `minimize` the `val_loss`.
+ * Use `hyperband` for early termination of poorly-performing runs
+
+ #### [Check out more on Sweep Configs $\rightarrow$](https://docs.wandb.com/sweeps/configuration)
+ """)
+ return
+
+
+@app.cell
+def _():
+ sweep_config = {
+ 'method': 'random',
+ 'metric': {
+ 'name': 'val_loss',
+ 'goal': 'minimize'
+ },
+ 'early_terminate':{
+ 'type': 'hyperband',
+ 'min_iter': 5
+ },
+ 'parameters': {
+ 'batch_size': {
+ 'values': [32, 64, 128, 256]
+ },
+ 'learning_rate':{
+ 'values': [0.01, 0.005, 0.001, 0.0005, 0.0001]
+ }
+ }
+ }
+ return (sweep_config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 5️⃣: Wrap the Training Loop
+
+ You'll need a function, like `sweep_train` below,
+ that uses `wandb.config` to set the hyperparameters
+ before `train` gets called.
+ """)
+ return
+
+
+@app.cell
+def _(Model, keras, tf, train, wandb, x_test, x_train, y_test, y_train):
+ def sweep_train(config_defaults=None):
+ # Set default values
+ config_defaults = {
+ "batch_size": 64,
+ "learning_rate": 0.01
+ }
+ # Initialize wandb with a sample project name
+ wandb.init(config=config_defaults) # this gets over-written in the Sweep
+
+ # Specify the other hyperparameters to the configuration, if any
+ wandb.config.epochs = 2
+ wandb.config.log_step = 20
+ wandb.config.val_log_step = 50
+ wandb.config.architecture_name = "MLP"
+ wandb.config.dataset_name = "MNIST"
+
+ # build input pipeline using tf.data
+ train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+ train_dataset = (train_dataset.shuffle(buffer_size=1024)
+ .batch(wandb.config.batch_size)
+ .prefetch(buffer_size=tf.data.AUTOTUNE))
+
+ val_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+ val_dataset = (val_dataset.batch(wandb.config.batch_size)
+ .prefetch(buffer_size=tf.data.AUTOTUNE))
+
+ # initialize model
+ model = Model()
+
+ # Instantiate an optimizer to train the model.
+ optimizer = keras.optimizers.SGD(learning_rate=wandb.config.learning_rate)
+ # Instantiate a loss function.
+ loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+
+ # Prepare the metrics.
+ train_acc_metric = keras.metrics.SparseCategoricalAccuracy()
+ val_acc_metric = keras.metrics.SparseCategoricalAccuracy()
+
+ train(train_dataset,
+ val_dataset,
+ model,
+ optimizer,
+ loss_fn,
+ train_acc_metric,
+ val_acc_metric,
+ epochs=wandb.config.epochs,
+ log_step=wandb.config.log_step,
+ val_log_step=wandb.config.val_log_step)
+
+ return (sweep_train,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 6️⃣: Initialize Sweep and Run Agent
+ """)
+ return
+
+
+@app.cell
+def _(sweep_config, wandb):
+ sweep_id = wandb.sweep(sweep_config, project="sweeps-tensorflow")
+ return (sweep_id,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can limit the number of total runs with the `count` parameter, we will limit a 10 to make the script run fast, feel free to increase the number of runs and see what happens.
+ """)
+ return
+
+
+@app.cell
+def _(sweep_id, sweep_train, wandb):
+ wandb.agent(sweep_id, function=sweep_train, count=10)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👀 Visualize Results
+
+ Click on the **Sweep URL** link above to see your live results.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🎨 Example Gallery
+
+ See examples of projects tracked and visualized with W&B in our [Gallery →](https://app.wandb.ai/gallery)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 📏 Best Practices
+ 1. **Projects**: Log multiple runs to a project to compare them. `wandb.init(project="project-name")`
+ 2. **Groups**: For multiple processes or cross validation folds, log each process as a runs and group them together. `wandb.init(group='experiment-1')`
+ 3. **Tags**: Add tags to track your current baseline or production model.
+ 4. **Notes**: Type notes in the table to track the changes between runs.
+ 5. **Reports**: Take quick notes on progress to share with colleagues and make dashboards and snapshots of your ML projects.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🤓 Advanced Setup
+ 1. [Environment variables](https://docs.wandb.com/library/environment-variables): Set API keys in environment variables so you can run training on a managed cluster.
+ 2. [Offline mode](https://docs.wandb.com/library/technical-faq#can-i-run-wandb-offline): Use `dryrun` mode to train offline and sync results later.
+ 3. [On-prem](https://docs.wandb.com/self-hosted): 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.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/tensorflow-tensorflow-initialization-methods/tensorflow_tensorflow_initialization_methods.py b/marimo/convert/tensorflow-tensorflow-initialization-methods/tensorflow_tensorflow_initialization_methods.py
new file mode 100644
index 00000000..71bcaef6
--- /dev/null
+++ b/marimo/convert/tensorflow-tensorflow-initialization-methods/tensorflow_tensorflow_initialization_methods.py
@@ -0,0 +1,127 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Packages 📦 and Basic Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install -Uq wandb
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import numpy as np
+ import tensorflow as tf
+ from wandb.keras import WandbCallback
+
+ return WandbCallback, tf, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 💿 Dataset
+
+ For the sake of simplicity, MNIST was chosen
+ """)
+ return
+
+
+@app.cell
+def _(tf):
+ mnist = tf.keras.datasets.mnist
+
+ (x_train, y_train),(x_test, y_test) = mnist.load_data()
+ x_train, x_test = x_train / 255.0, x_test / 255.0
+ return x_train, y_train
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # The Model 👷♀️
+
+ Initializers used (from `tf.keras.initializers`):-
+
+ * Zeros
+ * LeCunNormal
+ * GlorotNormal
+ * HeNormal
+ """)
+ return
+
+
+@app.cell
+def _(tf):
+ initializer = tf.keras.initializers.LecunNormal()
+
+ model = tf.keras.models.Sequential([
+ tf.keras.layers.Flatten(input_shape=(28, 28)),
+ tf.keras.layers.Dense(128, activation='tanh', kernel_initializer=initializer),
+ tf.keras.layers.Dense(64, activation='tanh', kernel_initializer=initializer),
+ tf.keras.layers.Dense(32, activation='tanh', kernel_initializer=initializer),
+ tf.keras.layers.Dense(10, activation='softmax')
+ ])
+
+ model.compile(optimizer='adam',
+ loss='sparse_categorical_crossentropy',
+ metrics=['accuracy'])
+ return (model,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training 💪🏻
+
+ The weights and gradients were logged using the helpful `log_gradients` and `log_weights` parameters of `WandbCallback()`
+ """)
+ return
+
+
+@app.cell
+def _():
+ PROJECT = "tensorflow_initialization_methods"
+ return (PROJECT,)
+
+
+@app.cell
+def _(PROJECT, WandbCallback, model, wandb, x_train, y_train):
+ run = wandb.init(project=PROJECT)
+
+ model.fit(x_train, y_train, epochs=20, callbacks=[WandbCallback(training_data = (x_train, y_train),log_weights = True, log_gradients = True, save_model = False)])
+
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/ultralytics-00-inference/ultralytics_00_inference.py b/marimo/convert/ultralytics-00-inference/ultralytics_00_inference.py
new file mode 100644
index 00000000..20e69ccc
--- /dev/null
+++ b/marimo/convert/ultralytics-00-inference/ultralytics_00_inference.py
@@ -0,0 +1,165 @@
+# /// script
+# dependencies = ["ultralytics", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install Dependencies
+
+ - Install Ultralytics using `pip install ultralytics`. In order to learn about more ways to install Ultralytics, you can check out the [official docs](https://docs.ultralytics.com/quickstart/#install-ultralytics).
+
+ - Then, you need to install the [`feat/ultralytics`](https://github.com/wandb/wandb/tree/feat/ultralytics) branch from W&B, which currently houses the out-of-the-box integration for Ultralytics.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install WandB and Ultralytics
+ # packages added via marimo's package management: wandb ultralytics !pip install -q -U wandb ultralytics
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Note:** The Ultralytcs integration will be soon available as a fully supported feature on Weights & Biases once [this pull request](https://github.com/wandb/wandb/pull/5867) is merged.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using Ultralytics with Weights & Biases
+
+ In order to use the W&B integration with Ultralytics, we need to import the `wandb.yolov8.add_wandb_callback` function.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.integration.ultralytics import add_wandb_callback
+
+ from ultralytics import YOLO
+
+ return YOLO, add_wandb_callback, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, let us download a few images to test the integration on. You can use your own images, videos or camera sources. For more information on inference sources, you can check out the [official docs](https://docs.ultralytics.com/modes/predict/).
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! wget https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img1.png
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img1.png'])
+ #! wget https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img2.png
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img2.png'])
+ #! wget https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img4.png
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img4.png'])
+ #! wget https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img5.png
+ subprocess.call(['wget', 'https://raw.githubusercontent.com/wandb/examples/ultralytics/colabs/ultralytics/assets/img5.png'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we initialize a W&B [run](https://docs.wandb.ai/guides/runs) using `wandb.init`.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # Initialize Weights & Biases run
+ wandb.init(project="ultralytics", job_type="inference")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we initialize the `YOLO` model of our choice, and invoke the `add_wandb_callback` function on it before performing inference with the model. This would ensure that when we perform inference, it would automatically log the images overlayed with our [interactive overlays for computer vision tasks](https://docs.wandb.ai/guides/track/log/media#image-overlays-in-tables) along with additional insights in a [`wandb.Table`](https://docs.wandb.ai/guides/data-vis).
+ """)
+ return
+
+
+@app.cell
+def _(YOLO, add_wandb_callback, wandb):
+ model_name = 'yolov8n' #@param {type:"string"}
+
+ # Initialize YOLO Model
+ model = YOLO(f"{model_name}.pt")
+
+ # Add Weights & Biases callback for Ultralytics
+ add_wandb_callback(model, enable_model_checkpointing=True)
+
+ # Perform prediction which automatically logs to a W&B Table
+ # with interactive overlays for bounding boxes, segmentation masks
+ model(["./assets/img1.jpeg", "./assets/img3.png", "./assets/img4.jpeg", "./assets/img5.jpeg"])
+
+ # Finish the W&B run
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, you can check out the following notebook to learn how to perform experiment tracking and visualize validation predictions during training using Weights & Biases in the following notebook:
+
+ [](http://wandb.me/ultralytics-train)
+
+ In order to learn more about using Weights & Biases with Ultralytics, you can also read the report: [**Supercharging Ultralytics with Weights & Biases**](https://wandb.ai/geekyrakshit/ultralytics/reports/Supercharging-Ultralytics-with-Weights-Biases--Vmlldzo0OTMyMDI4)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/ultralytics-01-train-val/ultralytics_01_train_val.py b/marimo/convert/ultralytics-01-train-val/ultralytics_01_train_val.py
new file mode 100644
index 00000000..74d2d613
--- /dev/null
+++ b/marimo/convert/ultralytics-01-train-val/ultralytics_01_train_val.py
@@ -0,0 +1,162 @@
+# /// script
+# dependencies = ["ultralytics", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🔥🔥 Explore Predictions from Ultralytics models using Weights & Biases 🪄🐝
+
+
+
+ This notebook demonstrates a typical workflow of using an [Ultralytics](https://docs.ultralytics.com/modes/predict/) model for training, fine-tuning, and validation and performing experiment tracking, model-checkpointing, and visualization of the model's performance using [Weights & Biases](https://wandb.ai/site).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install Dependencies
+
+ - Install Ultralytics using `pip install ultralytics`. In order to learn about more ways to install Ultralytics, you can check out the [official docs](https://docs.ultralytics.com/quickstart/#install-ultralytics).
+
+ - Then, you need to install the [`feat/ultralytics`](https://github.com/wandb/wandb/tree/feat/ultralytics) branch from W&B, which currently houses the out-of-the-box integration for Ultralytics.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Install WandB and Ultralytics
+ # packages added via marimo's package management: wandb ultralytics !pip install -q -U wandb ultralytics
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Note:** The Ultralytcs integration will be soon available as a fully supported feature on Weights & Biases once [this pull request](https://github.com/wandb/wandb/pull/5867) is merged.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using Ultralytics with Weights & Biases
+
+ In order to use the W&B integration with Ultralytics, we need to import the `wandb.yolov8.add_wandb_callback` function.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ from wandb.integration.ultralytics import add_wandb_callback
+
+ from ultralytics import YOLO
+
+ return YOLO, add_wandb_callback, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, we initialize the `YOLO` model of our choice, and invoke the `add_wandb_callback` function on it before performing inference with the model. This would ensure that when we perform training, fine-tuning, validation, or inference, it would automatically log the experiment logs and the images overlayed with both ground-truth and the respective prediction results using the [interactive overlays for computer vision tasks](https://docs.wandb.ai/guides/track/log/media#image-overlays-in-tables) on W&B along with additional insights in a [`wandb.Table`](https://docs.wandb.ai/guides/data-vis).
+ """)
+ return
+
+
+@app.cell
+def _(YOLO, add_wandb_callback, wandb):
+ model_name = "yolov8n" #@param {type:"string"}
+ dataset_name = "coco128.yaml" #@param {type:"string"}
+
+ # Initialize YOLO Model
+ model = YOLO(f"{model_name}.pt")
+
+ # Add Weights & Biases callback for Ultralytics
+ add_wandb_callback(model, enable_model_checkpointing=True)
+
+ # Train/fine-tune your model
+ # At the end of each epoch, predictions on validation batches are logged
+ # to a W&B table with insightful and interactive overlays for
+ # computer vision tasks
+ model.train(project="ultralytics", data=dataset_name, epochs=5, imgsz=640)
+ model.val()
+
+ # Finish the W&B run
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Sample Experiment Tracking
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Epoch-wise results visualized
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Next, you can check out the following notebook to learn how to perform inference and visualize predictions during training using Weights & Biases in the following notebook:
+
+ [](http://wandb.me/ultralytics-inference)
+
+ In order to learn more about using Weights & Biases with Ultralytics, you can also read the report: [**Supercharging Ultralytics with Weights & Biases**](https://wandb.ai/geekyrakshit/ultralytics/reports/Supercharging-Ultralytics-with-Weights-Biases--Vmlldzo0OTMyMDI4)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-artifact-fundamentals/wandb_artifacts_artifact_fundamentals.py b/marimo/convert/wandb-artifacts-artifact-fundamentals/wandb_artifacts_artifact_fundamentals.py
new file mode 100644
index 00000000..cf920d98
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-artifact-fundamentals/wandb_artifacts_artifact_fundamentals.py
@@ -0,0 +1,415 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use [Weights & Biases](https://wandb.com) for machine learning experiment tracking, dataset and model versioning and management, collaboration and more.
+
+
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Use W&B Artifacts to track and version data as the inputs and outputs of your W&B Runs. In addition to logging hyperparameters, metadata, and metrics to a run, you can use an artifact to log the dataset used to train the model as input and the resulting model checkpoints as outputs.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Set Up
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In order to use Weights & Biases, you will need the `wandb` package installed. You can install it as follows within Colab.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once it is installed, the next step is to import it into your script or notebook with `import wandb`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We also need to authenticate to the Weights & Biases server. There are various ways of doing this, including for [remote or non-interactice workflows](https://docs.wandb.ai/guides/track/environment-variables), but given this is running interactively, we can use `wandb.login()`.
+
+ If we are not already authenticated, a link will appear which you can use to do so.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Create a Dataset
+ Let's create some datasets that we can work with in this example.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ import numpy as np
+ import csv
+
+ directory = "dataset"
+ os.makedirs(directory, exist_ok=True)
+ file1, file2 = os.path.join(directory, "file1.csv"), os.path.join(directory, "file2.csv")
+
+ def generate_dummy_data(num_samples):
+ data = [
+ np.random.normal(50, 10, num_samples),
+ np.random.randint(1, 100, num_samples),
+ np.random.choice(['A', 'B', 'C', 'D'], num_samples),
+ np.random.uniform(0.0, 1.0, num_samples)
+ ]
+ return zip(*data)
+
+ def save_to_csv(file, data):
+ with open(file, 'w', newline='') as f:
+ writer = csv.writer(f)
+ writer.writerow(['feature1', 'feature2', 'feature3', 'feature4'])
+ writer.writerows(data)
+
+ num_samples = 100
+ save_to_csv(file1, generate_dummy_data(num_samples))
+ save_to_csv(file2, generate_dummy_data(num_samples))
+ return (directory,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Create An Artifact
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The general workflow for creating an Artifact is:
+
+ 1. Intialize a run.
+ 2. Create an Artifact.
+ 3. Add a any files or directories to the new Artifact that you want to track and version.
+ 4. Log the artifact in the W&B platform.
+
+ The most straightforward way of accomplishing this is the second line of code in the example below, which will log, track and version a new dataset (i.e. do points 2, 3, and 4 above in one step).
+ """)
+ return
+
+
+@app.cell
+def _(directory, wandb):
+ _run = wandb.init(project='artifact-basics')
+ _run.log_artifact(artifact_or_path=f'{directory}/file1.csv', name='my_first_artifact', type='dataset')
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In the above example we first initalize a run using [`wandb.init()`](https://docs.wandb.ai/ref/python/init) the `artifact-basics` project. If this project doesn't exist, it will be created. If it alreadt exists, a new W&B Run will be added to it.
+
+ In the second line we actually log the Artifact with [`run.log_artifact()`](https://docs.wandb.ai/ref/python/public-api/run#log_artifact). In this example, we use three common arguments to the function.
+ 1. With `artifact_or_path` we specifiy the path to where the data we want to version exists. Any file or directory can be added here.
+ 2. with `name` we give the artifact a name within Weights & Biases that we will use to access it.
+ 3. With `type` we give the artifact a higher level grouping. For example, we may have multiple artifacts of type data, and multiple artifacts of type model.
+
+ See the [Artifacts Reference](https://docs.wandb.ai/ref/python/artifact) guide for more information and other commonly used arguments, including how to store additional metadata.
+
+ Each time the above `log_artifact` is executed, wandb will create a new version of the Artifact within Weights & Biases if the underlying data has changed.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ An alternative approach that offers more control (at the expense of more lines of code) can be seen below.
+ """)
+ return
+
+
+@app.cell
+def _(directory, wandb):
+ _run = wandb.init(project='artifact-basics')
+ _artifact = wandb.Artifact('my_first_artifact', type='dataset')
+ _artifact.add_file(local_path=f'{directory}/file1.csv')
+ # the below will add two individual files to the artifact.
+ _artifact.add_file(local_path=f'{directory}/file2.csv')
+ _artifact.add_dir(local_path=f'{directory}')
+ # or the below if you wanted to add the entire directory contents.
+ _run.log_artifact(_artifact)
+ # explictly log the artifact to Weights & Biases.
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In the above example, lines 3-5 will create a new Artifact within your Weights & Biases project. With the resulting artifact object, you can call the [`artifact.add_file`](https://docs.wandb.ai/ref/python/artifact#add_file) or [`artifact.add_dir`](https://docs.wandb.ai/ref/python/artifact#add_dir) functions in order to add as many files and directories to the Artifact as you want. Once added, the artifact must then be explictly logged to Weights & Biases.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Use an Artifact
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ When you want to use a specific version of an Artifact in a downstream task, you can specify the specific version you would like to use via either `v0`, `v1`, `v2` and so on, or via specific aliases you may have added. The `latest` alias always refers to the most recent version of the Artifact logged.
+
+ The proceeding code snippet specifies that the W&B Run will use an artifact called `my_first_artifact` with the alias `latest`:
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _run = wandb.init(project='artifact-basics')
+ _artifact = _run.use_artifact(artifact_or_name='my_first_artifact:latest') # this creates a reference within Weights & Biases that this artifact was used by this run.
+ path = _artifact.download() # this downloads the artifact from Weights & Biases to your local system where the code is executing.
+ print(f'Data directory located at {path}')
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ For more information on ways to customize your Artifact download, including via the command line, see the [Download and Usage guide](https://docs.wandb.ai/guides/artifacts/download-and-use-an-artifact).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Create a new Artifact version
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's say we want to modify our dataset while also tracking and versioning these changes. In the below example we will subsample our dataset and save it as a new file. We will use the [Pandas](https://pandas.pydata.org/pandas-docs/stable/index.html) library to read our CSV file.
+
+ In the second block of code we will log it to Weights & Biases under the same Artifact name (*my_first_artifact*) so that Weights & Biases knows that this is a new version of an existing artifact.
+ """)
+ return
+
+
+@app.cell
+def _(directory):
+ import pandas
+ df = pandas.read_csv(f"{directory}/file1.csv")
+ # subsample to 50% of the original size
+ df_subsampled = df.sample(frac=0.5, random_state=1)
+ # save the subsampled dataframe to a new file.
+ df_subsampled.to_csv(f"{directory}/file1.csv", index=False)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now we have a new subsampled version of our dataset locally, we can log the new version to Weights & Biases.
+ """)
+ return
+
+
+@app.cell
+def _(directory, wandb):
+ _run = wandb.init(project='artifact-basics')
+ _run.log_artifact(artifact_or_path=f'{directory}/file1.csv', name='my_first_artifact', type='dataset', aliases=['subsampled'])
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now the sampled dataset will be logged to the `my_first_artifact` Artifact as a new version.
+
+ The Artifact has also been given a custom `alias`, which is a unique label for this Artifact version. While the `alias` is currently `subsampled`, the default aliases is `vN`, where `N` is the number of versions the Artifact has. This increments automatically. You can always access specific versions of an Artifact by using an alias.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Update Artifact version metadata
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can update the `description`, `metadata`, and `alias` of an artifact on the W&B platform during or outside a W&B Run.
+
+ This example changes the `description` of the `my_first_artifact` artifact inside a run:
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _run = wandb.init(project='artifact-basics')
+ _artifact = _run.use_artifact(artifact_or_name='my_first_artifact:subsampled')
+ _artifact.description = 'This is an edited description.'
+ _artifact.metadata = {'source': 'local disk', 'internal data owner': 'platform team'}
+ _artifact.save() # persists changes to an Artifact's properties
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Use the Artifact within your pipelines
+ Once the artifact is tracked and versioned within Weights & Biases it's now easy to integrate it into your ML workflows.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _run = wandb.init(project='artifact-basics')
+ _artifact = _run.use_artifact(artifact_or_name='my_first_artifact:latest')
+ # the below is left as an exercise to the reader
+ # train model
+ # log model as artifact
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Navigate the Artifacts UI
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can also manage your Artifacts via the W&B platform. This can give you insight into your model's performance or dataset versioning. To navigate to the relevant information, click this [link](https://wandb.ai/wandb/artifact-basics/overview), then click on the **Artifacts** tab.
+
+ Navigating to the **Lineage** section in the tab will show the dependency graph formed by calling `run.use_artifact()` when an Artifact is an input to a run, and `run.log_artifact()` when an Artifact is output to a run. This helps visualize the relationship between different model versions and other objects like datasets and jobs in your project. Click [this](https://wandb.ai/wandb/artifact-basics/artifacts/dataset/my_first_artifact/v0/lineage) link to navigate to the project's lineage page.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Naturally, as you integrate W&B Artifacts into your workflow, lineage graphs such as [this interactive example](https://wandb.ai/wandb-smle/artifact_workflow/artifacts/model/quant_model/v16/lineage) will be built up over time, giving you reproducibility, governance, and auditability.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Next steps
+ 1. [Artifacts Python reference documentation](https://docs.wandb.ai/ref/python/artifact): Deep dive into artifact parameters and advanced methods.
+ 2. [Lineage](https://docs.wandb.ai/guides/artifacts/explore-and-traverse-an-artifact-graph): View lineage graphs, which are automatically built when using W&B artifact system, providing an auditable visual overview of the relationships between specific artifact versions, datasets models and runs.
+ 3. [Model Registry](https://docs.wandb.ai/guides/model_registry): Learn how to centralize your best artifact versions in a shared registry.
+ 4. [Artifact Automations](https://docs.wandb.ai/guides/artifacts/project-scoped-automations): Automatically run specific Weights & Biases jobs based on changes to your artifacts, such as automatically training a new model each time a new version of the training data is logged.
+ 5. [Reference Artifacts](https://docs.wandb.ai/guides/artifacts/track-external-files#download-a-reference-artifact): Track files saved outside the W&B server, like Amazon S3 buckets, GCS buckets, Azure blobs, and more.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-artifacts-quickstart-with-w-b/wandb_artifacts_artifacts_quickstart_with_w_b.py b/marimo/convert/wandb-artifacts-artifacts-quickstart-with-w-b/wandb_artifacts_artifacts_quickstart_with_w_b.py
new file mode 100644
index 00000000..b70f1ae0
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-artifacts-quickstart-with-w-b/wandb_artifacts_artifacts_quickstart_with_w_b.py
@@ -0,0 +1,400 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # W&B Artifacts Quickstart
+
+ This tutorial shows how to get started with **W&B Artifacts** very quickly. I finetune a convnet in Keras to identify 10 types of living things in photos: plants, animals, insects, etc.
+
+ * [follow along in a W&B Report](https://wandb.ai/wandb/arttest/reports/Artifacts-Quickstart--VmlldzozNTAzMDM)
+ * [see the Artifacts API and documentation](https://docs.wandb.com/artifacts/api)
+
+ This demo will generate an experiment workflow like the following:
+
+ 
+
+ In this example we're using Google Colab as a convenient hosted environment, but you can run your own training scripts from anywhere and visualize metrics with W&B's experiment tracking tool.
+
+ ## Sign up or login
+
+ [Sign up or login](https://wandb.ai/login) to W&B to see and interact with your experiments in the browser.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download sample data: Choose 1 of 3 sizes
+
+ Choose one of the three dataset size options below to run the rest of the demo. With fewer images, you'll run through the demo much faster and use less storage space. With more images, you'll get more realistic model training and more interesting results and examples to explore.
+
+ Note: **for the largest dataset, this stage might take a few minutes**. If you end up needing to rerun a cell, comment out the first capture line (change ```%%capture``` to ```#%%capture``` ) so you can respond to the prompt about re-downloading the dataset (and see the progress bar).
+
+ Each zipped directory contains randomly sampled images from the [iNaturalist dataset](https://github.com/visipedia/inat_comp), evenly distributed across 10 classes of living things like birds, insects, plants, and mammals (names given in Latin—so Aves, Insecta, Plantae, etc :).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # set SIZE to "TINY", "MEDIUM", or "LARGE"
+ # to select one of these three datasets
+ # TINY dataset: 100 images, 30MB
+ # MEDIUM dataset: 1000 images, 312MB
+ # LARGE datast: 12,000 images, 3.6GB
+
+ SIZE = "TINY"
+ return (SIZE,)
+
+
+@app.cell
+def _(SIZE):
+ if SIZE == "TINY":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_100.zip"
+ src_zip = "nature_100.zip"
+ DATA_SRC = "nature_100"
+ IMAGES_PER_LABEL = 10
+ BALANCED_SPLITS = {"train" : 8, "val" : 1, "test": 1}
+ elif SIZE == "MEDIUM":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_1K.zip"
+ src_zip = "nature_1K.zip"
+ DATA_SRC = "nature_1K"
+ IMAGES_PER_LABEL = 100
+ BALANCED_SPLITS = {"train" : 80, "val" : 10, "test": 10}
+ elif SIZE == "LARGE":
+ src_url = "https://storage.googleapis.com/wandb_datasets/nature_12K.zip"
+ src_zip = "nature_12K.zip"
+ DATA_SRC = "inaturalist_12K/train" # (technically a subset of only 10K images)
+ IMAGES_PER_LABEL = 1000
+ BALANCED_SPLITS = {"train" : 800, "val" : 100, "test": 100}
+ return BALANCED_SPLITS, DATA_SRC, IMAGES_PER_LABEL
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !curl -SL $src_url > $src_zip
+ # !unzip $src_zip
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 0: Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Start out by installing the experiment tracking library and setting up your free W&B account:
+
+ * **pip install wandb** – Install the W&B library
+ * **import wandb** – Import the wandb library
+ * **wandb login** – Login to your W&B account so you can log all your metrics in one place
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qq
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(DATA_SRC, IMAGES_PER_LABEL):
+ import os
+ from random import shuffle
+
+ # source directory for all raw data
+ SRC = DATA_SRC
+ # number of images per class label
+ # the total number of images is 10X this (10 classes)
+ TOTAL_IMAGES = IMAGES_PER_LABEL * 10
+ PROJECT_NAME = "artifacts_demo"
+ PREFIX = "inat" # convenient for tracking local data
+ return PREFIX, PROJECT_NAME, SRC, TOTAL_IMAGES, os, shuffle
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 1: Upload raw data
+ """)
+ return
+
+
+@app.cell
+def _(
+ IMAGES_PER_LABEL,
+ PREFIX,
+ PROJECT_NAME,
+ SRC,
+ TOTAL_IMAGES,
+ os,
+ shuffle,
+ wandb,
+):
+ RAW_DATA_AT = '_'.join([PREFIX, 'raw_data', str(TOTAL_IMAGES)])
+ _run = wandb.init(project=PROJECT_NAME, job_type='upload')
+ raw_data_at = wandb.Artifact(RAW_DATA_AT, type='raw_data')
+ # create an artifact for all the raw data
+ _labels = os.listdir(SRC)
+ for _l in _labels:
+ # SRC_DIR contains 10 folders, one for each of 10 class labels
+ # each folder contains images of the corresponding class
+ _imgs_per_label = os.path.join(SRC, _l)
+ if os.path.isdir(_imgs_per_label):
+ _imgs = os.listdir(_imgs_per_label)
+ shuffle(_imgs)
+ img_file_ids = _imgs[:IMAGES_PER_LABEL]
+ for f in img_file_ids: # randomize the order
+ file_path = os.path.join(SRC, _l, f)
+ raw_data_at.add_file(file_path, name=_l + '/' + f)
+ _run.log_artifact(raw_data_at)
+ # save artifact to W&B
+ _run.finish() # add file to artifact by full path
+ return (RAW_DATA_AT,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 2: Prepare a data split
+ """)
+ return
+
+
+@app.cell
+def _(BALANCED_SPLITS, PREFIX, PROJECT_NAME, RAW_DATA_AT, os, shuffle, wandb):
+ _run = wandb.init(project=PROJECT_NAME, job_type='data_split')
+ data_at = _run.use_artifact(RAW_DATA_AT + ':latest')
+ # find the most recent ("latest") version of the full raw data
+ # you can of course pass around programmatic aliases and not string literals
+ data_dir = data_at.download()
+ # download it locally (for illustration purposes/across hardware; you can
+ # also sync/version artifacts by reference)
+ DATA_SPLITS = BALANCED_SPLITS
+ ats = {}
+ # create balanced train, val, test splits
+ # each count is the number of images per label
+ for split, count in DATA_SPLITS.items():
+ ats[split] = wandb.Artifact('_'.join([PREFIX, split, 'data', str(count * 10)]), '_'.join([split, 'data']))
+ _labels = os.listdir(data_dir)
+ # wrap artifacts in dictionary for convenience
+ for _l in _labels:
+ if _l.startswith('.'):
+ continue
+ _imgs_per_label = os.listdir(os.path.join(data_dir, _l))
+ shuffle(_imgs_per_label)
+ start_id = 0
+ for split, count in DATA_SPLITS.items(): # skip non-label file
+ split_imgs = _imgs_per_label[start_id:start_id + count]
+ for img_file in split_imgs:
+ full_path = os.path.join(data_dir, _l, img_file)
+ ats[split].add_file(full_path, name=os.path.join(_l, img_file))
+ start_id += count
+ for split, artifact in ats.items(): # take a subset
+ _run.log_artifact(artifact)
+ # save all three artifacts to W&B
+ # note: yes, in this example, we are cheating and have labels for the "test" data ;)
+ _run.finish() # add file to artifact by full path # note: pass the label to the name parameter to retain it in # the data structure
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 3: Train with artifacts and save model
+ """)
+ return
+
+
+@app.cell
+def _(BALANCED_SPLITS, PREFIX, PROJECT_NAME, os, wandb):
+ NUM_TRAIN = BALANCED_SPLITS['train'] * 10
+ NUM_VAL = BALANCED_SPLITS['val'] * 10
+ NUM_EPOCHS = 1
+ MODEL_NAME = 'iv3_trained'
+ INIT_MODEL_DIR = 'init_model_keras_iv3.keras'
+ FINAL_MODEL_DIR = 'trained_keras_model_iv3.keras'
+ import numpy as np
+ from sklearn.metrics import precision_recall_curve, roc_curve
+ from sklearn.metrics import average_precision_score
+ from sklearn.preprocessing import label_binarize
+ from tensorflow.keras.applications.inception_v3 import InceptionV3
+ from tensorflow.keras.callbacks import Callback
+ from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
+ from tensorflow.keras.models import Model
+ from tensorflow.keras.preprocessing.image import ImageDataGenerator
+ from wandb.integration.keras import WandbMetricsLogger, WandbModelCheckpoint
+ config_defaults = {'num_train': NUM_TRAIN, 'num_val': NUM_VAL, 'epochs': NUM_EPOCHS, 'num_classes': 10, 'fc_size': 1024, 'img_width': 299, 'img_height': 299, 'batch_size': 32}
+
+ def finetune_inception_model(fc_size, num_classes):
+ """Load InceptionV3 with ImageNet weights, freeze it,
+ and attach a finetuning top for this classification task"""
+ base = InceptionV3(weights='imagenet', include_top='False')
+ for layer in base.layers:
+ layer.trainable = False
+ x = base.get_layer('mixed10').output
+ x = GlobalAveragePooling2D()(x)
+ x = Dense(fc_size, activation='relu')(x)
+ guesses = Dense(num_classes, activation='softmax')(x)
+ model = Model(inputs=base.input, outputs=guesses)
+ model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])
+ return model
+
+ def train():
+ """ Main training loop. This is called pretrain because it freezes
+ the InceptionV3 layers of the model and only trains the new top layers # inceptionV3 settings
+ on the new data. subsequent training phase would unfreeze all the layers
+ and finetune the whole model on the new data"""
+ _run = wandb.init(project=PROJECT_NAME, job_type='train', config=config_defaults)
+ cfg = wandb.config
+ train_at = os.path.join(PROJECT_NAME, PREFIX + '_train_data_' + str(NUM_TRAIN)) + ':latest'
+ val_at = os.path.join(PROJECT_NAME, PREFIX + '_val_data_' + str(NUM_VAL)) + ':latest'
+ train_data = _run.use_artifact(train_at, type='train_data')
+ train_dir = train_data.download()
+ val_data = _run.use_artifact(val_at, type='val_data')
+ val_dir = val_data.download()
+ train_datagen = ImageDataGenerator(rescale=1.0 / 255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True)
+ val_datagen = ImageDataGenerator(rescale=1.0 / 255)
+ train_generator = train_datagen.flow_from_directory(train_dir, target_size=(cfg.img_width, cfg.img_height), batch_size=cfg.batch_size, class_mode='categorical')
+ val_generator = val_datagen.flow_from_directory(val_dir, target_size=(cfg.img_width, cfg.img_height), batch_size=cfg.batch_size, class_mode='categorical')
+ model = finetune_inception_model(cfg.fc_size, cfg.num_classes)
+ model_artifact = wandb.Artifact('iv3', type='model', description='unmodified inception v3', metadata=dict(cfg))
+ model.save(INIT_MODEL_DIR)
+ model_artifact.add_file(INIT_MODEL_DIR)
+ _run.log_artifact(model_artifact)
+ callbacks = [WandbMetricsLogger(), WandbModelCheckpoint('checkpoint.keras')]
+ model.fit(train_generator, steps_per_epoch=cfg.num_train // cfg.batch_size, epochs=cfg.epochs, validation_data=val_generator, callbacks=callbacks, validation_steps=cfg.num_val // cfg.batch_size)
+ trained_model_artifact = wandb.Artifact(MODEL_NAME, type='model', description='trained inception v3', metadata=dict(cfg))
+ model.save(FINAL_MODEL_DIR)
+ trained_model_artifact.add_file(FINAL_MODEL_DIR)
+ _run.log_artifact(trained_model_artifact)
+ _run.finish()
+
+ return MODEL_NAME, np, train
+
+
+@app.cell
+def _(train):
+ train()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 4: Load model for inference
+ """)
+ return
+
+
+@app.cell
+def _(BALANCED_SPLITS, MODEL_NAME, PREFIX, PROJECT_NAME, np, os, wandb):
+ from tensorflow import keras
+ from tensorflow.keras.preprocessing import image
+ import pathlib
+ _run = wandb.init(project=PROJECT_NAME, job_type='inference')
+ model_at = _run.use_artifact(MODEL_NAME + ':latest')
+ artifact_dir = pathlib.Path(model_at.download())
+ print('artifact directory:', artifact_dir)
+ model_path = artifact_dir / 'trained_keras_model_iv3.keras'
+ model = keras.models.load_model(model_path, compile=False)
+ print('loaded model from', model_path)
+ TEST_DATA_AT = PREFIX + '_test_data_' + str(BALANCED_SPLITS['test'] * 10) + ':latest'
+ test_data_at = _run.use_artifact(TEST_DATA_AT)
+ test_dir = test_data_at.download()
+ _imgs = []
+ class_labels = os.listdir(test_dir)
+ for _l in class_labels:
+ if _l.startswith('.'):
+ continue
+ imgs_per_class = os.listdir(os.path.join(test_dir, _l))
+ for img in imgs_per_class:
+ img_path = os.path.join(test_dir, _l, img)
+ img = image.load_img(img_path, target_size=(299, 299))
+ img = image.img_to_array(img)
+ img = np.expand_dims(img / 255.0, axis=0)
+ _imgs.append(img)
+ preds = {}
+ _imgs = np.vstack(_imgs)
+ classes = model.predict(_imgs, batch_size=32)
+ for c in classes:
+ class_id = np.argmax(c)
+ if class_id in preds:
+ preds[class_id] += 1
+ else:
+ preds[class_id] = 1
+ print(preds)
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # More about Weights & Biases
+ We're always free for academics and open source projects. Here are some more resources:
+
+ 1. [Documentation](http://docs.wandb.com) - Python docs
+ 2. [Gallery](https://app.wandb.ai/gallery) - example reports in W&B
+ 3. [Articles](https://www.wandb.com/articles) - blog posts and tutorials
+ 4. [Community](wandb.me/slack) - join our Slack community forum
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-basic-artifacts-with-w-b/wandb_artifacts_basic_artifacts_with_w_b.py b/marimo/convert/wandb-artifacts-basic-artifacts-with-w-b/wandb_artifacts_basic_artifacts_with_w_b.py
new file mode 100644
index 00000000..ef25deb2
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-basic-artifacts-with-w-b/wandb_artifacts_basic_artifacts_with_w_b.py
@@ -0,0 +1,356 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Artifacts Quickstart
+
+
+
+ This tutorial shows how to get started with W&B Artifacts very quickly. I finetune a convnet in Keras to identify 10 types of living things in photos: plants, animals, insects, etc.
+ [Check out the companion report on W&B](https://wandb.ai/wandb/arttest/reports/Artifacts-Quickstart--VmlldzozNTAzMDM)
+
+ In this example we're using Google Colab as a convenient hosted environment, but you can run your own training scripts from anywhere and visualize metrics with W&B's experiment tracking tool.
+
+ ## Sign up or login
+
+ [Sign up or login](https://wandb.ai/login) to W&B to see and interact with your experiments in the browser.
+
+ ### Note on Artifacts storage space and deletion
+
+ Running this colab end-to-end will create at least 7GB of artifacts in your wandb account (more if you try different experiments, increase the number of epochs or examples, etc). If you'd like to free up this space later, you can
+ * delete the whole project (top right menu at wandb.ai / USERNAME / PROJECT_NAME /overview), or
+ * delete individual artifacts (hover on the three vertical dots to the right of the artifact name in the sidebar), or
+ * delete specific artifact versions through the storage explorer at wandb.ai / storage / USERNAME /PROJECT_NAME.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Download sample data: Nature photos
+
+ Note: **this stage might take a few minutes (~3.6GB of data)**. If you end up needing to rerun this cell, comment out the first capture line (change ```%%capture``` to ```#%%capture``` ) so you can respond to the prompt about re-downloading the dataset (and see the progress bar).
+
+ Download subsampled data: 10,000 training images and 2,000 validation images from the [iNaturalist dataset](https://github.com/visipedia/inat_comp), evenly distributed across 10 classes of living things like birds, insects, plants, and mammals (names given in Latin—so Aves, Insecta, Plantae, etc :). We will fine-tune a convolutional neural network already trained on ImageNet on this task: given a photo of a living thing, correctly classify it into one of the 10 classes.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !curl -SL https://storage.googleapis.com/wandb_datasets/nature_12K.zip > nature_12K.zip
+ # !unzip nature_12K.zip
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Start out by installing the experiment tracking library and setting up your free W&B account:
+
+ * **pip install wandb** – Install the W&B library
+ * **import wandb** – Import the wandb library
+ * **wandb login** – Login to your W&B account so you can log all your metrics in one place
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qq
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import os
+ from random import shuffle
+
+ # source directory for all raw data (technically a subset of only 10K images)
+ SRC = "inaturalist_12K/train"
+
+ # number of images per class label
+ # The total number of images is
+ # 10 classes * 1000 images = 10,000 images in SRC
+ NUM_IMAGES = 1000 # per class label, set this lower for faster results/fewer files
+ PROJECT_NAME = "artifacts_demo"
+ PREFIX = "inat" # convenient for tracking local data
+ return NUM_IMAGES, PREFIX, PROJECT_NAME, SRC, os, shuffle
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 1: Upload raw data
+ """)
+ return
+
+
+@app.cell
+def _(NUM_IMAGES, PREFIX, PROJECT_NAME, SRC, os, shuffle, wandb):
+ _RAW_DATA_AT = '_'.join([PREFIX, 'raw_data_10K'])
+ _run = wandb.init(project=PROJECT_NAME, job_type='upload')
+ raw_data_at = wandb.Artifact(_RAW_DATA_AT, type='raw_data')
+ # create an artifact for all the raw data
+ _labels = os.listdir(SRC)
+ for _l in _labels:
+ # SRC_DIR contains 10 folders, one for each of 10 class labels
+ # each folder contains images of the corresponding class
+ _imgs_per_label = os.path.join(SRC, _l)
+ if os.path.isdir(_imgs_per_label):
+ _imgs = os.listdir(_imgs_per_label)
+ shuffle(_imgs)
+ img_file_ids = _imgs[:NUM_IMAGES]
+ for f in img_file_ids: # randomize the order
+ file_path = os.path.join(SRC, _l, f)
+ raw_data_at.add_file(file_path, name=_l + '/' + f)
+ _run.log_artifact(raw_data_at)
+ # save artifact to W&B
+ _run.finish() # add file to artifact by full path
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 2: Split raw data to prepare for training
+ """)
+ return
+
+
+@app.cell
+def _(PREFIX, PROJECT_NAME, os, shuffle, wandb):
+ _RAW_DATA_AT = 'inat_raw_data_10K'
+ _run = wandb.init(project=PROJECT_NAME, job_type='data_split')
+ data_at = _run.use_artifact(_RAW_DATA_AT + ':latest')
+ # find the most recent ("latest") version of the full raw data
+ # you can of course pass around programmatic aliases and not string literals
+ data_dir = data_at.download()
+ # download it locally (for illustration purposes/across hardware; you can
+ # also sync/version artifacts by reference)
+ DATA_SPLITS = {'train': 800, 'val': 100, 'test': 100}
+ ats = {}
+ # create balanced train, val, test splits
+ # each count is the number of images per label
+ for split, count in DATA_SPLITS.items():
+ ats[split] = wandb.Artifact('_'.join([PREFIX, split, 'data', str(count * 10)]), '_'.join([split, 'data']))
+ _labels = os.listdir(data_dir)
+ # wrap artifacts in dictionary for convenience
+ for _l in _labels:
+ if _l.startswith('.'):
+ continue
+ _imgs_per_label = os.listdir(os.path.join(data_dir, _l))
+ shuffle(_imgs_per_label)
+ start_id = 0
+ for split, count in DATA_SPLITS.items(): # skip non-label file
+ split_imgs = _imgs_per_label[start_id:start_id + count]
+ for img_file in split_imgs:
+ full_path = os.path.join(data_dir, _l, img_file)
+ ats[split].add_file(full_path, name=os.path.join(_l, img_file))
+ start_id += count
+ for split, artifact in ats.items(): # take a subset
+ _run.log_artifact(artifact)
+ # save all three artifacts to W&B
+ # note: yes, in this example, we are cheating and have labels for the "test" data ;)
+ _run.finish() # add file to artifact by full path # note: pass the label to the name parameter to retain it in # the data structure
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 3: Train with artifacts and save model
+ """)
+ return
+
+
+@app.cell
+def _(PREFIX, PROJECT_NAME, os, wandb):
+ # EXPERIMENT CONFIG
+ #---------------------------
+ # number of training and validation examples
+ # set these lower for fewer files/faster results
+ # if you set these higher, make sure the total count is less than or equal to
+ # the number of files uploaded for that split in the train/val data artifact
+ NUM_TRAIN = 800 # try 500, 1000, 2000, or max 10000
+ NUM_VAL = 100
+ NUM_EPOCHS = 1 # set low for demo purposes; try 3, 5, or as many as you like
+ MODEL_NAME = 'iv3_trained'
+ # model name
+ # if you want to train a sufficiently different model, give this a new name
+ # to start a new lineage for the model, instead of just incrementing the
+ # version of the old model
+ INIT_MODEL_DIR = 'init_model_keras_iv3.keras'
+ FINAL_MODEL_DIR = 'trained_keras_model_iv3.keras'
+ # folder in which to save initial, untrained model
+ import numpy as np
+ from sklearn.metrics import precision_recall_curve, roc_curve
+ # folder in which to save the final, trained model
+ from sklearn.metrics import average_precision_score
+ from sklearn.preprocessing import label_binarize
+ from tensorflow.keras.applications.inception_v3 import InceptionV3
+ from tensorflow.keras.callbacks import Callback
+ from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
+ from tensorflow.keras.models import Model
+ from tensorflow.keras.preprocessing.image import ImageDataGenerator
+ from wandb.integration.keras import WandbMetricsLogger, WandbModelCheckpoint
+ config_defaults = {'num_train': NUM_TRAIN, 'num_val': NUM_VAL, 'num_classes': 10, 'fc_size': 1024, 'img_width': 299, 'img_height': 299, 'batch_size': 32, 'epochs': NUM_EPOCHS}
+
+ def finetune_inception_model(fc_size, num_classes):
+ """Load InceptionV3 with ImageNet weights, freeze it,
+ and attach a finetuning top for this classification task"""
+ base = InceptionV3(weights='imagenet', include_top='False')
+ for layer in base.layers:
+ layer.trainable = False
+ x = base.get_layer('mixed10').output
+ x = GlobalAveragePooling2D()(x)
+ x = Dense(fc_size, activation='relu')(x)
+ guesses = Dense(num_classes, activation='softmax')(x)
+ # experiment configuration saved to W&B
+ model = Model(inputs=base.input, outputs=guesses)
+ model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])
+ return model
+
+ def train():
+ """ Main training loop. This is called pretrain because it freezes
+ the InceptionV3 layers of the model and only trains the new top layers # inceptionV3 settings
+ on the new data. subsequent training phase would unfreeze all the layers
+ and finetune the whole model on the new data"""
+ _run = wandb.init(project=PROJECT_NAME, job_type='train', config=config_defaults)
+ cfg = wandb.config
+ train_at = os.path.join(PROJECT_NAME, PREFIX + '_train_data_8000') + ':latest'
+ val_at = os.path.join(PROJECT_NAME, PREFIX + '_val_data_1000') + ':latest'
+ train_data = _run.use_artifact(train_at, type='train_data')
+ train_dir = train_data.download()
+ val_data = _run.use_artifact(val_at, type='val_data')
+ val_dir = val_data.download()
+ train_datagen = ImageDataGenerator(rescale=1.0 / 255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True) # load InceptionV3 as base
+ val_datagen = ImageDataGenerator(rescale=1.0 / 255)
+ train_generator = train_datagen.flow_from_directory(train_dir, target_size=(cfg.img_width, cfg.img_height), batch_size=cfg.batch_size, class_mode='categorical') # freeze base layers
+ val_generator = val_datagen.flow_from_directory(val_dir, target_size=(cfg.img_width, cfg.img_height), batch_size=cfg.batch_size, class_mode='categorical')
+ model = finetune_inception_model(cfg.fc_size, cfg.num_classes)
+ model_artifact = wandb.Artifact('iv3', type='model', description='unmodified inception v3', metadata=dict(cfg))
+ model.save(INIT_MODEL_DIR)
+ model_artifact.add_file(INIT_MODEL_DIR) # attach a fine-tuning layer
+ _run.log_artifact(model_artifact)
+ callbacks = [WandbMetricsLogger(), WandbModelCheckpoint('checkpoint.keras')]
+ model.fit(train_generator, steps_per_epoch=cfg.num_train // cfg.batch_size, epochs=cfg.epochs, validation_data=val_generator, callbacks=callbacks, validation_steps=cfg.num_val // cfg.batch_size)
+ trained_model_artifact = wandb.Artifact(MODEL_NAME, type='model', description='trained inception v3', metadata=dict(cfg))
+ model.save(FINAL_MODEL_DIR)
+ trained_model_artifact.add_file(FINAL_MODEL_DIR)
+ _run.log_artifact(trained_model_artifact)
+ _run.finish() # track this experiment with wandb: all runs will be sent # to the given project name # artifact names # create train and validation data generators # instantiate model and callbacks # log model # train! # save trained model as artifact
+
+ return MODEL_NAME, np, train
+
+
+@app.cell
+def _(train):
+ train()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 4: Load model for inference
+ """)
+ return
+
+
+@app.cell
+def _(MODEL_NAME, PROJECT_NAME, np, os, wandb):
+ from tensorflow.keras.preprocessing import image
+ from tensorflow import keras
+ import pathlib
+ _run = wandb.init(project=PROJECT_NAME, job_type='inference')
+ model_at = _run.use_artifact(MODEL_NAME + ':latest')
+ model_dir = pathlib.Path(model_at.download())
+ print('model: ', model_dir)
+ model_path = model_dir / 'trained_keras_model_iv3.keras'
+ model = keras.models.load_model(model_path, compile=False)
+ test_data_at = _run.use_artifact('inat_test_data_1000:latest')
+ test_dir = test_data_at.download()
+ _imgs = []
+ class_labels = os.listdir(test_dir)
+ for _l in class_labels:
+ if _l.startswith('.'):
+ continue
+ imgs_per_class = os.listdir(os.path.join(test_dir, _l))
+ for img in imgs_per_class:
+ img_path = os.path.join(test_dir, _l, img)
+ img = image.load_img(img_path, target_size=(299, 299))
+ img = image.img_to_array(img)
+ img = np.expand_dims(img / 255.0, axis=0)
+ _imgs.append(img)
+ preds = {}
+ _imgs = np.vstack(_imgs)
+ classes = model.predict(_imgs, batch_size=32)
+ for c in classes:
+ class_id = np.argmax(c)
+ preds[class_id] = preds.get(class_id, 0) + 1
+ print(preds)
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # More about Weights & Biases
+ We're always free for academics and open source projects. Here are some more resources:
+
+ 1. [Documentation](http://docs.wandb.com) - Python docs
+ 2. [Gallery](https://app.wandb.ai/gallery) - example reports in W&B
+ 3. [Articles](https://www.wandb.com/articles) - blog posts and tutorials
+ 4. [Community](wandb.me/slack) - join our Slack community forum
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-model-management-guide/wandb_artifacts_model_management_guide.py b/marimo/convert/wandb-artifacts-model-management-guide/wandb_artifacts_model_management_guide.py
new file mode 100644
index 00000000..98411aa8
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-model-management-guide/wandb_artifacts_model_management_guide.py
@@ -0,0 +1,275 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Model Management Guide Companion Notebook
+
+ This is a companion notebook to the [W&B Model Management Guide](https://docs.wandb.ai/guides/models).
+
+ **Table of Contents**
+ * **Cell 1**: Installs `wandb` python library
+ * **Cell 2** (Form): Allows you to specify some parameters and defines a handful of helper functions. Note: there is not any `wandb` specific library calls in these helper functions - they are purely used to allow the example cells to be more terse and focus on the key aspects of Model Management
+ * **Cell 3**: (Train, Log, & Link Models) Covers steps [2. Traing & log a Model](https://docs.wandb.ai/guides/models#2.-train-and-log-model-versions) and [3. Link Model Versions to the Collection](https://docs.wandb.ai/guides/models#3.-link-model-versions-to-the-portfolio)
+ * **Cell 4**: (Use, Evaluate, and Promote a Model) Covers steps [4. Using a Model Version](https://docs.wandb.ai/guides/models#4.-use-a-model-version), [5. Evaluate Model Performance](https://docs.wandb.ai/guides/models#5.-evaluate-model-performance), and [6. Promote a Version to Production](https://docs.wandb.ai/guides/models#6.-promote-a-version-to-production)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ **Stop! 🛑** Please complete [Step 1 of the tutorial](https://docs.wandb.ai/guides/models/walkthrough#1-create-a-new-registered-model) before continuing. This will ensure you have a **Model Collection** defined in your project. Enter the Project name where you created the Collection in `project_name` and the name of the Collection in the `model_collection_name` fields respectively.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -U -qqq
+ return
+
+
+@app.cell
+def _():
+ project_name = 'quickstart-model-registry' #@param {type:"string"}
+ # dataset_name = "mnist" #@param {type:"string"}
+ dataset_name = 'mnist'
+ model_collection_name = 'MNIST Grayscale 28x28' #@param {type:"string"}
+ import wandb
+ import torch
+ from torchvision import datasets, transforms
+ import torch.nn as nn
+ import torch.nn.functional as F
+ import torch.optim as optim
+ from torch.utils.data import DataLoader
+ from torch.optim.lr_scheduler import OneCycleLR
+
+ class Net(nn.Module):
+
+ def __init__(self):
+ super(Net, self).__init__()
+ self.conv1 = nn.Conv2d(1, 32, 3, 1)
+ self.conv2 = nn.Conv2d(32, 64, 3, 1)
+ self.dropout1 = nn.Dropout(0.25)
+ self.dropout2 = nn.Dropout(0.5)
+ self.fc1 = nn.Linear(9216, 128)
+ self.fc2 = nn.Linear(128, 10)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = F.relu(x)
+ x = self.conv2(x)
+ x = F.relu(x)
+ x = F.max_pool2d(x, 2)
+ x = self.dropout1(x)
+ x = torch.flatten(x, 1)
+ x = self.fc1(x)
+ x = F.relu(x)
+ x = self.dropout2(x)
+ output = self.fc2(x)
+ return output
+
+ def _sample_mnist(split0, split1, is_train=True):
+ """Sample MNIST dataset"""
+ mnist_data = datasets.MNIST('train_data/' if is_train else 'test_data/', download=True, train=is_train, transform=transforms.Compose([transforms.ToTensor()]))
+ extra = len(mnist_data) - split0 - split1
+ assert extra >= 0
+ splits = torch.utils.data.random_split(mnist_data, [split0, split1, extra])
+ return (splits[0], splits[1])
+
+ def build_train_data(train_size, val_size, batch_size=128):
+ splits = _sample_mnist(train_size, val_size)
+ return (DataLoader(splits[0], batch_size=batch_size), DataLoader(splits[1], batch_size=val_size))
+
+ def build_test_data(test_size):
+ splits = _sample_mnist(test_size, 0, is_train=False)
+ return DataLoader(splits[0], batch_size=test_size)
+
+ def build_model(learning_rate, total_steps):
+ device = torch.device('cpu')
+ _model = Net().to(device)
+ optimizer = optim.Adam(_model.parameters())
+ scheduler = OneCycleLR(optimizer, max_lr=learning_rate, total_steps=total_steps)
+ return (_model, optimizer, scheduler)
+
+ def train_step(model, optimizer, scheduler, batch_x, batch_y):
+ _model.train()
+ batch_x, batch_y = (batch_x.to('cpu'), batch_y.to('cpu'))
+ optimizer.zero_grad()
+ preds = _model(batch_x)
+ loss = F.cross_entropy(preds, batch_y)
+ loss.backward()
+ optimizer.step()
+ scheduler.step()
+ return (loss.item(), preds)
+
+ @torch.no_grad()
+ def evaluate_model(model, eval_dl):
+ device = torch.device('cpu')
+ _model.eval()
+ test_loss = 0
+ correct = 0
+ preds = []
+ for data, target in eval_dl:
+ data, target = (data.to(device), target.to(device))
+ output = _model(data)
+ test_loss += F.cross_entropy(output, target, reduction='sum').item()
+ pred = output.argmax(dim=1, keepdim=True)
+ preds += list(pred.flatten().tolist())
+ correct += pred.eq(target.view_as(pred)).sum().item()
+ test_loss /= len(eval_dl.dataset)
+ accuracy = 100.0 * correct / len(eval_dl.dataset)
+ return (test_loss, accuracy, preds) # sum up batch loss # get the index of the max log-probability
+
+ return (
+ Net,
+ build_model,
+ build_test_data,
+ build_train_data,
+ dataset_name,
+ evaluate_model,
+ model_collection_name,
+ project_name,
+ torch,
+ train_step,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train, Log, & Link Models
+ """)
+ return
+
+
+@app.cell
+def _(evaluate_model, torch, train_step, wandb):
+ def save_model(model, is_best=False):
+ """Save model to W&B and locally if it's the best one so far."""
+ _art = wandb.Artifact(f'mnist-{wandb.run.id}', 'model') ##### W&B MODEL MANAGEMENT SPECIFIC CALLS ######
+ torch.save(_model.state_dict(), 'model.pt')
+ _art.add_file('model.pt')
+ wandb.log_artifact(_art, aliases=['best', 'latest'] if is_best else None)
+ return _art
+
+ def train_model(model, optimizer, scheduler, train_loader, val_loader, num_epochs=5):
+ """A simple training loop"""
+ best_val_loss = 10000000000.0
+ best_model_art = None
+ for epoch in range(num_epochs):
+ for batch_x, batch_y in train_loader:
+ train_loss, _ = train_step(_model, optimizer, scheduler, batch_x, batch_y)
+ wandb.log({'epoch': epoch, 'train_loss': train_loss, 'learning_rate': optimizer.param_groups[0]['lr']})
+ val_loss, val_acc, _ = evaluate_model(_model, val_loader)
+ wandb.log({'val_loss': val_loss, 'val_acc': val_acc})
+ best_val_loss = min(best_val_loss, val_loss)
+ model_art = save_model(_model, is_best=val_loss <= best_val_loss)
+ if val_loss <= best_val_loss:
+ best_model_art = model_art
+ print('New best model saved!')
+ print(f'Epoch {epoch}: val_loss: {val_loss}, val_acc: {val_acc}')
+ return best_model_art
+
+ return (train_model,)
+
+
+@app.cell
+def _(
+ build_model,
+ build_train_data,
+ dataset_name,
+ model_collection_name,
+ project_name,
+ train_model,
+ wandb,
+):
+ # Startup a W&B Run
+ wandb.init(project=project_name, job_type='model_trainer', config={'train_size': 2000, 'val_size': 200, 'batch_size': 64, 'learning_rate': 0.001, 'epochs': 5})
+ _config = wandb.config
+ train_dl, val_dl = build_train_data(_config.train_size, _config.val_size, _config.batch_size)
+ _art = wandb.Artifact(f'{dataset_name}-train', 'dataset')
+ _art.add_dir('./train_data')
+ wandb.use_artifact(_art)
+ _model, optimizer, scheduler = build_model(_config.learning_rate, total_steps=len(train_dl) * _config.epochs)
+ best_model_art = train_model(_model, optimizer, scheduler, train_dl, val_dl, _config.epochs)
+ wandb.run.link_artifact(best_model_art, model_collection_name, ['latest'])
+ # Load in the training data
+ # (Optional) Declare dataset dependency
+ # Define a model
+ # Train the Model
+ ##### W&B MODEL MANAGEMENT SPECIFIC CALLS ######
+ # Finish the Run
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Use, Evaluate, and Promote a Model
+ """)
+ return
+
+
+@app.cell
+def _(
+ Net,
+ build_test_data,
+ dataset_name,
+ evaluate_model,
+ model_collection_name,
+ project_name,
+ torch,
+ wandb,
+):
+ wandb.init(project=project_name, job_type='model_evaluator', config={'test_size': 100})
+ _config = wandb.config
+ test_dl = build_test_data(test_size=_config.test_size)
+ _art = wandb.Artifact(f'{dataset_name}-test', 'dataset')
+ _art.add_dir('./test_data')
+ wandb.use_artifact(_art)
+ model_art = wandb.use_artifact(f'{model_collection_name}:latest')
+ model_path = model_art.get_path('model.pt').download()
+ _model = Net().cpu()
+ checkpt = torch.load(model_path)
+ _model.load_state_dict(checkpt)
+ val_loss, val_acc, preds = evaluate_model(_model, test_dl)
+ table = wandb.Table(data=[], columns=[])
+ table.add_column('image', [wandb.Image(i.numpy()) for i in list(test_dl)[0][0]])
+ table.add_column('label', list(test_dl)[0][1].tolist())
+ table.add_column('pred', preds)
+ wandb.log({'test_loss': val_loss, 'test_acc': val_acc, 'predictions': table})
+ wandb.run.link_artifact(model_art, model_collection_name, ['latest', 'production'])
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-train-val-test-split-with-tabular-data/wandb_artifacts_train_val_test_split_with_tabular_data.py b/marimo/convert/wandb-artifacts-train-val-test-split-with-tabular-data/wandb_artifacts_train_val_test_split_with_tabular_data.py
new file mode 100644
index 00000000..4ab8af20
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-train-val-test-split-with-tabular-data/wandb_artifacts_train_val_test_split_with_tabular_data.py
@@ -0,0 +1,566 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ # Tabular Data Versioning and Deduplication with Weights & Biases
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Introduction
+ This walkthrough focuses on using W&B to version control and iterate on tabular data. We will use [Artifacts](https://docs.wandb.ai/guides/artifacts) and [Tables](https://docs.wandb.ai/guides/data-vis/tables-quickstart) to load in a dataset and split it into train, validation, and test subsets. Thanks to the versioning capability of Artifacts, we will use minimal storage space and have persistent version labels to easily share dataset iterations with colleagues.
+
+ For this project, we will be working with tabular medical data that has great potential for predicting outcomes of heart attack patients. If you'd rather learn about similar features applied to classification tasks on image data, see this [other example](https://wandb.ai/stacey/mendeleev/reports/Tables-Tutorial-Visualize-Data-for-Image-Classification--VmlldzozNjE3NjA).
+
+ You can also find an overview [report on this topic here](https://wandb.ai/dpaiton/splitting-tabular-data/reports/Dataset-Version-Control-and-Deduplication-with-Tabular-Data-with-W-B-Artifacts-and-Tables--VmlldzoxNDIzOTA1).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ #### About the data
+ Our dataset is a collection of measurements and survey answers from hospital patients that have heart attack-related symptoms. It's not a stretch to say datasets like these -- and the models built from them -- can improve patient outcomes and save lives. On a purely machine learning level, this particular dataset is interesting because it has:
+ * **Mixed types**: entries can be binary, ordinal, numeric, or categorical
+ * **Missing data**: almost all features have some fraction of missing data
+ * **Real-world complexity**: feature values are not uniformly distributed and have outliers
+ * **Limited size**: collecting data is difficult and so the dataset is small
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Helper functions
+ Below we will keep all of our imports and helper functions. Reading through them is optional, and only recommended after you have looked over the rest of the report.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install --upgrade wandb -qqq
+ import wandb
+
+ import random
+ from collections import OrderedDict
+ import json
+ import requests
+ import csv
+ import os
+
+ import numpy as np
+ import torch
+ from torch.utils.data import Dataset, DataLoader
+ from torchvision import transforms
+
+ DEVICE = 'cpu'
+ PROJECT_NAME = 'splitting-tabular-data'
+
+ # Set the random seeds to improve reproducibility by removing stochasticity
+ def set_seeds(seed):
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed(seed)
+ torch.backends.cudnn.benchmark = False # Force cuDNN to use a consistent convolution algorithm
+ torch.backends.cudnn.deterministic = True # Force cuDNN to use deterministic algorithms if available
+ torch.use_deterministic_algorithms(True) # Force torch to use deterministic algorithms if available
+
+ set_seeds(0)
+ return (
+ DEVICE,
+ DataLoader,
+ Dataset,
+ OrderedDict,
+ PROJECT_NAME,
+ csv,
+ json,
+ np,
+ os,
+ requests,
+ torch,
+ transforms,
+ wandb,
+ )
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _(
+ DataLoader,
+ Dataset,
+ OrderedDict,
+ PROJECT_NAME,
+ json,
+ np,
+ os,
+ torch,
+ transforms,
+ wandb,
+):
+ def make_split_artifact(run, raw_data_table, train_rows, val_rows, test_rows):
+ """
+ Creates a w&b artifact that contains a singular reference table (aka a ForeignIndex table).
+ The ForeignIndex table has a single column that we are naming 'source'.
+ It contains references to the original table (raw_data_table) for each of the splits.
+ Arguments:
+ run (wandb run) returned from wandb.init()
+ raw_data_table (wandb Table) that contains your original tabular data
+ train_rows (list of ints) indices that reference the training rows in the raw_data_table
+ val_rows (list of ints) indices that reference the validation rows in the raw_data_table
+ test_rows (list of ints) indices that reference the test rows in the raw_data_table
+ """
+ split_artifact = wandb.Artifact('data-splits', type='dataset', description='Train, validation, test dataset splits')
+ data_table_pointer = raw_data_table.get_index()
+ split_artifact.add(wandb.Table(columns=['source'], data=[[data_table_pointer[i]] for i in train_rows]), 'train-data')
+ split_artifact.add(wandb.Table(columns=['source'], data=[[data_table_pointer[i]] for i in val_rows]), 'val-data')
+ split_artifact.add(wandb.Table(columns=['source'], data=[[data_table_pointer[i]] for i in test_rows]), 'test-data') # Our data split artifact will only store index references to the original dataset table to save space
+ _run.log_artifact(split_artifact) # ForeignIndex automatically references the source table
+
+ def make_loaders(config):
+ """
+ Makes data loaders using a artifact containing the dataset splits (created using the make_split_artifact() function)
+ The function assumes that you have created a data-splits artifact and a data-transforms artifact
+ Arguments:
+ config [dict] containing keys:
+ data_columns (list of ints) referencing which columns are to be treated as data
+ label_columns (list of ints) referencing which columns are to be treated as labels
+ num_classes (int) number of possible label classes in the dataset
+ batch_size (int) amount of rows (i.e. data instances) to be delivered in a single batch
+ Returns:
+ train_loader (PyTorch DataLoader) containing the training data
+ val_loader (PyTorch DataLoader) containing the validation data
+ test_loader (PyTorch DataLoader) containing the test data
+ """
+ with wandb.init(project=PROJECT_NAME, job_type='package-data', config=config) as _run:
+ transform_dir = _run.use_artifact('data-transforms:latest').download()
+ transform_dict = json.load(open(os.path.join(transform_dir, 'transforms.txt')), object_pairs_hook=OrderedDict)
+ composed_transforms = get_transforms(transform_dict)
+ split_artifact = _run.use_artifact('data-splits:latest')
+ train_loader = DataLoader(MyocardialInfarctionDataset(split_artifact.get('train-data'), config['data_columns'], config['label_columns'], config['num_classes'], composed_transforms), batch_size=config['batch_size'], drop_last=True, shuffle=True, num_workers=0)
+ val_loader = DataLoader(MyocardialInfarctionDataset(split_artifact.get('val-data'), config['data_columns'], config['label_columns'], config['num_classes'], composed_transforms), batch_size=config['batch_size'], batch_sampler=None, shuffle=False, num_workers=0)
+ test_loader = DataLoader(MyocardialInfarctionDataset(split_artifact.get('test-data'), config['data_columns'], config['label_columns'], config['num_classes'], composed_transforms), batch_size=config['batch_size'], batch_sampler=None, shuffle=False, num_workers=0)
+ return (train_loader, val_loader, test_loader)
+
+ def get_table_row(table, ndx):
+ """
+ Given a table and index, return the corresponding row # Load the transforms
+ Arguments:
+ table (wandb.Table) can be a standard table of data or a pointer to a reference table
+ ndx (int) row index to slice
+ Returns:
+ ref_row (list) of data entries for the row referenced by ndx # Reformat data to (inputs, labels)
+ """
+ linked_table = np.all([type(value) is wandb.data_types._ForeignIndexType for value in table._column_types.params['type_map'].values()])
+ if linked_table:
+ ref_table = table.get_column(table.columns[0])
+ if type(ndx) is list:
+ ref_row = [list(ref_table[i].get_row().values()) for i in ndx]
+ elif type(ndx) is int:
+ ref_row = list(ref_table[ndx].get_row().values())
+ else:
+ raise ValueError(f'Input argument ndx must be of type int or list, not {type(ndx)}')
+ return ref_row
+ else:
+ return table.data[ndx]
+
+ class MyocardialInfarctionDataset(Dataset):
+ """
+ Myocardial Infarction Dataset
+ In general columns 2-112 can be used as input data for prediction.
+ Possible complications (outputs) are listed in columns 113-124.
+
+ There are four possible time moments for complication prediction: on base of the information known at
+ 1. The time of admission to hospital: all input columns (2-112) except 93, 94, 95, 100, 101, 102, 103, 104, 105 can be used for prediction;
+ 2. The end of the first day (24 hours after admission to the hospital): all input columns (2-112) except 94, 95, 101, 102, 104, 105 can be used for prediction;
+ 3. The end of the second day (48 hours after admission to the hospital) all input columns (2-112) except 95, 102, 105 can be used for prediction;
+ 4. The end of the third day (72 hours after admission to the hospital) all input columns (2-112) can be used for prediction.
+
+ All of the above column numbers are 1-indexed.
+ """
+
+ def __init__(self, table, data_columns, label_columns, num_classes, transform=None):
+ """
+ Args:
+ table (wandb.Table): table containing the dataset
+ data_columns (list): list of column indices corresponding to the data (X)
+ label_columns (list): list of column indices corresponding to the labels (Y)
+ num_classes (int): number of possible output classes (for one-hot encoding)
+ transform (function): receives (data, label) tuple as input and produces transformed (data, label) tuple as output
+ """
+ super(MyocardialInfarctionDataset, self).__init__()
+ self.table = table
+ self.data_columns = data_columns # Check if the table's contents are pointers to another table or not
+ self.label_columns = label_columns
+ self.num_classes = num_classes
+ self.transform = transform
+
+ def __len__(self): # The table entries reference another table
+ return len(self.table.data) # There should only be one reference column
+ # The pointers are dereferenced using the get_row() function
+ def __getitem__(self, idx):
+ if torch.is_tensor(idx):
+ idx = idx.tolist()
+ label_row = np.array(get_table_row(self.table, idx), dtype=np.float32).take(self.label_columns)
+ data_row = np.array(get_table_row(self.table, idx), dtype=np.float32).take(self.data_columns)
+ if self.transform:
+ data_row, label_row = self.transform((data_row, label_row))
+ return (data_row, label_row)
+
+ class NoneToVal(object): # Standard w&b Table containing the data
+ """Convert None or NaN entries to usable values
+ """
+
+ def __init__(self, fill_value):
+ self.fill_value = fill_value
+
+ def __call__(self, data_tuple):
+ data, label = data_tuple
+ data = np.ma.masked_invalid(data).filled(fill_value=self.fill_value)
+ return (data, label)
+
+ class ToTensor(object):
+ """Convert numpy arrays to tensor arrays
+ """
+
+ def __init__(self, device=None):
+ if device is None:
+ device = 'cpu'
+ self.device = device
+
+ def __call__(self, data_tuple):
+ data, labels = data_tuple
+ return (torch.from_numpy(data).to(self.device), torch.from_numpy(labels).to(self.device))
+
+ class OneHot(object):
+ """Convert input tensor to one-hot array
+ """
+
+ def __init__(self, num_classes):
+ self.num_classes = int(num_classes)
+
+ def __call__(self, data_tuple):
+ data, labels = data_tuple
+ device = labels.device
+ dtype = labels.dtype
+ num_datapoints = int(labels.ndim)
+ labels_one_hot = torch.zeros((num_datapoints, self.num_classes), dtype=dtype).to(device)
+ labels_one_hot[:, labels.long()] = 1
+ return (data, labels_one_hot.squeeze())
+
+ def get_transforms(transform_dict):
+ """
+ Given a dictionary of transform parameters, return a list of class instances for each transform
+ Arguments:
+ transform_dict (OrderedDict) with optional keys:
+ NoneToVal (dict) if present, requires the 'value' key that None/nan will be replaced with
+ ToTensor (dict) if present, requires the 'device' key that indicates the PyTorch device
+ OneHot (dict) if present, requires the 'num_classes' key that has an int value for the number of possible data labels
+ Returns:
+ composed_transforms (PyTorch composed transform class) containing the requested transform steps in order
+ """
+ transform_functions = []
+ for key in transform_dict.keys():
+ if key == 'NoneToVal':
+ transform_functions.append(NoneToVal(transform_dict[key]['value']))
+ elif key == 'ToTensor':
+ transform_functions.append(ToTensor(transform_dict[key]['device']))
+ elif key == 'OneHot':
+ transform_functions.append(OneHot(transform_dict[key]['num_classes']))
+ composed_transforms = transforms.Compose(transform_functions)
+ return composed_transforms # Replace None/nan entries with a given value or array of values # Convert array to a PyTorch Tensor # Convert class labels to a one-hot representation
+
+ return make_loaders, make_split_artifact
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The following list contains the data types for each feature in the Myocardial Infarction dataset.
+
+ More detailed information can be found via this [pdf file](https://s3-eu-west-1.amazonaws.com/pstorage-leicester-213265548798/22803695/Descriptivestatistics.pdf)
+ """)
+ return
+
+
+@app.cell
+def _():
+ data_types = [
+ int, # 001. ID; numeric, [0, --]
+ int, # 002. AGE; numeric, [26, 92]
+ int, # 003. SEX; binary
+ int, # 004. INF_ANAM; ordinal, 4 levels
+ int, # 005. STENOK_AN; ordinal, 7 levels
+ int, # 006. FK_STENOK; ordinal, 5 levels
+ int, # 007. IBS_POST; ordinal, 3 levels
+ int, # 008. IBS_NASL; binary
+ int, # 009. GB; ordinal, 4 levels
+ int, # 010. SIM_GIPERT; binary
+ int, # 011. DLT_AG; ordinal, 8 levels
+ int, # 012. ZSN_A; partially ordered, 5 levels
+ int, # 013. nr11; binary
+ int, # 014. nr01; binary
+ int, # 015. nr02; binary
+ int, # 016. nr03; binary
+ int, # 017. nr04; binary
+ int, # 018. nr07; binary
+ int, # 019. nr08; binary
+ int, # 020. np01; binary
+ int, # 021. np04; binary
+ int, # 022. np05; binary
+ int, # 023. np07; binary
+ int, # 024. np08; binary
+ int, # 025. np09; binary
+ int, # 026. np10; binary
+ int, # 027. endocr_01; binary
+ int, # 028. endocr_02; binary
+ int, # 029. endocr_03; binary
+ int, # 030. zab_leg_01; binary
+ int, # 031. zab_leg_02; binary
+ int, # 032. zab_leg_03; binary
+ int, # 033. zab_leg_04; binary
+ int, # 034. zab_leg_06; binary
+ float, # 035. S_AD_KBRIG; numeric, [0, 260] mmHg
+ float, # 036. D_AD_KBRIG; numeric, [0, 190] mmHg
+ float, # 037. S_AD_ORIT; numeric, [0, 260] mmHg
+ float, # 038. D_AD_ORIT; numeric, [0, 190] mmHg
+ int, # 039. O_L_POST; binary
+ int, # 040. K_SH_POST; binary
+ int, # 041. MP_TP_POST; binary
+ int, # 042. SVT_POST; binary
+ int, # 043. GT_POST; binary
+ int, # 044. FIB_G_POST; binary
+ int, # 045. ant_im; ordinal, 5 levels
+ int, # 046. lat_im; ordinal, 5 levels
+ int, # 047. inf_im; ordinal, 5 levels
+ int, # 048. post_im; ordinal, 5 levels
+ int, # 049. IM_PG_P; binary
+ int, # 050. ritm_ecg_p_01; binary
+ int, # 051. ritm_ecg_p_02; binary
+ int, # 052. ritm_ecg_p_04; binary
+ int, # 053. ritm_ecg_p_06; binary
+ int, # 054. ritm_ecg_p_07; binary
+ int, # 055. ritm_ecg_p_08; binary
+ int, # 056. n_r_ecg_p_01; binary
+ int, # 057. n_r_ecg_p_02; binary
+ int, # 058. n_r_ecg_p_03; binary
+ int, # 059. n_r_ecg_p_04; binary
+ int, # 060. n_r_ecg_p_05; binary
+ int, # 061. n_r_ecg_p_06; binary
+ int, # 062. n_r_ecg_p_08; binary
+ int, # 063. n_r_ecg_p_09; binary
+ int, # 064. n_r_ecg_p_10; binary
+ int, # 065. n_p_ecg_p_01; binary
+ int, # 066. n_p_ecg_p_03; binary
+ int, # 067. n_p_ecg_p_04; binary
+ int, # 068. n_p_ecg_p_05; binary
+ int, # 069. n_p_ecg_p_06; binary
+ int, # 070. n_p_ecg_p_07; binary
+ int, # 071. n_p_ecg_p_08; binary
+ int, # 072. n_p_ecg_p_09; binary
+ int, # 073. n_p_ecg_p_10; binary
+ int, # 074. n_p_ecg_p_11; binary
+ int, # 075. n_p_ecg_p_12; binary
+ int, # 076. fibr_ter_01; binary
+ int, # 077. fibr_ter_02; binary
+ int, # 078. fibr_ter_03; binary
+ int, # 079. fibr_ter_05; binary
+ int, # 080. fibr_ter_06; binary
+ int, # 081. fibr_ter_07; binary
+ int, # 082. fibr_ter_08; binary
+ int, # 083. GIPO_K; binary
+ float, # 084. K_BLOOD; numeric, [2.3, 8.2] mmol/L
+ int, # 085. GIPER_Na; binary
+ float, # 086. Na_BLOOD; numeric, [117, 169] mmol/L
+ float, # 087. ALT_BLOOD; numeric, [0.03, 0.48] IU/L
+ float, # 088. AST_BLOOD; numeric, [0.04, 2.15] IU/L
+ float, # 089. KFK_BLOOD; numeric, [1.2, 3.6] IU/L
+ float, # 090. L_BLOOD; numeric, [2, 27.9] billions per liter
+ float, # 091. ROE; numeric, [1, 140] mm
+ int, # 092. TIME_B_S; ordinal, 10 levels
+ int, # 093. R_AB_1_n; ordinal, 4 levels
+ int, # 094. R_AB_2_n; ordinal, 4 levels
+ int, # 095. R_AB_3_n; ordinal, 4 levels
+ int, # 096. NA_KB; binary
+ int, # 097. NOT_NA_KB; binary
+ int, # 098. LID_KB; binary
+ int, # 099. NITR_S; binary
+ int, # 100. NA_R_1_n; ordinal, 5 levels
+ int, # 101. NA_R_2_n; ordinal, 4 levels
+ int, # 102. NA_R_3_n; ordinal, 3 levels
+ int, # 103. NOT_NA_1_n; ordinal, 5 levels
+ int, # 104. NOT_NA_2_n; ordinal, 4 levels
+ int, # 105. NOT_NA_3_n; ordinal, 3 levels
+ int, # 106. LID_S_n; binary
+ int, # 107. B_BLOCK_S_n; binary
+ int, # 108. ANT_CA_S_n; binary
+ int, # 109. GEPAR_S_n; binary
+ int, # 110. ASP_S_n; binary
+ int, # 111. TIKL_S_n; binary
+ int, # 112. TRENT_S_n; binary
+ int, # 113. FIBR_PREDS; binary
+ int, # 114. PREDS_TAH; binary
+ int, # 115. JELUD_TAH; binary
+ int, # 116. FIBR_JELUD; binary
+ int, # 117. A_V_BLOCK; binary
+ int, # 118. OTEK_LANC; binary
+ int, # 119. RAZRIV; binary
+ int, # 120. DRESSLER, binary
+ int, # 121. ZSN; binary
+ int, # 122. REC_IM; binary
+ int, # 123. P_IM_STEN; binary
+ int, # 124. LET_IS; categorical, 8 categories
+ ]
+ return (data_types,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Loading the dataset into an artifact
+ Our first step is to load in the dataset from a CSV file, which we accomplish with Artifacts and Tables. The wandb Artifact has two very useful features for our application: 1) it supports versioning, which will allow us to track changes we make to the original datset and 2) it supports deduplication, which will minimize the amount of storage space we use when generating modified versions of the dataset.
+
+ In the code below we use the python `requests` and `csv` libararies to load in each line of the CSV file into a wandb Table. Then we store the table in a wandb Artifact.
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT_NAME, csv, data_types, np, requests, wandb):
+ # Load raw dataset into a table and store it as an artifact
+ with wandb.init(project=PROJECT_NAME, job_type='load-data') as _run:
+ dataset_url = 'https://s3-eu-west-1.amazonaws.com/pstorage-leicester-213265548798/23581310/MyocardialinfarctioncomplicationsDatabase.csv' # Load data row-by-row & add the rows to the table (Note: the whole table will be stored in memory)
+ with requests.get(dataset_url, stream=True) as r:
+ lines = (line.decode('utf-8') for line in r.iter_lines())
+ column_headings = next(csv.reader(lines)) # Load each line into the table
+ data_table = wandb.Table(columns=column_headings)
+ for index, row in enumerate(csv.reader(lines)): # This assumes that the first CSV line contains the column headings
+ row = [data_types[entry_index](entry) if entry != '' else np.nan for entry_index, entry in enumerate(row)] # Initialize the table
+ if len(row) == len(column_headings): # Starting at the second row
+ data_table.add_data(*row)
+ dataset_artifact = wandb.Artifact('data-library', type='dataset', description='Table containing the CSV dataset', metadata={'MD5_checksum': 'd409a89bd7e566da4b82232c3956f576', 'filename': 'MyocardialinfarctioncomplicationsDatabase.csv', 'filesize': '427.31 kB', 'dataset_host': 'University of Leicester', 'dataset_url': dataset_url, 'project_url': 'https://doi.org/10.25392/leicester.data.12045261.v3', 'reference_doi': '10.25392/leicester.data.12045261.v3'})
+ dataset_artifact.add(data_table, 'data-table')
+ _run.log_artifact(dataset_artifact) # Create an artifact for our dataset # Add the table to the artifact & log the artifact
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Splitting the data into training, validation, and test sets
+ As is typical with most machine learning applications, we want to grab a majority of our data for training and use a smaller subset for validation and testing. The validation set will be used to help us tune our parameters and modify the preprocessing steps.
+
+ To avoid saving multiple copies of the dataset, we will only store corresponding indices for the train/val/test splits. This is also version controlled in case you decide you need more training data or you want to redo the shuffles.
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT_NAME, make_split_artifact, np, wandb):
+ config = {'train_val_test_split': [0.8, 0.1, 0.1], 'data_columns': [i for i in range(1, 112)], 'label_columns': [123], 'num_classes': 8, 'batch_size': 20}
+ with wandb.init(project=PROJECT_NAME, job_type='split-data', config=config) as _run: # These must sum to 1.0
+ raw_data_table = _run.use_artifact('data-library:latest').get('data-table') # All possible training features
+ num_samples = len(raw_data_table.data) # Lethal outcome
+ shuffled_rows = np.random.choice(np.arange(num_samples), num_samples, replace=False) # Num classes after preprocessing
+ train_rows, val_rows, test_rows = np.split(shuffled_rows, np.cumsum([num_samples * split for split in config['train_val_test_split'][:-1]], dtype=int)) # Num samples to average over for gradient updates
+ make_split_artifact(_run, raw_data_table, train_rows, val_rows, test_rows)
+ test_slicing = True
+ # Split data into train, val, test tables
+ if test_slicing:
+ print('num total: ', num_samples) # Define the data splits
+ print('num train:', len(train_rows))
+ print('num val:', len(val_rows)) # One of many methods to extract random train/val/test rows
+ print('num test:', len(test_rows))
+ print('shuffle duplicates: ', len(set(shuffled_rows)) != len(shuffled_rows))
+ print('val in train:', np.any([row in train_rows for row in val_rows]))
+ print('test in train:', np.any([row in train_rows for row in test_rows]))
+ print('val in test:', np.any([row in val_rows for row in test_rows])) # Construct a new artifact for the data splits
+ print('train dupliates: ', len(set(train_rows)) != len(train_rows))
+ print('val dupliates: ', len(set(val_rows)) != len(val_rows))
+ # Quick test to make sure the slicing worked properly
+ print('test dupliates: ', len(set(test_rows)) != len(test_rows))
+ print('all samples accounted for in the shuffled set:', len(shuffled_rows) == num_samples)
+ return (config,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Packaging the data into PyTorch data loaders
+ Next we will package the data into a PyTorch DataLoader to make it easier to work with. The DataLoader includes a list of preprocessing steps that are to be performed on the data. We want to be able to iterate and version control preprocessing pipeline, so we also have to write some code to store it as an artifact.
+ """)
+ return
+
+
+@app.cell
+def _(DEVICE, OrderedDict, PROJECT_NAME, config, json, make_loaders, wandb):
+ with wandb.init(project=PROJECT_NAME, job_type='define-transforms', config=config) as _run:
+ transform_dict = OrderedDict() # Define an initial set of transforms that we think will be useful
+ transform_dict['NoneToVal'] = {'value': 0}
+ transform_dict['ToTensor'] = {'device': DEVICE}
+ transform_dict['OneHot'] = {'num_classes': config['num_classes']} # for the first pass we will replace missing values with 0
+ for key_idx, key in enumerate(transform_dict.keys()):
+ transform_dict[key]['order'] = key_idx
+ data_transform_artifact = wandb.Artifact('data-transforms', type='parameters', description='Data preprocessing functions and parameters.', metadata=transform_dict)
+ with data_transform_artifact.new_file('transforms.txt') as f:
+ f.write(json.dumps(transform_dict, indent=4))
+ _run.log_artifact(data_transform_artifact)
+ config.update(transform_dict)
+ # Now we can make the data loaders with the preprocessing pipeline
+ train_loader, val_loader, test_loader = make_loaders(config) # Include an operational index to verify the order # Create an artifact for logging the transforms # Optional for viewing on the web app; the data is also stored in the txt file below # Log the transforms in JSON format # Log the transforms in the config so that we can sweep over them in future iterations
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Done! Time to train
+
+ That concludes this part of the tutorial. In a future tutorial we will use this same data to train and iterate on our model. This will also use Artifacts to version control iterations on the preprocessing pipeline, parameters, and model architecture.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-w-b-artifacts-for-auditing-purposes/wandb_artifacts_w_b_artifacts_for_auditing_purposes.py b/marimo/convert/wandb-artifacts-w-b-artifacts-for-auditing-purposes/wandb_artifacts_w_b_artifacts_for_auditing_purposes.py
new file mode 100644
index 00000000..d2c63ce1
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-w-b-artifacts-for-auditing-purposes/wandb_artifacts_w_b_artifacts_for_auditing_purposes.py
@@ -0,0 +1,545 @@
+# /// script
+# dependencies = ["awscli", "boto3", "six", "timm", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [Weights & Biases](https://wandb.ai/site) makes running collaborative machine learning projects a breeze. You can focus on what you're trying to experiment with, and W&B will take on the burden of keeping track of everything. If you want to review a loss plot, download the latest model for production, or just see which configurations produced a certain model, W&B is your friend. There's also a bunch of features to help you and your team collaborate, like having a shared dashboard and sharing interactive reports.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # How Weights and Biases can help you with Audits and Regulatory Guidelines
+
+ This notebook accompanies and implements a
+ [blog post](http://wandb.me/audit-artifacts-report)
+ on using W&B Artifacts to help teams in regulation-heavy industries share their Machine Learning models with clients.
+
+ Run the cells below to train an image classifier and upload the model checkpoints as W&B Artifacts. Then you can reliably know which models you've given to your clients and happily share this information with any regulators.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Please make sure that you set CUDA device before running the following colab. This can be done by changing `Runtime Type` to use GPU hardware accelerator.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: awscli six !pip install awscli --ignore-installed six
+ # packages added via marimo's package management: timm wandb boto3 !pip install timm wandb boto3
+ return
+
+
+@app.cell
+def _(subprocess):
+ # install packages and prepare dataset
+ #! wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz -q
+ subprocess.call(['wget', 'https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz', '-q'])
+ #! tar -xf imagenette2-160.tgz
+ subprocess.call(['tar', '-xf', 'imagenette2-160.tgz'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## ✍️ Login to W&B
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import timm
+ import boto3
+ import torch
+ import operator
+ import os
+ import logging
+ import warnings
+ import tempfile
+ import torchvision
+ import torch.nn as nn
+ from tqdm import tqdm
+ from torchvision import transforms
+ from timm.utils.log import setup_default_logging
+
+ return (
+ boto3,
+ logging,
+ nn,
+ operator,
+ os,
+ setup_default_logging,
+ tempfile,
+ timm,
+ torch,
+ torchvision,
+ tqdm,
+ transforms,
+ warnings,
+ )
+
+
+@app.cell
+def _(logging):
+ _logger = logging.getLogger('TrainEval')
+ return
+
+
+@app.cell
+def _(transforms):
+ Config = dict(
+ PROJECT='artifacts',
+ DATA_DIR="./imagenette2-160",
+ TRAIN_DATA_DIR="./imagenette2-160/train",
+ TEST_DATA_DIR="./imagenette2-160/val",
+ DEVICE="cuda",
+ MODEL="efficientnet_b3",
+ PRETRAINED=False,
+ LR=3e-4,
+ EPOCHS=3,
+ IMG_SIZE=160,
+ FILENAME='checkpoint-1.pth.tar',
+ ALIAS='v0',
+ BS=96,
+ TRAIN_AUG=transforms.Compose(
+ [
+ transforms.RandomCrop(160),
+ transforms.RandomHorizontalFlip(p=0.5),
+ transforms.ToTensor(),
+ transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+ ]
+ ),
+ TEST_AUG=transforms.Compose(
+ [
+ transforms.CenterCrop(160),
+ transforms.ToTensor(),
+ transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+ ]
+ ),
+ NUM_CHECKPOINTS=2,
+ BUCKET='test-bucket-wandb'
+ )
+ return (Config,)
+
+
+@app.cell
+def _(torch):
+ assert torch.cuda.is_available()
+ DEVICE = torch.device('cuda')
+ return (DEVICE,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🏋️♀️ Model Training and Evaluation
+ """)
+ return
+
+
+@app.cell
+def _(DEVICE, nn, tqdm):
+ def train_fn(model, train_data_loader, optimizer, epoch):
+ model.train()
+ fin_loss = 0.0
+ tk = tqdm(train_data_loader, desc="Epoch" + " [TRAIN] " + str(epoch + 1))
+
+ for t, data in enumerate(tk):
+ data[0] = data[0].to(DEVICE)
+ data[1] = data[1].to(DEVICE)
+
+ optimizer.zero_grad()
+ out = model(data[0])
+ loss = nn.CrossEntropyLoss()(out, data[1])
+ loss.backward()
+ optimizer.step()
+
+ fin_loss += loss.item()
+ tk.set_postfix(
+ {
+ "loss": "%.6f" % float(fin_loss / (t + 1)),
+ "LR": optimizer.param_groups[0]["lr"],
+ }
+ )
+ return fin_loss / len(train_data_loader), optimizer.param_groups[0]["lr"]
+
+ return (train_fn,)
+
+
+@app.cell
+def _(DEVICE, nn, torch, tqdm):
+ def eval_fn(model, eval_data_loader, epoch):
+ model.eval()
+ fin_loss = 0.0
+ tk = tqdm(eval_data_loader, desc="Epoch" + " [VALID] " + str(epoch + 1))
+
+ with torch.no_grad():
+ for t, data in enumerate(tk):
+ data[0] = data[0].to(DEVICE)
+ data[1] = data[1].to(DEVICE)
+ out = model(data[0])
+ loss = nn.CrossEntropyLoss()(out, data[1])
+ fin_loss += loss.item()
+ tk.set_postfix({"loss": "%.6f" % float(fin_loss / (t + 1))})
+ return fin_loss / len(eval_data_loader)
+
+ return (eval_fn,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🏁 Checkpoint Saver
+
+ Track top-n training checkpoints and maintain recovery checkpoints on specified intervals.
+ Hacked together by / Copyright 2020 Ross Wightman
+
+ This script has been adapted from `pytorch-image-models` checkpoint saver script
+ written by Ross Wightman.
+ This script adds Weights and Biases artifact integration on top.
+ (https://github.com/rwightman/pytorch-image-models/blob/master/timm/utils/checkpoint_saver.py)
+ """)
+ return
+
+
+@app.cell
+def _(Config, operator, os, torch, wandb):
+ class CheckpointSaver:
+ def __init__(
+ self,
+ model,
+ optimizer,
+ config=None,
+ checkpoint_prefix='checkpoint',
+ checkpoint_dir='',
+ decreasing=False,
+ max_history=2,
+ wandb_run=None):
+
+ # wandb run
+ self.wandb_run = wandb_run if wandb_run is not None else wandb.init(job_type='model-artifact')
+
+ # objects to save state_dicts of
+ self.model = model
+ self.optimizer = optimizer
+ self.config = config
+
+ # state
+ self.checkpoint_files = [] # (filename, metric) tuples in order of decreasing betterness
+ self.best_epoch = None
+ self.best_metric = None
+ self.curr_recovery_file = ''
+ self.last_recovery_file = ''
+
+ # config
+ self.checkpoint_dir = checkpoint_dir
+ self.save_prefix = checkpoint_prefix
+ self.extension = '.pth.tar'
+ self.decreasing = decreasing # a lower metric is better if True
+ self.cmp = operator.lt if decreasing else operator.gt # True if lhs better than rhs
+ self.max_history = max_history
+ assert self.max_history >= 1
+
+ def save_checkpoint(self, epoch, metric=None):
+ assert epoch >= 0
+ tmp_save_path = os.path.join(self.checkpoint_dir, 'tmp' + self.extension)
+ last_save_path = os.path.join(self.checkpoint_dir, 'last' + self.extension)
+ self._save(tmp_save_path, epoch, metric)
+ if os.path.exists(last_save_path):
+ os.unlink(last_save_path) # required for Windows support.
+ os.rename(tmp_save_path, last_save_path)
+ worst_file = self.checkpoint_files[-1] if self.checkpoint_files else None
+ if (len(self.checkpoint_files) < self.max_history
+ or metric is None or self.cmp(metric, worst_file[1])):
+ if len(self.checkpoint_files) >= self.max_history:
+ self._cleanup_checkpoints(1)
+ filename = '-'.join([self.save_prefix, str(epoch)]) + self.extension
+ save_path = os.path.join(self.checkpoint_dir, filename)
+ os.link(last_save_path, save_path)
+ self.log_artifact(filename, save_path)
+ self.checkpoint_files.append((save_path, metric))
+ self.checkpoint_files = sorted(
+ self.checkpoint_files, key=lambda x: x[1],
+ reverse=not self.decreasing) # sort in descending order if a lower metric is not better
+
+ checkpoints_str = "Current checkpoints:\n"
+ for c in self.checkpoint_files:
+ checkpoints_str += ' {}\n'.format(c)
+ _logger.info(checkpoints_str)
+
+ if metric is not None and (self.best_metric is None or self.cmp(metric, self.best_metric)):
+ self.best_epoch = epoch
+ self.best_metric = metric
+ best_save_path = os.path.join(self.checkpoint_dir, 'model_best' + self.extension)
+ if os.path.exists(best_save_path):
+ os.unlink(best_save_path)
+ os.link(last_save_path, best_save_path)
+
+ return (None, None) if self.best_metric is None else (self.best_metric, self.best_epoch)
+
+ def _save(self, save_path, epoch, metric=None):
+ save_state = {
+ 'epoch': epoch,
+ 'arch': type(self.model).__name__.lower(),
+ 'state_dict': self.model.state_dict(),
+ 'optimizer': self.optimizer.state_dict(),
+ }
+ if metric is not None:
+ save_state['metric'] = metric
+ torch.save(save_state, save_path)
+
+ def _cleanup_checkpoints(self, trim=0):
+ trim = min(len(self.checkpoint_files), trim)
+ delete_index = self.max_history - trim
+ if delete_index < 0 or len(self.checkpoint_files) <= delete_index:
+ return
+ to_delete = self.checkpoint_files[delete_index:]
+ for d in to_delete:
+ try:
+ _logger.debug("Cleaning checkpoint: {}".format(d))
+ # Optionally, only keep top N artifacts in W&B.
+ # self.delete_artifact(os.path.basename(d[0]))
+ os.remove(d[0])
+ except Exception as e:
+ _logger.error("Exception '{}' while deleting checkpoint".format(e))
+ self.checkpoint_files = self.checkpoint_files[:delete_index]
+
+ def log_artifact(self, filename, save_path):
+ try:
+ artifact = wandb.Artifact(filename, type='model')
+ artifact.add_file(save_path)
+ self.wandb_run.log_artifact(artifact)
+ except Exception as e:
+ _logger.error("Exception '{}' while logging wandb artifact".format(e))
+
+ def delete_artifact(self, filename, alias='v0'):
+ api = wandb.Api()
+ artifact = api.artifact(f'{Config["PROJECT"]}/{filename}:{alias}')
+ try:
+ artifact.delete(delete_aliases=True)
+ except Exception as e:
+ _logger.error("Exception '{}' while deleting wandb artifact {}".format(e, filename))
+
+ return (CheckpointSaver,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 💡Bring it all together!
+ """)
+ return
+
+
+@app.cell
+def _(
+ CheckpointSaver,
+ Config,
+ eval_fn,
+ timm,
+ torch,
+ torchvision,
+ train_fn,
+ wandb,
+):
+ def main(wandb_run=None):
+ # train and eval datasets
+ train_dataset = torchvision.datasets.ImageFolder(
+ Config["TRAIN_DATA_DIR"], transform=Config["TRAIN_AUG"]
+ )
+ eval_dataset = torchvision.datasets.ImageFolder(
+ Config["TEST_DATA_DIR"], transform=Config["TEST_AUG"]
+ )
+
+ # train and eval dataloaders
+ train_dataloader = torch.utils.data.DataLoader(
+ train_dataset,
+ batch_size=Config["BS"],
+ shuffle=True,
+ )
+ eval_dataloader = torch.utils.data.DataLoader(
+ eval_dataset, batch_size=Config["BS"],
+ )
+
+ # model
+ model = timm.create_model(Config["MODEL"], pretrained=Config["PRETRAINED"])
+ model = model.cuda()
+
+ # optimizer
+ optimizer = torch.optim.Adam(model.parameters(), lr=Config["LR"])
+
+ # setup checkpoint saver
+ saver = CheckpointSaver(model=model, optimizer=optimizer, config=Config, decreasing=True,
+ wandb_run=wandb_run, max_history=Config['NUM_CHECKPOINTS'])
+
+ for epoch in range(Config["EPOCHS"]):
+ avg_loss_train, lr = train_fn(
+ model, train_dataloader, optimizer, epoch
+ )
+ avg_loss_eval = eval_fn(model, eval_dataloader, epoch)
+ wandb.run.log({
+ "epoch": epoch,
+ "learning rate": lr,
+ "train loss": avg_loss_train,
+ "evaluation loss": avg_loss_eval
+ })
+ saver.save_checkpoint(epoch, metric=avg_loss_eval)
+
+ return (main,)
+
+
+@app.cell
+def _(setup_default_logging):
+ setup_default_logging()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train Model and Log Artifacts to W&B
+ """)
+ return
+
+
+@app.cell
+def _(Config, main, wandb):
+ run = wandb.init(project=Config['PROJECT'], config=Config)
+ wandb.config = Config
+ main(wandb_run=run)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Upload Artifact to S3
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # Setup AWSCLI https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-quickstart.html
+ #! aws configure
+ subprocess.call(['aws', 'configure'])
+ return
+
+
+@app.cell
+def _(boto3, wandb):
+ s3 = boto3.client('s3')
+ api = wandb.Api()
+ return api, s3
+
+
+@app.cell
+def _(Config, api, boto3, os, s3, tempfile, warnings):
+ def upload_artifact_to_s3(config):
+ artifact = api.artifact(f"{config['PROJECT']}/{config['FILENAME']}:{config['ALIAS']}")
+ digest = artifact.digest
+
+ with tempfile.TemporaryDirectory() as tmpdir:
+ path = artifact.download(tmpdir)
+ fname = os.listdir(path)[0]
+ fpath = path + '/' + fname
+
+ _logger.info(f"Downloaded artifact {fname} to {fpath} locally.")
+
+ try:
+ metadata = s3.head_object(Bucket=Config['BUCKET'], Key=fname)['Metadata']
+ except:
+ warnings.warn(f"""File {fname} does not already exist in Bucket {Config['BUCKET']} on AWS.\
+ Cleaning up AWS bucket for any existing files, and uploading new \
+ artifact.""")
+ bucket = boto3.resource('s3').Bucket(Config['BUCKET'])
+ bucket.objects.all().delete()
+ metadata = {'digest': -1}
+
+ # upload files to S3 if digests are different
+ if metadata['digest']!=digest:
+ s3.upload_file(fpath, Config['BUCKET'], fname, ExtraArgs={"Metadata": {"digest": digest}})
+ else:
+ _logger.info(f"File {fname} already exists in Bucket {Config['BUCKET']} on AWS with same digest. Nothing to do.")
+
+ return (upload_artifact_to_s3,)
+
+
+@app.cell
+def _(Config, upload_artifact_to_s3):
+ upload_artifact_to_s3(config=Config)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-wandb-artifact-tags/wandb_artifacts_wandb_artifact_tags.py b/marimo/convert/wandb-artifacts-wandb-artifact-tags/wandb_artifacts_wandb_artifact_tags.py
new file mode 100644
index 00000000..48b08a8d
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-wandb-artifact-tags/wandb_artifacts_wandb_artifact_tags.py
@@ -0,0 +1,288 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## W&B Artifact version tags
+
+ W&B Models supports finding and retrieving artifacts using version tags from Registry or ML projects. Filtering directly by artifact version tags via the SDK offers a more efficient way to retrieve only the artifacts you need instead of grabbing every artifact from a collection and parsing them manually or using aliases which can present challenges due to enforced uniqueness within each collection.
+
+ In this notebook, you will create multiple model versions in a collection in the Model registry. Each of the versions has been assigned multiple tags making discoverability and retrieval simple using the SDK.
+
+ You can read more about organizing artifacts with tags in the W&B documentation.
+
+ You can find Registry in your W&B account by clicking on the Registry link in the left sidebar in the Applications section.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Prerequisites
+
+ Install the W&B Python SDK and log in:
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qU
+ return
+
+
+@app.cell
+def _():
+ # Log in to your W&B account
+ import wandb
+ wandb.login()
+ return (wandb,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Initialize a W&B run
+
+ Import additional Python libraries and initialize a W&B run to generate demo artifacts[link text](https://):
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ import random
+ import math
+ import pandas as pd
+ import numpy as np
+ import os
+
+ PROJECT = "artifacts-example"
+ JOB_TYPE = "generate_artifacts"
+
+ run = wandb.init(
+ project=PROJECT,
+ job_type=JOB_TYPE
+ )
+ return np, pd, run
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Generate demo artifact versions with tags
+
+ The following code block will generate four artifact versions for each AWS region. Each artifact version represents the result of fine-tuning a base model. The artifact versions will be tagged with the following:
+
+ 1. Region tag: Specifies the AWS region where the model is stored or deployed
+ 2. Base model tag: Indicates which base model was fine-tuned
+ 3. Status tag: Specifies the model's status, when applicable (e.g., production, candidate, archived)
+
+ To create the artifact versions, simply execute the code block.
+ """)
+ return
+
+
+@app.cell
+def _(np, pd, run, wandb):
+ regions = ['us-west-2', 'eu-central-1', 'me-central-1', 'ap-east-1']
+ artifact_base_model_by_region = {'us-west-2': ['Llama 3 Instruct - 70B', 'Llama 3 Instruct - 70B', 'Gemini 1_5 Pro', 'Claude 3_5 Sonnet'], 'eu-central-1': ['Llama 3 Instruct - 70B', 'Gemini 1_5 Pro', 'Claude 3_5 Sonnet', 'Llama 3 Instruct - 70B'], 'me-central-1': ['Llama 3 Instruct - 70B', 'Gemini 1_5 Pro', 'Gemini 1_5 Pro', 'Claude 3_5 Sonnet'], 'ap-east-1': ['Claude 3_5 Sonnet', 'Gemini 1_5 Pro', 'Claude 3_5 Sonnet', 'Gemini 1_5 Pro']}
+ artifact_metadata_by_region = {'us-west-2': [0.7, 0.73, 0.77, 0.71], 'eu-central-1': [0.76, 0.81, 0.79, 0.77], 'me-central-1': [0.77, 0.68, 0.74, 0.83], 'ap-east-1': [0.82, 0.79, 0.76, 0.76]}
+ artifact_status_by_region = {'us-west-2': ['production', 'candidate', 'archived', 'archived'], 'eu-central-1': ['production', 'candidate', 'archived', 'archived'], 'me-central-1': ['production', 'candidate', 'candidate', 'archived'], 'ap-east-1': ['production', 'candidate', 'archived', 'archived']}
+ for region in regions:
+ # set up tag values
+ # AWS regions (us-west-2 = Oregon, eu-central-1 = Frankfurt, me-central-1 = UAE, ap-east-1 = Hong Kong)
+ i = 0
+ # add base model tags to each artifact version
+ for base_model in artifact_base_model_by_region[region]:
+ df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)), columns=list('ABCD'))
+ df.to_json('test_model.pt', orient='records', lines=True)
+ model_name = 'test_model_' + region + '_' + str(i)
+ at = wandb.Artifact(name=model_name, type='model')
+ at.add_file('test_model.pt')
+ # add performance metadata to each artifact version
+ arti = run.log_artifact(at)
+ arti.wait()
+ if i == 3:
+ arti.tags = [region, base_model]
+ arti.metadata = {'accuracy': artifact_metadata_by_region[region][i]}
+ else:
+ # add model status tags to each artifact version
+ arti.tags = [region, base_model, artifact_status_by_region[region][i]]
+ arti.metadata = {'accuracy': artifact_metadata_by_region[region][i]}
+ arti.save()
+ registered_at = run.link_artifact(at, f'wandb-registry-model/Artifact Demo Models')
+ i = i + 1
+ # create 4 artifact versions for each region
+ # mark the run as finished
+ run.finish() # create random dataset to use as a model version # assign a status tag to each artifact except the 4th which has no status tag assigned # Provide one or more tags in a list # link the artifact to the model registry
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Filter artifacts using version tags
+
+ After generating the 16 artifacts using the code block above, the Artifact Demo Models collection in your model registry should look like this:
+
+
+ 
+
+ This colab presumes that you have not one production model in a collection, but multiple production models, one for each AWS region. In cases where only a single production model exists in a collection, using a unique alias to identify this model is generally the right approach. The enforced uniqueness of aliases makes them extremely valuable when searching for and retrieving a specific artifact version. Aliases can also be used in W&B Models to trigger automated workflows, or Automations, which are often used for model testing and deployment as part of a CI/CD pipeline. But event-based triggers are not always necessary and there are times when you want to track down and retrieve multiple artifact versions based on a search filter, such as when multiple production models exist. This is when tags are the answer.
+
+ An example where version artifact tags come in handy is when deploying models using Amazon Sagemaker. The Amazon S3 bucket where the model artifacts are stored must be in the same AWS Region as the model that you are creating. In cases where it is a requirement to deploy specific models to specific regions, retrieving those models using an AWS region artifact tag ensures that the right model exists and that the right model is deployed.
+
+ Attaching tags to artifact versions also helps with compliance requirements. Filtering by artifact tag retruns a specific model of interest and, from there, it is easy to track the detailed lineage of this model, including all input and output artifacts, using W&B Registry. For example, during an audit, it might be required to produce the exact dataset used for training the model deployed in the Central Europe region to ensure that it did not contain any PII data or other sensitive information.
+
+ We have compiled a number of use cases that require filtering by artifact version tags to retrieve the right artifacts. To see version tag filtering using the SDK in action, just execute the following code block:
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ api = wandb.Api()
+ # if you belong to multiple orgs, prefix the 'name' value in the api.artifacts call with the org you are fetching from:
+ # name=f"{INSERT_ORG_NAME}/wandb-registry-model/Artifact Demo Models"
+
+ ##########################################################
+ ##### Artifact Version Tag Use Cases #####################
+ ##########################################################
+
+ # 1. As an ML Platform Engineer, I need to ensure the correct models are deployed in each region.
+ # Action: Filter by `production`, `us-west-2`, and base_model tags. (Base model should be `Llama 3 Instruct - 70B`).
+ # Scenario: Verify that only production-ready models are deployed in the US-West region, using the correct base model.
+ # Benefit: Ensures accurate deployment for specific regions and base models.
+
+ print("\nArtifact Version Tag Use Case #1")
+ print("As an ML Platform Engineer, I need to ensure the correct models are deployed in each region.")
+ print("---------------------------------------\n")
+ artifact_versions = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['production', 'us-west-2'])
+ for av in artifact_versions:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+
+ # 2. As a Data Scientist, I want to compare model performance in US-West and EU-Central regions.
+ # Action: Filter by `us-west-2`, `eu-central-1` tags.
+ # Scenario: Compare models deployed in US-West and EU-Central based on performance metrics.
+ # Benefit: Identifies performance variations between regions for optimization.
+
+ print("\n\nArtifact Version Tag Use Case #2")
+ print("As a Data Scientist, I want to compare model performance in US-West and EU-Central regions.")
+ print("---------------------------------------\n")
+ artifact_versions_us_west_2 = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['production', 'us-west-2'])
+ artifact_versions_eu_central_1 = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['production', 'eu-central-1'])
+ for av in artifact_versions_us_west_2:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+ for av in artifact_versions_eu_central_1:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+
+ # 3. As an ML Ops Engineer, I need to check the production model in EU-Central to troubleshoot an issue.
+ # Action: Filter by `production`, `eu-central-1`, and model_version tags.
+ # Scenario: Quickly find the production model version deployed in EU-Central for debugging.
+ # Benefit: Saves time and ensures the correct model is under investigation.
+
+ print("\n\nArtifact Version Tag Use Case #3")
+ print("As an ML Ops Engineer, I need to check the production model in EU-Central to troubleshoot an issue.")
+ print("---------------------------------------\n")
+ artifact_versions = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['production', 'eu-central-1'])
+ for av in artifact_versions:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+
+ # 4. As a Product Manager, I want to review all candidate models in ME-Central for possible promotion.
+ # Action: Filter by `candidate`, `me-central-1`, and base_model tags.
+ # Scenario: Gather models tagged as candidates for the ME-Central region and evaluate for production.
+ # Benefit: Streamlines the process of selecting models for promotion.
+
+ print("\n\nArtifact Version Tag Use Case #4")
+ print("As a Product Manager, I want to review all candidate models in ME-Central for possible promotion.")
+ print("---------------------------------------\n")
+ artifact_versions = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['candidate', 'me-central-1'])
+ for av in artifact_versions:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+
+ # 5. As a Compliance Officer, I need to audit archived models in US-West and ME-Central.
+ # Action: Filter by `archived`, `us-west-2`, and `me-central-1` tags.
+ # Scenario: Find older models in these regions to ensure they meet compliance requirements.
+ # Benefit: Efficiently audits the models without needing manual searches.
+
+ print("\n\nArtifact Version Tag Use Case #5")
+ print("As a Compliance Officer, I need to audit archived models in US-West and ME-Central.")
+ print("---------------------------------------\n")
+ artifact_versions_us_west_2 = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['archived', 'us-west-2'])
+ artifact_versions_me_central_1 = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['archived', 'me-central-1'])
+ for av in artifact_versions_us_west_2:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+ for av in artifact_versions_me_central_1:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+
+ # 6. As an ML Engineer, I want to evaluate models using the `Claude 3_5 Sonnet` base model in AP-East and EU-Central regions.
+ # Action: Filter by `Claude 3_5 Sonnet`, `ap-east-1`, and `eu-central-1` tags.
+ # Review and compare models fine-tuned on `Claude 3_5 Sonnet` across these regions.
+ # Benefit: Helps assess and improve performance for models based on the `Claude 3_5 Sonnet` base.
+
+ print("\n\nArtifact Version Tag Use Case #6")
+ print("As an ML Engineer, I want to evaluate models using the `Claude 3_5 Sonnet` base model in AP-East and EU-Central regions.")
+ print("---------------------------------------\n")
+ artifact_versions_ap_east_1 = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['Claude 3_5 Sonnet', 'ap-east-1'])
+ artifact_versions_eu_central_1 = api.artifacts(type_name="model", name="wandb-registry-model/Artifact Demo Models", tags=['Claude 3_5 Sonnet', 'eu-central-1'])
+ for av in artifact_versions_ap_east_1:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+ for av in artifact_versions_eu_central_1:
+ print("Artifact Version: " + str(av.name) + "\nTags: " + str(av.tags) + "\nVersion: " + str(av.version) + "\nMetadata: " + str(av.metadata))
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Artifact version filtering results
+
+ As you can see from the output of the script, it is possible to use one or more tags to retrieve one or more artifact versions from your registry collection.
+
+ It is also possible to search by collection names, tags, and version tags from within the UI. Just use the search bar inside of any registry to find the artifacts that you need. The following screenshot shows the results for a search on "ap-east-1" in our Model registry.
+
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-artifacts-wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_wandb_artifacts_time_to_live_ttl_walkthrough.py b/marimo/convert/wandb-artifacts-wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_wandb_artifacts_time_to_live_ttl_walkthrough.py
new file mode 100644
index 00000000..42a73351
--- /dev/null
+++ b/marimo/convert/wandb-artifacts-wandb-artifacts-time-to-live-ttl-walkthrough/wandb_artifacts_wandb_artifacts_time_to_live_ttl_walkthrough.py
@@ -0,0 +1,345 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Weights & Biases Artifacts Time-to-live (TTL) Walkthrough
+ W&B Artifacts now supports setting time-to-live policies on each version of an Artifact. The feature is currently available in W&B SaaS Cloud and will be released to Enterprise customers using W&B Server in version 0.42.0. The following examples show the use TTL policy in a common Artifact logging workflow. We'll cover:
+
+ - Setting a TTL policy when creating an Artifact
+ - Retroactively setting TTL for a specific Artifact aliases
+ - Using the W&B API to set a TTL for all versions of an Artifact
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ Let's do a few things before we get started. Below we will:
+
+ - Install the wandb library and download a dataset
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ log to wandb
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ wandb.login()
+ return (wandb,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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 Colab session. The block below:
+
+ - 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
+
+ *Note: we overwrite the files every time we execute this so we get new Artifact versions.*
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ imagenette_url = "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz"
+ #! wget {imagenette_url} -O "imagenette.tgz"
+ subprocess.call(['wget', str(imagenette_url), '-O', 'imagenette.tgz'])
+ return
+
+
+@app.cell
+def _():
+ def untar_file(file_path, dest_path):
+ import tarfile
+ with tarfile.open(file_path, "r:gz") as tar:
+ tar.extractall(dest_path)
+
+ untar_file("imagenette.tgz", "./")
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import random
+ from pathlib import Path
+
+ dataset_dir = Path("imagenette2-160")
+
+ # 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 = Path("imagenette2-160/train")
+ val_source_dir = Path("imagenette2-160/val")
+ return Path, random, train_source_dir, val_source_dir
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Image Preview
+ Quick block to view some of the images in the sampled dataset.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import os
+ from PIL import Image
+ import matplotlib.pyplot as plt
+
+ 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))
+
+ # Iterate over the images and display them
+ for i, img_path in enumerate(images):
+ img = Image.open(img_path)
+ axes[i].imshow(img)
+ axes[i].axis('off') # Turn off axis labels
+
+ plt.tight_layout()
+ plt.show()
+
+ return Image, show_sample_images
+
+
+@app.cell
+def _(show_sample_images, train_source_dir):
+ show_sample_images(train_source_dir)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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:
+
+ - 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
+
+ > We will use the train dataset as our real data and the validation dataset as our fake data.
+ """)
+ return
+
+
+@app.cell
+def _(train_source_dir, val_source_dir, wandb):
+ from datetime import timedelta
+ with wandb.init(entity='wandb-smle', project='artifacts-ttl-demo', job_type='raw-data') as _run:
+ raw_real_art = wandb.Artifact('real-raw', type='dataset', description='Raw sample train Imagenette')
+ raw_real_art.add_dir(train_source_dir)
+ raw_real_art.ttl = timedelta(days=10)
+ _run.log_artifact(raw_real_art)
+ raw_fake_art = wandb.Artifact('fake-raw', type='dataset', description='Raw sample from val Imagenette')
+ raw_fake_art.add_dir(val_source_dir)
+ _run.log_artifact(raw_fake_art)
+ _run.finish()
+ return (timedelta,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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:
+
+ - 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
+ """)
+ return
+
+
+@app.cell
+def _(timedelta, wandb):
+ with wandb.init(entity='wandb-smle', project='artifacts-ttl-demo', job_type='modify-ttl') as _run:
+ extended_art = _run.use_artifact('wandb-smle/artifacts-ttl-demo/real-raw:extended')
+ extended_art.ttl = timedelta(days=365) # Delete in a year
+ extended_art.save()
+ compliant_art = _run.use_artifact('wandb-smle/artifacts-ttl-demo/real-raw:compliant')
+ compliant_art.ttl = None
+ compliant_art.save()
+ print(extended_art.ttl)
+ print(compliant_art.ttl)
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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:
+
+ - 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
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # Artifact metadata extraction
+ _api = wandb.Api()
+ entity, project = ('wandb-smle', 'artifacts-ttl-demo')
+ # 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:
+ _version_names.append(f'{_artifact.name}/{_artifact.version}')
+ with wandb.init(entity='wandb-smle', project='artifacts-ttl-demo', job_type='modify-ttl') as _run:
+ for _version in _version_names:
+ _version_art = _run.use_artifact(f"wandb-smle/artifacts-ttl-demo/{'/'.join(_version.split('/')[:-1])}") # Can be edited to just display individual elements
+ _version_art.ttl = None
+ _version_art.save()
+ print(_version_art.ttl)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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/guides/artifacts/ttl#set-default-ttl-policies-for-a-team) section of the W&B documentation.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Preprocess and log a new Artifact
+ """)
+ return
+
+
+@app.cell
+def _(Image, Path, wandb):
+ real_prepro_dir = Path('data/prepro/real')
+ real_prepro_dir.mkdir(parents=True, exist_ok=True)
+
+ def preprocess_image(image_path):
+ """Resize the image to 64x64"""
+ return Image.open(image_path).resize((64, 64))
+ with wandb.init(entity='wandb-smle', project='artifacts-ttl-demo', job_type='preprocessing') as _run:
+ real_art = _run.use_artifact('wandb-smle/artifacts-ttl-demo/real-raw:latest')
+ real_images = Path(real_art.download())
+ 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 CIFAR')
+ prepro_real_art.add_dir(real_prepro_dir)
+ _run.log_artifact(prepro_real_art)
+ _run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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
+
+
+@app.cell
+def _(random, timedelta, wandb):
+ _api = wandb.Api()
+ _artifact = _api.artifact('wandb-smle/artifacts-ttl-demo/real-raw:latest')
+ # For demo purposes we'll just do this on the latest version of the real dataset
+ consumer_runs = _artifact.used_by()
+ _version_names = []
+ for _run in consumer_runs:
+ # Same pattern from above to get all downstream versions
+ for _artifact in iter(_run.logged_artifacts()):
+ if _artifact.type == 'dataset':
+ _version_names.append(f'{_artifact.name}/{_artifact.version}')
+ with wandb.init(entity='wandb-smle', project='artifacts-ttl-demo', job_type='modify-ttl') as _run: # filter for datasets only
+ for _version in _version_names:
+ _version_art = _run.use_artifact(f"wandb-smle/artifacts-ttl-demo/{'/'.join(_version.split('/')[:-1])}") # Can be edited to just display individual elements
+ _version_art.ttl = timedelta(days=random.randint(1, 100))
+ _version_art.save()
+ _run.finish() # set ttl to a random integer so we can see changes in the UI after we run this
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-configs-in-w-b/wandb_log_configs_in_w_b.py b/marimo/convert/wandb-log-configs-in-w-b/wandb_log_configs_in_w_b.py
new file mode 100644
index 00000000..c908228c
--- /dev/null
+++ b/marimo/convert/wandb-log-configs-in-w-b/wandb_log_configs_in_w_b.py
@@ -0,0 +1,366 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ # Quickstart
+ Use [Weights & Biases](https://wandb.ai)
+ for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To get started, just `pip install` the package and log using `wandb.login()`
+ If this is your first time using `wandb`, you'll need to sign up. It's easy!
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -Uq wandb
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # What's a `config` for?
+
+ Set [`wandb.config`](https://docs.wandb.ai/guides/track/config)
+ once at the beginning of your script to save your training configuration: hyperparameters, input settings like dataset name or model type, and include any other independent variables or metadata for your experiments.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Why does that matter?
+
+ This is useful for analyzing your experiments and reproducing your work in the future. You'll be able to group by `config` values in our web interface, comparing the settings of different runs and seeing how these affect the output.
+
+ > Note that output metrics or dependent variables (like loss and accuracy) should be saved with `wandb.log` instead.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # How do I set up a `config`?
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Your `config` should be set just once at the beginning of your training experiment.
+
+ But workflows differ, so we offer a number of ways to set up your config.
+
+ Let's look at all the ways you can create and send the config dictionary to the Dashboard!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setting the `config` at `init`ialization
+
+ The best time to set the `config` values is when you call [`wandb.init`](https://docs.wandb.ai/guides/track/launch),
+ by passing a dictionary as the `config` keyword argument.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.init(project='config_example', config={'dataset': 'CelebA', 'type': 'baseline'})
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Head to the [Run page](https://docs.wandb.ai/ref/app/pages/run-page)
+ linked in the output of `wandb.init`
+ and head to the [Overview tab](https://docs.wandb.ai/ref/app/pages/run-page#overview-tab)
+ (top of the list of panels on the left-most side of the screen).
+ You'll see a "Config" section that looks like this:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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`.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Adding to the `config` by hand
+ You can add more parameters to the `config` later if you want:
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.config.epochs = 4
+ wandb.config["batch_size"] = 32
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Now, your Config section on the dashboard has been updated:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ import argparse
+
+ 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=[])
+ wandb.config.update(args)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's the updated Config panel:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ api = wandb.Api()
+
+ # pulling the relevant info automatically from the run object
+ # this can also be found on the website
+ username = wandb.run.entity
+ project = wandb.run.project
+ run_id = wandb.run.id
+
+ run = api.run(f"{username}/{project}/{run_id}")
+ run.config["bar"] = 32
+ run.update()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's what the final Config panel looks like:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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".
+
+ Run the cell below to see this in action!
+ """)
+ return
+
+
+@app.cell
+def _():
+ from IPython import display
+
+ display.YouTubeVideo("aSMXwOSPtJE", rel=0, width=450)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ 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.
+
+ 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.
+
+ 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.
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-customize-metric-logging-with-define-metric/wandb_log_customize_metric_logging_with_define_metric.py b/marimo/convert/wandb-log-customize-metric-logging-with-define-metric/wandb_log_customize_metric_logging_with_define_metric.py
new file mode 100644
index 00000000..d1f3425a
--- /dev/null
+++ b/marimo/convert/wandb-log-customize-metric-logging-with-define-metric/wandb_log_customize_metric_logging_with_define_metric.py
@@ -0,0 +1,156 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+
+ # Define your custom metrics with `define_metric`
+
+ Use `define_metric` to set custom x-axes or capture the min and max values of your metrics.
+
+ For more details, [see the docs](http://wandb.me/define-metric-docs).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -Uq
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import random
+
+ return random, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Custom X Axis
+
+ Here's how to set a custom step so your charts have a custom x-axis:
+ ```python
+ wandb.define_metric("my-metric", step_metric='my-custom-x-axis')
+ ```
+ """)
+ return
+
+
+@app.cell
+def _(random, wandb):
+ random.seed(1)
+ wandb.init(project='define-metric-demo', notes='custom step')
+ # Initalize a new run
+ wandb.define_metric('custom_step')
+ wandb.define_metric('validation/loss', step_metric='custom_step')
+ # Define the custom x axis metric
+ for _i in range(10):
+ _log_dict = {'train/loss': 1 / (_i + 1), 'custom_step': _i ** 2, 'validation/loss': 1 / (_i + 1)}
+ # Define which metrics to plot against that x-axis
+ wandb.log(_log_dict)
+ # Use this in the context of a jupyter notebook to mark a run finished
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Run the cell above and click on the link that prints out to see the dashboard. It will look something like this:
+ - `train_loss` is plotted against the standard W&B internal step
+ - `custom_step` is plotted too, so you can see how it increases over the W&B internal step
+ - `validation_loss` is plotted against the `custom_step`, replacing the default with the x-axis as the W&B internal step
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Min/Max of Metrics
+
+ Each time you call `wandb.log()` to log a metric, you're writing to run `history`. The run `summary` saves a single value for each metric. By default, `summary` captures the final step of `history`. So if you log accuracy for 100 steps, your `history` will have all 100 steps and your `summary` will have just the final value for accuracy.
+
+ Sometimes, you want to get the _best_ value instead of the _last_ value for a metric and save that to `summary`. That's where `define_metric` comes in.
+
+ Here, you can set `summary=` to either `max` or `min`.
+
+ ```python
+ wandb.define_metric("my-metric", summary="max")
+ ```
+ """)
+ return
+
+
+@app.cell
+def _(random, wandb):
+ random.seed(1)
+ wandb.init(project='define-metric-demo', notes='min of loss, max of acc')
+ # Start a new run
+ wandb.define_metric('loss', summary='min')
+ wandb.define_metric('acc', summary='max')
+ # For loss, capture the min value from history in summary
+ for _i in range(10):
+ _log_dict = {'loss': random.uniform(0, 1 / (_i + 1)), 'acc': random.uniform(1 / (_i + 1), 1)}
+ # For acc, capture the max value from history in summary
+ wandb.log(_log_dict)
+ # Simulate a training loop where we're logging metrics
+ # Mark the run as finished, useful in the context of Jupyter notebooks
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Run the cell above and click on the project page link that prints out to see the dashboard. It will look something like this:
+ - `acc.max` is visible in the sidebar, saved in the run summary
+ - `loss.min` is visible in the sidebar, saved in the run summary
+
+ You can see the summary values in the Project Page Table. Here I've pinned two columns in the sidebar, which you can see on the left.
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-image-logging-de-duplication/wandb_log_image_logging_de_duplication.py b/marimo/convert/wandb-log-image-logging-de-duplication/wandb_log_image_logging_de_duplication.py
new file mode 100644
index 00000000..5bbe8089
--- /dev/null
+++ b/marimo/convert/wandb-log-image-logging-de-duplication/wandb_log_image_logging_de_duplication.py
@@ -0,0 +1,81 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ import PIL
+ import numpy as np
+
+ def write_img(path):
+ PIL.Image.fromarray(np.random.rand(128,128), mode="L").save(path)
+
+ def setup_demo(num_images=10):
+ paths = []
+ for ndx in range(num_images):
+ path = f"./img_{ndx}.png"
+ write_img(path)
+ paths.append(path)
+ return paths
+
+ return (setup_demo,)
+
+
+@app.cell
+def _(setup_demo):
+ IMAGE_PATHS = setup_demo()
+ return (IMAGE_PATHS,)
+
+
+@app.cell
+def _(IMAGE_PATHS):
+ import wandb
+ wandb.init(project='image_docs')
+ # Step 1: Add your Images to an Artifact
+ _art = wandb.Artifact('my_images', 'dataset')
+ for path in IMAGE_PATHS:
+ _art.add(wandb.Image(path), path)
+ wandb.log_artifact(_art)
+ wandb.finish()
+ return (wandb,)
+
+
+@app.cell
+def _(IMAGE_PATHS, wandb):
+ run = wandb.init(project='image_docs')
+ _art = wandb.use_artifact('my_images:latest')
+ img_1 = _art.get(IMAGE_PATHS[0])
+ wandb.log({'image': img_1})
+ wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-log-a-confusion-matrix-with-w-b/wandb_log_log_a_confusion_matrix_with_w_b.py b/marimo/convert/wandb-log-log-a-confusion-matrix-with-w-b/wandb_log_log_a_confusion_matrix_with_w_b.py
new file mode 100644
index 00000000..bb9717c0
--- /dev/null
+++ b/marimo/convert/wandb-log-log-a-confusion-matrix-with-w-b/wandb_log_log_a_confusion_matrix_with_w_b.py
@@ -0,0 +1,289 @@
+# /// script
+# dependencies = ["tensorflow", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ # Plot a Confusion Matrix with W&B
+
+ How to log a [confusion matrix](https://scikit-learn.org/stable/auto_examples/model_selection/plot_confusion_matrix.html) with [Vega](https://vega.github.io/vega/docs/) in [Weights & Biases](https://www.wandb.com).
+
+ ## Method: wandb.plot.confusion_matrix()
+
+ - 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).
+
+ This Colab 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).
+
+ 
+
+ 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.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup: Download data
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Note: **this stage might take a few minutes (~3.6GB of data)**. If you end up needing to rerun this cell, comment out the first capture line (change ```%%capture``` to ```#%%capture``` ) so you can respond to the prompt about re-downloading the dataset (and see the progress bar).
+
+ Download sample data: 10,000 training images and 2,000 validation images from the [iNaturalist dataset](https://github.com/visipedia/inat_comp), evenly distributed across 10 classes of living things like birds, insects, plants, and mammals (names given in Latin—so Aves, Insecta, Plantae, etc :). We will fine-tune a convolutional neural network already trained on ImageNet on this task: given a photo of a living thing, correctly classify it into one of the 10 classes.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !curl -SL https://storage.googleapis.com/wandb_datasets/nature_12K.zip > nature_12K.zip
+ # !unzip nature_12K.zip
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install dependencies
+
+ Install tensorflow and wandb; log in to wandb.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: tensorflow !pip install tensorflow -qqq
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training code
+
+ Feel free to try different values for "NUM_TRAIN" and "NUM_EPOCHS" below so you can see a variety of PR curves (generally better ones with more training examples/longer training time)
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # this determines the name of your wandb project, where all your
+ # runs will be logged
+ PROJECT_NAME = "confusion_matrix"
+
+ # EXPERIMENT CONFIG
+ #---------------------------
+ # try changing the number of training examples
+ # to generate a range of different models
+ NUM_TRAIN = 100 # try 500, 1000, 2000, or max 10000
+ NUM_EPOCHS = 1 # try 3, 5, or as many as you like
+
+ import numpy as np
+ from sklearn.metrics import precision_recall_curve, roc_curve
+ from sklearn.metrics import average_precision_score
+ from sklearn.preprocessing import label_binarize
+
+ from tensorflow.keras.applications.inception_v3 import InceptionV3
+ from tensorflow.keras.callbacks import Callback
+ from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
+ from tensorflow.keras.models import Model
+ from tensorflow.keras.preprocessing.image import ImageDataGenerator
+ import tensorflow as tf
+
+ from wandb.integration.keras import WandbMetricsLogger, WandbModelCheckpoint
+
+ # local paths to data
+ train_data = "inaturalist_12K/train"
+ val_data = "inaturalist_12K/val"
+
+ # 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 10000 for the full dataset
+ # GOOD CONFIG TO TRY: 100, 500, 1000, 2000
+ "num_train" : NUM_TRAIN, # up to 10000,
+ # number of images used to validate--set low for demo training speed
+ # you can set this up to 2000 for the full dataset
+ "num_val" : 500, #2000,
+ "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
+ # you can set this up to 5, 10, or more for better results
+ # GOOD CONFIG TO TRY: 3, 5, 10
+ "pretrain_epochs" : NUM_EPOCHS, #5,
+ # number of validation data batches to use when computing metrics
+ # at the end of each epoch
+ "num_log_batches": 15,
+ # random seed
+ "random_seed": 23
+ }
+
+ def build_model(fc_size, num_classes):
+ """Load InceptionV3 with ImageNet weights, freeze it,
+ and attach a finetuning top for this classification task"""
+ # load InceptionV3 as base
+ base = InceptionV3(weights="imagenet", include_top="False")
+ # freeze base layers
+ for layer in base.layers:
+ layer.trainable = False
+ x = base.get_layer('mixed10').output
+
+ # attach a fine-tuning layer
+ x = GlobalAveragePooling2D()(x)
+ x = Dense(fc_size, activation='relu')(x)
+ guesses = Dense(num_classes, activation='softmax')(x)
+
+ model = Model(inputs=base.input, outputs=guesses)
+ model.compile(optimizer='rmsprop', loss='categorical_crossentropy',
+ metrics=['accuracy'])
+ return model
+
+ def pretrain():
+ """ 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"""
+ # track this experiment with wandb: all runs will be sent
+ # to the given project name
+ run = wandb.init(project=PROJECT_NAME, config=config_defaults)
+ cfg = run.config
+
+ # set random seed
+ tf.random.set_seed(cfg.random_seed)
+ # also set numpy seed to control train/val dataset split
+ np.random.seed(cfg.random_seed)
+
+ # create train and validation data generators
+ train_datagen = ImageDataGenerator(
+ rescale=1. / 255,
+ shear_range=0.2,
+ zoom_range=0.2,
+ horizontal_flip=True)
+ val_datagen = ImageDataGenerator(rescale=1. / 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 = val_datagen.flow_from_directory(
+ val_data,
+ target_size=(cfg.img_width, cfg.img_height),
+ batch_size=cfg.batch_size,
+ class_mode='categorical')
+
+ # instantiate model and callbacks
+ model = build_model(cfg.fc_size, cfg.num_classes)
+ callbacks = [WandbMetricsLogger(), WandbModelCheckpoint("checkpoint.keras"), PRMetrics(val_generator, cfg.num_log_batches)]
+
+ # train!
+ model.fit(
+ train_generator,
+ steps_per_epoch = cfg.num_train // cfg.batch_size,
+ epochs=cfg.pretrain_epochs,
+ validation_data=val_generator,
+ callbacks = callbacks,
+ validation_steps=cfg.num_val // cfg.batch_size)
+
+ run.finish()
+
+ class PRMetrics(Callback):
+ """ Custom callback to compute metrics at the end of each training epoch"""
+ def __init__(self, generator=None, num_log_batches=1):
+ self.generator = generator
+ self.num_batches = num_log_batches
+ # store full names of classes
+ self.flat_class_names = [k for k, v in generator.class_indices.items()]
+
+ def on_epoch_end(self, epoch, logs={}):
+ # collect validation data and ground truth labels from generator
+ val_data, val_labels = zip(*(self.generator[i] for i 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)
+ 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 confusion matrix
+ # the key "conf_mat" is the id of the plot--do not change
+ # this if you want subsequent runs to show up on the same plot
+ wandb.log({"conf_mat" : wandb.plot.confusion_matrix(probs=None,
+ preds=top_pred_ids, y_true=ground_truth_class_ids,
+ class_names=self.flat_class_names)})
+
+ return (pretrain,)
+
+
+@app.cell
+def _(pretrain):
+ # run this cell to launch your experiment!
+ # charts will show up in your run page under the heading "Media" or
+ # "Custom Charts", which you may need to click on to expand
+ pretrain()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-log-almost-anything-with-w-b-media/wandb_log_log_almost_anything_with_w_b_media.py b/marimo/convert/wandb-log-log-almost-anything-with-w-b-media/wandb_log_log_almost_anything_with_w_b_media.py
new file mode 100644
index 00000000..c9921ac0
--- /dev/null
+++ b/marimo/convert/wandb-log-log-almost-anything-with-w-b-media/wandb_log_log_almost_anything_with_w_b_media.py
@@ -0,0 +1,427 @@
+# /// script
+# dependencies = ["soundfile", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+ Use [Weights & Biases](https://wandb.com) for machine learning experiment tracking, dataset versioning, and project collaboration.
+
+
+
+
+
+
+
+ # Log (Almost) Anything with W&B Media
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this notebook, we'll show you how to visualize a model's predictions with Weights & Biases – images, videos, audio, tables, HTML, metrics, plots, 3D objects and point clouds.
+
+ ### Follow along with a [video tutorial →](http://wandb.me/media-video)!
+ #### View plots in interactive [dashboard →](https://app.wandb.ai/lavanyashukla/visualize-predictions/reports/Visualize-Model-Predictions--Vmlldzo1NjM4OA).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # packages added via marimo's package management: wandb !pip install wandb -qq
+
+ # Fetch audio, video and other data files to log
+ #! git clone --depth 1 https://github.com/wandb/examples.git
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/wandb/examples.git'])
+ # packages added via marimo's package management: soundfile !pip install soundfile -qq
+
+ import warnings
+ warnings.filterwarnings("ignore", category=UserWarning)
+ return
+
+
+@app.cell
+def _():
+ import pandas as pd
+ import numpy as np
+ import wandb
+
+ return np, pd, wandb
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log metrics
+ """)
+ return
+
+
+@app.cell
+def _(pd):
+ # Get Apple stock price data from
+ # https://www.macrotrends.net/stocks/charts/AAPL/apple/stock-price-history
+ # Read in dataset
+ apple = pd.read_csv("examples/examples/data/apple.csv")
+ apple = apple[-1000:]
+ return (apple,)
+
+
+@app.cell
+def _(apple, wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="metrics")
+
+ # Log the metric on each step
+ for price in apple['close']:
+ wandb.log({"Stock Price": price})
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log plots
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ import matplotlib.pyplot as plt
+ wandb.init(project='visualize-predictions', name='plots')
+ # Initialize a new run
+ _fibonacci = [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
+ fig, ax = plt.subplots()
+ # Make the plot
+ ax.plot(_fibonacci)
+ ax.set_ylabel('some interesting numbers')
+ wandb.log({'plot': fig})
+ wandb.finish()
+ # Log the plot
+ fig
+ return (plt,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log Histograms
+ """)
+ return
+
+
+@app.cell
+def _(np, wandb):
+ # Initialize a new run
+ wandb.init(project='visualize-predictions', name='histograms')
+ _fibonacci = np.array([0, 1, 1, 2, 3, 5, 8, 13, 21, 34])
+ for i in range(1, 10):
+ wandb.log({'histograms': wandb.Histogram(_fibonacci / i)})
+ # Log a histogram on each step
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log images
+ """)
+ return
+
+
+@app.cell
+def _(plt, wandb):
+ wandb.init(project='visualize-predictions', name='images')
+ path_to_img = 'examples/examples/data/cafe.jpg'
+ # Initialize a new run
+ im = plt.imread(path_to_img)
+ wandb.log({'img': [wandb.Image(im, caption='Cafe')]})
+ # Generate an image
+ # Log the image
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log videos
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="videos")
+
+ # Generate a video
+ path_to_video = "examples/examples/data/openai-gym.mp4"
+
+ # Log the video
+ wandb.log({"video": wandb.Video(path_to_video, fps=4, format="gif")})
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log audio
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="audio")
+
+ # Generate audio data
+ path_to_audio = "examples/examples/data/piano.wav"
+
+ # Log that audio data
+ wandb.log({"examples":
+ [wandb.Audio(path_to_audio, caption="Piano", sample_rate=32)]})
+
+ wandb.finish()
+ return
+
+
+@app.cell
+def _(np, wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="audio")
+
+ # Generate audio data
+ fs = 44100 # sampling frequency, Hz
+ length = 3 # length, seconds
+ xs = np.linspace(0, length, num=fs * length)
+ waveform = np.sin(fs * 2 * np.pi / 40 * xs ** 2)
+
+ # Log audio data
+ wandb.log({"examples":
+ [wandb.Audio(waveform, caption="Boop", sample_rate=fs)]})
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log tables
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="tables")
+
+ # Create tabular data, method 1
+ data = [["I love my phone", "1", "1"],["My phone sucks", "0", "-1"]]
+ wandb.log({"a_table": wandb.Table(data=data, columns=["Text", "Predicted Label", "True Label"])})
+
+ # Create tabular data, method 2
+ table = wandb.Table(columns=["Text", "Predicted Label", "True Label"])
+ table.add_data("I love my phone", "1", "1")
+ table.add_data("My phone sucks", "0", "-1")
+ wandb.log({"another_table": table})
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log HTML
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="html")
+
+ # Generate HTML data
+ path_to_html = "examples/examples/data/some_html.html"
+
+ # Log an HTML file
+ wandb.log({"custom_file": wandb.Html(open(path_to_html))})
+
+ # Log raw HTML strings
+ wandb.log({"custom_string": wandb.Html('Link')})
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log 3D Objects
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="3d_objects")
+
+ # Generate 3D object data
+ path_to_obj = "examples/examples/data/wolf.obj"
+
+ # Log the 3D object
+ wandb.log({"3d_object": wandb.Object3D(open(path_to_obj))})
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log Point Clouds
+ """)
+ return
+
+
+@app.cell
+def _(np, wandb):
+ # Initialize a new run
+ wandb.init(project="visualize-predictions", name="point_clouds")
+
+ # Generate a cloud of points
+ points = np.random.uniform(size=(250, 3))
+
+ # Log points and boxes in W&B
+ wandb.log(
+ {
+ "point_scene": wandb.Object3D(
+ {
+ "type": "lidar/beta",
+ "points": points,
+ "boxes": np.array(
+ [
+ {
+ "corners": [
+ [0,0,0],
+ [0,1,0],
+ [0,0,1],
+ [1,0,0],
+ [1,1,0],
+ [0,1,1],
+ [1,0,1],
+ [1,1,1]
+ ],
+ "label": "Box",
+ "color": [123,321,111],
+ },
+ {
+ "corners": [
+ [0,0,0],
+ [0,2,0],
+ [0,0,2],
+ [2,0,0],
+ [2,2,0],
+ [0,2,2],
+ [2,0,2],
+ [2,2,2]
+ ],
+ "label": "Box-2",
+ "color": [111,321,0],
+ }
+ ]
+ ),
+ "vectors": np.array([])
+ }
+ )
+ }
+ )
+
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## More Resources
+ Check out some other cool things you can do with Weights & Biases:
+ * [Track model performance](https://app.wandb.ai/lavanyashukla/visualize-models/reports/Visualize-Model-Performance--Vmlldzo1NTk2MA)
+ * [Visualize sklearn models](https://app.wandb.ai/lavanyashukla/visualize-sklearn/reports/Visualize-Sklearn-Model-Performance--Vmlldzo0ODIzNg)
+ * [Visualize model predictions](https://app.wandb.ai/lavanyashukla/visualize-predictions/reports/Visualize-Model-Predictions--Vmlldzo1NjM4OA/)
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-logging-strategies-for-high-frequency-data/wandb_log_logging_strategies_for_high_frequency_data.py b/marimo/convert/wandb-log-logging-strategies-for-high-frequency-data/wandb_log_logging_strategies_for_high_frequency_data.py
new file mode 100644
index 00000000..9a329f5d
--- /dev/null
+++ b/marimo/convert/wandb-log-logging-strategies-for-high-frequency-data/wandb_log_logging_strategies_for_high_frequency_data.py
@@ -0,0 +1,427 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🪵 Logging Strategies for High-Frequency Data
+
+ Sometimes, you need to log metrics that are generated
+ **at high frequency**.
+ This might happen if you log:
+ - loss and metrics on _every batch_ during training,
+ - reward on _each step_ of each episode during simulation, or
+ - outputs, media, and metrics on _every input_ during analysis.
+
+ This can lead to substantial slowdown
+ if the logging calls,
+ which may need to write to disk or communicate over a network,
+ end up much slower than the iterations --
+ training will finish, but `wandb` will still be catching up.
+
+ This limitation is fundamental:
+ our experiments generate large quantities of information,
+ and if we need [lossless](https://en.wikipedia.org/wiki/Lossless_compression)
+ access to all of that information,
+ we will need to [pay a cost](https://en.wikipedia.org/wiki/Landauer%27s_principle).
+
+ But if we are clever, we can create a summary
+ that's much smaller but doesn't lose too much information, just as
+ [a JPEG or mp3 file](https://en.wikipedia.org/wiki/Lossy_compression)
+ summarizes a TIFF image or a WAV file into a smaller amount of disk space
+ while preserving its contents as much as possible.
+
+ In this notebook we'll cover how to use three
+ common summarization strategies with `wandb.log`,
+ along with their benefits and drawbacks:
+ 1. **downsampling**, which takes a (hopefully) representative subset,
+ 1. **summarization**, which uses descriptive statistics, like mean or median, and
+ 1. **batching**, which uses histograms.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ To get started, we install wandb and log in.
+
+ If this is your first time using wandb, you'll need to sign up as well.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # # install the package
+ # !pip install wandb -qq
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell
+def _():
+ import random
+
+ timing_header = "=" * 5 + " Timing Results " + "=" * 5
+ return random, timing_header
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🧲 Downsampling
+
+ The simplest solution is to just log less frequently.
+ It is also very common:
+ many logging libraries include an option like
+ [`log_every_n_steps`](https://pytorch-lightning.readthedocs.io/en/latest/extensions/logging.html#logging-frequency),
+ which triggers logging less frequently.
+
+ This is a basic form of [down-sampling a signal](https://en.wikipedia.org/wiki/Downsampling_(signal_processing)).
+ It has the advantage of being simple to implement.
+
+ It has two disadvantages. First,
+ [aliasing of high-frequency components](https://en.wikipedia.org/wiki/Aliasing)
+ like noise might make the logged quantities appear to oscillate.
+ This is typically not a serious issue for ML experiment metrics.
+ Second, all information about other iterations is completely lost.
+
+ > There's one special consideration
+ for all of these summary methods when logging with W&B:
+ if you have metrics being logged at different frequencies,
+ [the "step" value](https://docs.wandb.ai/guides/track/log#incremental-logging)
+ used as the x-axis for default charts
+ becomes less interpretable.
+ Even if you don't, it's a good idea to include timing information,
+ like batch index, epoch, or optimization step index,
+ when logging.
+ Below, we track the `iter`ation count of the main for loop.
+ """)
+ return
+
+
+@app.cell
+def _(random, timing_header, wandb):
+ # magic command not supported in marimo; please file an issue to add support
+ # %%time
+ _total_steps = 100000
+ _log_every = 500
+ _run = wandb.init(project='perf-log', name='downsample')
+ with _run:
+ random.seed(117)
+ _metric = 0
+ for _iter in range(_total_steps):
+ _metric += random.randint(-1, 1)
+ if (_iter + 1) % _log_every == 0:
+ _run.log({'metric': _metric, 'iter': _iter})
+ print('Run Finished!')
+ print('Logging Finished!')
+ print('\n' + timing_header) # subsampling
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The run and logging finish in a few seconds on commodity hardware.
+
+ At the bottom of the notebook,
+ there's a cell that logs the value of `metric`
+ on every iteration.
+ That cell takes roughly 25x longer to run.
+
+ The screenshot below compares the results logged by these two cells
+ in the W&B interface.
+ Note how small the differences are,
+ even though we threw out every 500th entry
+ -- this is typical for ML experiment metrics.
+ You can get away with throwing out a lot!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 📝 Summarization
+
+ Downsampling approaches our problem from a signal processing perspective.
+ From the perspective of statistics, on the other hand,
+ by calculating a single number to represent our metric data
+ we have chosen a
+ [summary statistic](https://en.wikipedia.org/wiki/Summary_statistics).
+
+ This point of view suggests an alternative approach:
+ use common statistical estimators, like the mean, median, or mode.
+ If our goal is to capture the "average" or "typical" behavior
+ of the metric over a period,
+ these summaries get
+ [closer to the truth](https://en.wikipedia.org/wiki/Minimum-variance_unbiased_estimator)
+ than does a randomly or intermittently chosen data point.
+
+ Below, we'll take the mean after each `log_every` steps.
+ As before, we make sure to log the `iter`ation count
+ along with the mean when we call `wandb.log`.
+ """)
+ return
+
+
+@app.cell
+def _(random, timing_header, wandb):
+ # magic command not supported in marimo; please file an issue to add support
+ # %%time
+ _total_steps = 100000
+ _log_every = 500
+ _run = wandb.init(project='perf-log', name='summarizing')
+ with _run:
+ random.seed(117)
+ _metric, running_average = (0, 0)
+ for _iter in range(_total_steps):
+ _metric += random.randint(-1, 1)
+ running_average += _metric / _log_every
+ if (_iter + 1) % _log_every == 0:
+ _run.log({'metric': running_average, 'iter': _iter})
+ running_average = 0
+ print('Run Finished!')
+ print('Logging Finished!')
+ print('\n' + timing_header)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Once again, we can compare the result to the true signal.
+ The differences are again small:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In signal processing,
+ what we did would be called a
+ [low-pass filter](https://en.wikipedia.org/wiki/Low-pass_filter).
+ In the terminology of deep neural networks,
+ it's a type of
+ [strided, one-dimensional average pooling](https://pytorch.org/docs/stable/generated/torch.nn.AvgPool1d.html#torch.nn.AvgPool1d).
+
+ Both of these give good intuition for the disadvantages of this method:
+ rapid changes are missed and the reported signal lags slightly.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👯 Batching with `wandb.Histogram`
+
+ But wait -- why stop at a single number?
+ Summary statistics and estimators come from early 20th century statistics,
+ which allowed strong modeling assumptions to make up for
+ small data and limited compute.
+
+ Contemporary ML, and some contemporary statistics,
+ operates in the opposite context.
+ We aim to presume as little as possible about our data
+ and tolerate laborious calculations,
+ which we outsource to computers.
+
+ So instead of selecting a single summary statistic,
+ as in classical parametric statistics,
+ why not use a technique from
+ [non-parametric statistics](https://en.wikipedia.org/wiki/Nonparametric_statistics)?
+
+ This may sound like an over-complication, but
+ [the humble histogram is actually a simple, non-parametric method](https://en.wikipedia.org/wiki/Nonparametric_statistics#Non-parametric_models)
+ for estimating distributions.
+ The histogram captures not only the average or typical behavior
+ of the metric at a given moment,
+ but also the variability and spread.
+
+ With `wandb`, using it is as easy as
+ swapping out the calculation of the mean
+ for a call to `wandb.Histogram`.
+ The resulting chart will also include a
+ [kernel density estimate](https://en.wikipedia.org/wiki/Kernel_density_estimation)
+ (or "KDE"), which is a fancier, continuous version of a histogram --
+ and also a method from non-parametric statistics.
+ """)
+ return
+
+
+@app.cell
+def _(random, timing_header, wandb):
+ # magic command not supported in marimo; please file an issue to add support
+ # %%time
+ _total_steps = 100000
+ _log_every = 500
+ _run = wandb.init(project='perf-log', name='batch')
+ with _run:
+ random.seed(117)
+ _metric, batch = (0, [])
+ for _iter in range(_total_steps):
+ _metric += random.randint(-1, 1)
+ batch.append(_metric)
+ if (_iter + 1) % _log_every == 0:
+ _run.log({'metric_hist': wandb.Histogram(batch), 'iter': _iter})
+ batch = []
+ print('Run Finished!')
+ print('Logging Finished!')
+ print('\n' + timing_header)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Though it takes a bit more effort to read,
+ the histogram chart captures far more than the
+ scalar-valued summary charts did.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_:
+ the histogram chart will only show up on the
+ [run page](https://docs.wandb.ai/ref/app/pages/run-page),
+ which collects metrics from individual runs,
+ not on the
+ [project page](https://docs.wandb.ai/ref/app/pages/project-page),
+ which compares metrics across runs.
+ A chart of histograms changing over time has enough going on already
+ before you start comparing histograms to each other!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Notice also that this run, logging histograms,
+ finishes in approximately the same amount of time
+ as do the runs which log scalars,
+ even though it seems like we're doing a lot more.
+ That's because here,
+ as in most contemporary computing settings,
+ the primary bottlenecks are in reading and writing files
+ and communicating across networks.
+
+ That's a good thing to keep in mind when
+ you're writing and profiling your ML code!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🐌 Log Everything
+
+ For purposes of comparison,
+ we also include the code to log the metric on every step,
+ without any summarization.
+
+ > Warning! This cell takes a few minutes to run.
+ We put it at the end for a reason.
+ """)
+ return
+
+
+@app.cell
+def _(random, timing_header, wandb):
+ # magic command not supported in marimo; please file an issue to add support
+ # %%time
+ _total_steps = 100000
+ _run = wandb.init(project='perf-log', name='log-everything')
+ with _run:
+ random.seed(117)
+ _metric = 0
+ for _iter in range(_total_steps):
+ _metric += random.randint(-1, 1)
+ _run.log({'metric': _metric, 'iter': _iter})
+ print('Run Finished!')
+ print('Logging Finished!')
+ print('\n' + timing_header)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-plot-precision-recall-curves-with-w-b/wandb_log_plot_precision_recall_curves_with_w_b.py b/marimo/convert/wandb-log-plot-precision-recall-curves-with-w-b/wandb_log_plot_precision_recall_curves_with_w_b.py
new file mode 100644
index 00000000..3b76779d
--- /dev/null
+++ b/marimo/convert/wandb-log-plot-precision-recall-curves-with-w-b/wandb_log_plot_precision_recall_curves_with_w_b.py
@@ -0,0 +1,276 @@
+# /// script
+# dependencies = ["tensorflow", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Plot Precision-Recall Curves with W&B
+
+ How to log [Precision-Recall curves](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.precision_recall_curve.html#sklearn.metrics.precision_recall_curve) with [Vega](https://vega.github.io/vega/docs/) in [Weights & Biases](https://www.wandb.com).
+
+ ## Method: wandb.plot.pr_curve()
+
+ - More info and customization details: [Plot Precision Recall Curves](https://wandb.ai/wandb/plots/reports/Plot-Precision-Recall-Curves--VmlldzoyNjk1ODY)
+ - More examples in this W&B project: [Custom Charts](https://app.wandb.ai/demo-team/custom-charts).
+
+ These are simple cases to explain the basics—you can build much more sophisticated custom charts with our powerful new query editor.
+
+ This Colab explores a transfer learning problem: finetuning InceptionV3 with ImageNet weights to identify 10 types of living things (birds, plants, insects, etc) from 10K photos 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.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup: Download data
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Note: **this stage might take a few minutes (~3.6GB of data)**. If you end up needing to rerun this cell, comment out the first capture line (change ```%%capture``` to ```#%%capture``` ) so you can respond to the prompt about re-downloading the dataset (and see the progress bar).
+
+ Download sample data: 10,000 training images and 2,000 validation images from the [iNaturalist dataset](https://github.com/visipedia/inat_comp), evenly distributed across 10 classes of living things like birds, insects, plants, and mammals (names given in Latin—so Aves, Insecta, Plantae, etc :). We will fine-tune a convolutional neural network already trained on ImageNet on this task: given a photo of a living thing, correctly classify it into one of the 10 classes.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !curl -SL https://storage.googleapis.com/wandb_datasets/nature_12K.zip > nature_12K.zip
+ # !unzip nature_12K.zip
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install dependencies
+
+ Install tensorflow and wandb; log in to wandb.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: tensorflow !pip install tensorflow -qqq
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training code
+
+ Feel free to try different values for "NUM_TRAIN" and "NUM_EPOCHS" below so you can see a variety of PR curves (generally better ones with more training examples/longer training time)
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # this determines the name of your wandb project, where all your
+ # runs will be loggeed
+ PROJECT_NAME = "custom_pr_curve"
+
+ # EXPERIMENT CONFIG
+ #---------------------------
+ # try changing the number of training examples
+ # to generate a range of different PR curves
+ NUM_TRAIN = 100 # try 500, 1000, 2000, or max 10000
+ NUM_EPOCHS = 1 # try 3, 5, or as many as you like
+
+ import numpy as np
+ from sklearn.metrics import precision_recall_curve, roc_curve
+ from sklearn.metrics import average_precision_score
+ from sklearn.preprocessing import label_binarize
+
+ from tensorflow.keras.applications.inception_v3 import InceptionV3
+ from tensorflow.keras.callbacks import Callback
+ from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
+ from tensorflow.keras.models import Model
+ from tensorflow.keras.preprocessing.image import ImageDataGenerator
+
+ from wandb.keras import WandbCallback
+
+ # local paths to data
+ train_data = "inaturalist_12K/train"
+ val_data = "inaturalist_12K/val"
+
+ # 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 10000 for the full dataset
+ # GOOD CONFIG TO TRY: 100, 500, 1000, 2000
+ "num_train" : NUM_TRAIN, # up to 10000,
+ # number of images used to validate--set low for demo training speed
+ # you can set this up to 2000 for the full dataset
+ "num_val" : 500, #2000,
+ "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
+ # you can set this up to 5, 10, or more for better results
+ # GOOD CONFIG TO TRY: 3, 5, 10
+ "pretrain_epochs" : NUM_EPOCHS, #5,
+ # number of validation data batches to use when computing metrics
+ # at the end of each epoch
+ "num_log_batches": 15
+ }
+
+ def build_model(fc_size, num_classes):
+ """Load InceptionV3 with ImageNet weights, freeze it,
+ and attach a finetuning top for this classification task"""
+ # load InceptionV3 as base
+ base = InceptionV3(weights="imagenet", include_top="False")
+ # freeze base layers
+ for layer in base.layers:
+ layer.trainable = False
+ x = base.get_layer('mixed10').output
+
+ # attach a fine-tuning layer
+ x = GlobalAveragePooling2D()(x)
+ x = Dense(fc_size, activation='relu')(x)
+ guesses = Dense(num_classes, activation='softmax')(x)
+
+ model = Model(inputs=base.input, outputs=guesses)
+ model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])
+ return model
+
+ def pretrain():
+ """ 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. subsequent training phase would unfreeze all the layers
+ and finetune the whole model on the new data"""
+ # track this experiment with wandb: all runs will be sent
+ # to the given project name
+ wandb.init(project=PROJECT_NAME, config=config_defaults)
+ cfg = wandb.config
+
+ # create train and validation data generators
+ train_datagen = ImageDataGenerator(
+ rescale=1. / 255,
+ shear_range=0.2,
+ zoom_range=0.2,
+ horizontal_flip=True)
+ val_datagen = ImageDataGenerator(rescale=1. / 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 = val_datagen.flow_from_directory(
+ val_data,
+ target_size=(cfg.img_width, cfg.img_height),
+ batch_size=cfg.batch_size,
+ class_mode='categorical')
+
+ # instantiate model and callbacks
+ model = build_model(cfg.fc_size, cfg.num_classes)
+ callbacks = [WandbCallback(), PRMetrics(val_generator, num_log_batches=15)]
+
+ # train!
+ model.fit(
+ train_generator,
+ steps_per_epoch = cfg.num_train // cfg.batch_size,
+ epochs=cfg.pretrain_epochs,
+ validation_data=val_generator,
+ callbacks = callbacks,
+ validation_steps=cfg.num_val // cfg.batch_size)
+
+ wandb.run.finish()
+
+ class PRMetrics(Callback):
+ """ Custom callback to compute per-class PR & ROC curves
+ at the end of each training epoch"""
+ def __init__(self, generator=None, num_log_batches=1):
+ self.generator = generator
+ self.num_batches = num_log_batches
+ # store full names of classes
+ self.class_names = { v: k for k, v in generator.class_indices.items() }
+ self.flat_class_names = [k for k, v in generator.class_indices.items()]
+
+ def on_epoch_end(self, epoch, logs={}):
+ # collect validation data and ground truth labels from generator
+ val_data, val_labels = zip(*(self.generator[i] for i 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)
+ ground_truth_class_ids = val_labels.argmax(axis=1)
+
+ # Log precision-recall curve
+ # the key "pr_curve" is the id of the plot--do not change
+ # this if you want subsequent runs to show up on the same plot
+ wandb.log({"pr_curve" : wandb.plot.pr_curve(ground_truth_class_ids,
+ val_predictions,
+ labels=self.flat_class_names)})
+
+ return (pretrain,)
+
+
+@app.cell
+def _(pretrain):
+ # run this cell to launch your experiment!
+ # charts will show up in your run page under the heading "Custom Charts",
+ # which you may need to click on to expand
+ pretrain()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-plot-roc-curves-with-w-b/wandb_log_plot_roc_curves_with_w_b.py b/marimo/convert/wandb-log-plot-roc-curves-with-w-b/wandb_log_plot_roc_curves_with_w_b.py
new file mode 100644
index 00000000..2f03dd42
--- /dev/null
+++ b/marimo/convert/wandb-log-plot-roc-curves-with-w-b/wandb_log_plot_roc_curves_with_w_b.py
@@ -0,0 +1,274 @@
+# /// script
+# dependencies = ["tensorflow", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Plot ROC Curves with W&B
+
+ How to log [ROC curves](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_curve.html#sklearn.metrics.roc_curve) with [Vega](https://vega.github.io/vega/docs/) in [Weights & Biases](https://www.wandb.com).
+
+ ## Method: wandb.plot.roc_curve()
+
+ - More info and customization details: [Plot ROC Curves](https://wandb.ai/wandb/plots/reports/Plot-ROC-Curves--VmlldzoyNjk3MDE)
+ - More examples in this W&B project: [Custom Charts](https://app.wandb.ai/demo-team/custom-charts).
+
+ These are simple cases to explain the basics—you can build much more sophisticated custom charts with our powerful new query editor.
+
+ This Colab explores a transfer learning problem: finetuning InceptionV3 with ImageNet weights to identify 10 types of living things (birds, plants, insects, etc) from 10K photos 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.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Setup: Download data
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Note: **this stage might take a few minutes (~3.6GB of data)**. If you end up needing to rerun this cell, comment out the first capture line (change ```%%capture``` to ```#%%capture``` ) so you can respond to the prompt about re-downloading the dataset (and see the progress bar).
+
+ Download sample data: 10,000 training images and 2,000 validation images from the [iNaturalist dataset](https://github.com/visipedia/inat_comp), evenly distributed across 10 classes of living things like birds, insects, plants, and mammals (names given in Latin—so Aves, Insecta, Plantae, etc :). We will fine-tune a convolutional neural network already trained on ImageNet on this task: given a photo of a living thing, correctly classify it into one of the 10 classes.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !curl -SL https://storage.googleapis.com/wandb_datasets/nature_12K.zip > nature_12K.zip
+ # !unzip nature_12K.zip
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Install dependencies
+
+ Install tensorflow and wandb; log in to wandb.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: tensorflow !pip install tensorflow -qqq
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training code
+
+ Feel free to try different values for "NUM_TRAIN" and "NUM_EPOCHS" below so you can see a variety of PR curves (generally better ones with more training examples/longer training time)
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ # this determines the name of your wandb project, where all your
+ # runs will be loggeed
+ PROJECT_NAME = "custom_roc_curve"
+
+ # EXPERIMENT CONFIG
+ #---------------------------
+ # try changing the number of training examples
+ # to generate a range of different PR curves
+ NUM_TRAIN = 100 # try 500, 1000, 2000, or max 10000
+ NUM_EPOCHS = 1 # try 3, 5, or as many as you like
+
+ import numpy as np
+ from sklearn.metrics import precision_recall_curve, roc_curve
+ from sklearn.metrics import average_precision_score
+ from sklearn.preprocessing import label_binarize
+
+ from tensorflow.keras.applications.inception_v3 import InceptionV3
+ from tensorflow.keras.callbacks import Callback
+ from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
+ from tensorflow.keras.models import Model
+ from tensorflow.keras.preprocessing.image import ImageDataGenerator
+
+ from wandb.keras import WandbCallback
+
+ # local paths to data
+ train_data = "inaturalist_12K/train"
+ val_data = "inaturalist_12K/val"
+
+ # 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 10000 for the full dataset
+ # GOOD CONFIG TO TRY: 100, 500, 1000, 2000
+ "num_train" : NUM_TRAIN, # up to 10000,
+ # number of images used to validate--set low for demo training speed
+ # you can set this up to 2000 for the full dataset
+ "num_val" : 500, #2000,
+ "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
+ # you can set this up to 5, 10, or more for better results
+ # GOOD CONFIG TO TRY: 3, 5, 10
+ "pretrain_epochs" : NUM_EPOCHS, #5,
+ # number of validation data batches to use when computing metrics
+ # at the end of each epoch
+ "num_log_batches": 15
+ }
+
+ def build_model(fc_size, num_classes):
+ """Load InceptionV3 with ImageNet weights, freeze it,
+ and attach a finetuning top for this classification task"""
+ # load InceptionV3 as base
+ base = InceptionV3(weights="imagenet", include_top="False")
+ # freeze base layers
+ for layer in base.layers:
+ layer.trainable = False
+ x = base.get_layer('mixed10').output
+
+ # attach a fine-tuning layer
+ x = GlobalAveragePooling2D()(x)
+ x = Dense(fc_size, activation='relu')(x)
+ guesses = Dense(num_classes, activation='softmax')(x)
+
+ model = Model(inputs=base.input, outputs=guesses)
+ model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])
+ return model
+
+ def pretrain():
+ """ 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. subsequent training phase would unfreeze all the layers
+ and finetune the whole model on the new data"""
+ # track this experiment with wandb: all runs will be sent
+ # to the given project name
+ wandb.init(project=PROJECT_NAME, config=config_defaults)
+ cfg = wandb.config
+
+ # create train and validation data generators
+ train_datagen = ImageDataGenerator(
+ rescale=1. / 255,
+ shear_range=0.2,
+ zoom_range=0.2,
+ horizontal_flip=True)
+ val_datagen = ImageDataGenerator(rescale=1. / 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 = val_datagen.flow_from_directory(
+ val_data,
+ target_size=(cfg.img_width, cfg.img_height),
+ batch_size=cfg.batch_size,
+ class_mode='categorical')
+
+ # instantiate model and callbacks
+ model = build_model(cfg.fc_size, cfg.num_classes)
+ callbacks = [WandbCallback(), PRMetrics(val_generator, num_log_batches=15)]
+
+ # train!
+ model.fit(
+ train_generator,
+ steps_per_epoch = cfg.num_train // cfg.batch_size,
+ epochs=cfg.pretrain_epochs,
+ validation_data=val_generator,
+ callbacks = callbacks,
+ validation_steps=cfg.num_val // cfg.batch_size)
+
+ wandb.run.finish()
+
+ class PRMetrics(Callback):
+ """ Custom callback to compute per-class PR & ROC curves
+ at the end of each training epoch"""
+ def __init__(self, generator=None, num_log_batches=1):
+ self.generator = generator
+ self.num_batches = num_log_batches
+ # store full names of classes
+ self.class_names = { v: k for k, v in generator.class_indices.items() }
+ self.flat_class_names = [k for k, v in generator.class_indices.items()]
+
+ def on_epoch_end(self, epoch, logs={}):
+ # collect validation data and ground truth labels from generator
+ val_data, val_labels = zip(*(self.generator[i] for i 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)
+ ground_truth_class_ids = val_labels.argmax(axis=1)
+
+ # Log precision-recall curve
+ # the key "pr_curve" is the id of the plot--do not change
+ # this if you want subsequent runs to show up on the same plot
+ wandb.log({"roc_curve" : wandb.plot.roc_curve(ground_truth_class_ids, val_predictions, labels=self.flat_class_names)})
+
+ return (pretrain,)
+
+
+@app.cell
+def _(pretrain):
+ # run this cell to launch your experiment!
+ # charts will show up in your run page under the heading "Custom Charts",
+ # which you may need to click on to expand
+ pretrain()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-run-names-visualized-using-min-dalle/wandb_log_run_names_visualized_using_min_dalle.py b/marimo/convert/wandb-log-run-names-visualized-using-min-dalle/wandb_log_run_names_visualized_using_min_dalle.py
new file mode 100644
index 00000000..194937b2
--- /dev/null
+++ b/marimo/convert/wandb-log-run-names-visualized-using-min-dalle/wandb_log_run_names_visualized_using_min_dalle.py
@@ -0,0 +1,136 @@
+# /// script
+# dependencies = ["flax", "torch", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Visualize W&B Run Names using Craiyon
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We use min-dalle, https://github.com/kuprel/min-dalle, a repo with the bare essentials necessary for doing inference on the Craiyon model. We use `wandb.Api` to get the project and run names from your account.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Setup
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #@title Setup
+ #! git clone --depth 1 https://github.com/kuprel/min-dalle
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/kuprel/min-dalle'])
+ #! git lfs install
+ subprocess.call(['git', 'lfs', 'install'])
+ #! git clone https://huggingface.co/dalle-mini/vqgan_imagenet_f16_16384 /content/min-dalle/pretrained/vqgan
+ subprocess.call(['git', 'clone', 'https://huggingface.co/dalle-mini/vqgan_imagenet_f16_16384', '/content/min-dalle/pretrained/vqgan'])
+ # packages added via marimo's package management: torch flax==0.4.2 wandb !pip install torch flax==0.4.2 wandb
+ #! wandb artifact get --root=/content/min-dalle/pretrained/dalle_bart_mini dalle-mini/dalle-mini/mini-1:v0
+ subprocess.call(['wandb', 'artifact', 'get', '--root=/content/min-dalle/pretrained/dalle_bart_mini', 'dalle-mini/dalle-mini/mini-1:v0'])
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Generate Images using W&B Run Names
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import os
+ os.chdir('/content/min-dalle')
+ from min_dalle.min_dalle_torch import MinDalleTorch
+ import ipywidgets as widgets
+ from IPython.display import display, clear_output
+
+ mega = False
+ model = MinDalleTorch(mega)
+ seed = 7
+ api = wandb.Api()
+ projects = [project.name for project in api.projects()]
+
+ project_dropdown = widgets.Dropdown(
+ options = projects,
+ description = 'Projects:',
+ )
+
+ run_dropdown = widgets.Dropdown(
+ options=[run.name for run in api.runs(project_dropdown.value)],
+ description = 'Run names:',
+ )
+
+ image_output = widgets.Output()
+
+ def on_project_value_change(change):
+ run_dropdown.options = [run.name for run in api.runs(project_dropdown.value)]
+
+ button = widgets.Button(
+ description='Generate',
+ tooltip='Click me to create an image from the currently selected run name',
+ )
+ def generate(b):
+ button.disabled = True
+ run_name = run_dropdown.value.replace('-', ' ')
+ image = model.generate_image(run_name, seed=seed)
+ with image_output:
+ clear_output(wait=True)
+ display(image, run_name)
+ button.disabled = False
+ button.on_click(generate)
+
+ project_dropdown.observe(on_project_value_change, 'value')
+ widgets.VBox([project_dropdown, run_dropdown, button, image_output])
+ return (wandb,)
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-saving-code-with-w-b/wandb_log_saving_code_with_w_b.py b/marimo/convert/wandb-log-saving-code-with-w-b/wandb_log_saving_code_with_w_b.py
new file mode 100644
index 00000000..cf9a1b68
--- /dev/null
+++ b/marimo/convert/wandb-log-saving-code-with-w-b/wandb_log_saving_code_with_w_b.py
@@ -0,0 +1,214 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👩🏽💻 Never Lose Track of What Code You Ran Ever Again
+
+ It happens to the best of us:
+ chasing down a good idea for your model,
+ you run the Jupyter cells out of order
+ or you forget to commit while trying several different models.
+ The code is unrecoverable, but the accuracy
+ was higher than anything you've seen!
+ **A really cool result has been severed
+ from the code used to generate it**,
+ and it might as well have been told
+ to you by [Mr. Snuffleupagus](https://muppet.fandom.com/wiki/Mr._Snuffleupagus)
+ or [The Great Gazoo](https://en.wikipedia.org/wiki/The_Great_Gazoo).
+
+ With Weights & Biases, you **won't need to worry** about that happening again!
+ We'll **save the code you ran along with results and hyperparameters**, all in a centralized location with easy comparison and visualization tools.
+ Even better, in a Jupyter notebbok this feature **tracks all the cells you executed!**
+
+ Here is a simple implementation where we simulate logging some metrics and the code that generates them -- the same procedure works for both notebooks and Python scripts.
+
+ _Note_: We don't train any models here 😔. If you'd like to see how W&B integrates with other tools, check out
+ one of our integration demo
+ [colabs](https://github.com/wandb/examples/tree/master/colabs)
+ or [videos](https://www.youtube.com/playlist?list=PLD80i8An1OEGajeVo15ohAQYF1Ttle0lk).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🚀 Installs and Imports
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install -q wandb
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import math
+ import random
+
+ return math, random, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 🙅♂️ Logging without Saving Code
+
+ First, let's review how W&B logging works without code saving.
+
+ We need to `init`ialize a `run`, a unit of computation
+ (model training, data preprocessing, etc.)
+ in our `project`.
+ That's also where we set up the `config`uration of the `run`,
+ e.g. hyperparameters like learning rate.
+
+ Note here we specify the argument `save_code` as `False`, to specifically say that we do not need to save the code for the run.
+
+ Then, we'll wrap our training `for` loop in a `with` block,
+ ensuring our run closes out when training finishes,
+ and `log` our metrics.
+ """)
+ return
+
+
+@app.cell
+def _(math, random, wandb):
+ _run = wandb.init(project='code_save', config={'hyperparameter': 4}, save_code=False)
+ with _run:
+ for _step in range(100):
+ wandb.log({'acc': math.log(0.1 + random.random() + _step * 0.01), 'val_acc': math.log(0.1 + random.random() + _step * 0.01), 'loss': wandb.config.hyperparameter - math.log(0.1 + random.random() + _step * 0.01), 'val_loss': wandb.config.hyperparameter - math.log(0.1 + random.random() + _step * 0.01)}) # insert training process here
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📊 Dashboard Screenshot
+
+ Here is an example of what the dashboard looks like
+ after running the snippet above.
+
+ Metrics and hyperparameters are logged, but not the code!
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 💾 Logging Metrics and Saving Code
+
+ Adding code saving is easy: we just pass the argument
+ `save_code=True` to `wandb.init`. That's it!
+
+ _Hot Tip_: If you don't want to worry about setting this on every project,
+ just change your default on the [settings page](https://wandb.ai/settings), as below:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell
+def _(math, random, wandb):
+ _run = wandb.init(project='code_save', config={'hyperparameter': 4}, save_code=True)
+ with _run:
+ for _step in range(100):
+ wandb.log({'acc': math.log(0.1 + random.random() + _step * 0.01), 'val_acc': math.log(0.1 + random.random() + _step * 0.01), 'loss': wandb.config.hyperparameter - math.log(0.1 + random.random() + _step * 0.01), 'val_loss': wandb.config.hyperparameter - math.log(0.1 + random.random() + _step * 0.01)}) # insert training process here
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 📊 Dashboard Screenshot
+ In the left-hand panel on the dashboard, we can see a new icon pop up: `{}`.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If you click it, you'll see the entire history of the notebook session! That includes the code we ran for the section without code logging.
+
+ W&B also automatically catches the standard out and standard error,
+ plus system metrics!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 👨🏿🏫 Learn more about W&B
+ That's not all!
+ Check out the links below to learn how to use W&B to ...
+ - [log rich media like audio, video, and 3D point clouds](http://wandb.me/media-colab)
+ - [visualize datasets and model predictions](http://wandb.me/dsviz-nature-colab)
+ - [coordinate hyperparameter sweeps](http://wandb.me/sweeps-colab)
+
+ ...or check out our
+ [repository of examples](https://github.com/wandb/examples)
+ for even more features!
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-log-set-alerts-with-w-b/wandb_log_set_alerts_with_w_b.py b/marimo/convert/wandb-log-set-alerts-with-w-b/wandb_log_set_alerts_with_w_b.py
new file mode 100644
index 00000000..2d1472ad
--- /dev/null
+++ b/marimo/convert/wandb-log-set-alerts-with-w-b/wandb_log_set_alerts_with_w_b.py
@@ -0,0 +1,135 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+ # Using `wandb.alert()` to Send Alert Messages
+
+ Use W&B Alerts to send yourself a Slack message or email when something happens in your Python script. Follow the steps below to send your first Alert. See the [Alerts docs](https://docs.wandb.com/app/features/alerts) for a more detailed description.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 1. Turn on Alerts in your W&B User Settings
+ - Go to your **[User Settings](https://wandb.ai/settings)**
+ - Turn on **Scriptable Alerts**
+ - Select whether you'd like to get Alerts via email or Slack
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 2. Install the W&B Library and Login
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ # Log in to your W&B account
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 3. Launch a script that triggers alerts
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ import random
+ from wandb import AlertLevel
+
+ # Initialize a new run in Weights & Biases
+ wandb.init(project="test_alerts",
+ config={
+ "threshold": 0.3, # The minimum acceptable accuracy
+ "max_steps": 1000, # The max number of steps for this run
+ })
+ config = wandb.config
+
+ # Simulating a model training loop
+ for training_step in range(config.max_steps):
+
+ # Generate a random number for accuracy
+ accuracy = round(random.random() + random.random(), 3)
+ wandb.log({"Accuracy": accuracy})
+
+ # If the accuracy is below the threshold, fire an alert and stop the run
+ if accuracy <= config.threshold:
+ wandb.alert(
+ title='Low Accuracy',
+ text=f'Accuracy {accuracy} at step {training_step} is below the acceptable theshold',
+ level=AlertLevel.WARN,
+ wait_duration=5
+ )
+ print(f"Script stopped as accuracy is below threshold, {accuracy} vs {config.threshold}")
+ break
+
+ # Mark the run as finished (useful in Jupyter notebooks)
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 4. Check your Slack or email
+
+ Check your Slack or emails for the alert message. If you didn't receive any, check your [Settings](https://wandb.ai/settings) to make sure you've got emails or Slack turned on for **Scriptable Alerts**. More details in the
+ [Alerts docs](https://docs.wandb.com/app/features/alerts).
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-model-registry-ensembles-in-the-w-b-model-registry/wandb_model_registry_ensembles_in_the_w_b_model_registry.py b/marimo/convert/wandb-model-registry-ensembles-in-the-w-b-model-registry/wandb_model_registry_ensembles_in_the_w_b_model_registry.py
new file mode 100644
index 00000000..39af856c
--- /dev/null
+++ b/marimo/convert/wandb-model-registry-ensembles-in-the-w-b-model-registry/wandb_model_registry_ensembles_in_the_w_b_model_registry.py
@@ -0,0 +1,353 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Storing Ensemble Models with W&B Model Registry
+
+ This colab assumes you are familiar with Artifact, Collections, and the Model registry. Please see https://docs.wandb.ai/guides/models for details on such topics
+
+ `ensemble` - noun: "a group of items viewed as a whole rather than individually."
+
+ Ensemble Modeling is the class of solutions which have multiple sub-models which comprise a composite "model" of prediction. It is helpful to think of an Ensemble Model as a DAG (Directed Acyclical Graph) of models. A common pattern is to have individual engineers or teams working on different pieces of the composite, especially if each piece is complex. For example,
+
+ 
+
+ Now, there are many considersations and challenges with training such models:
+ - Tracking and versioning training data for child models generated by the parent models
+ - Evaluating & optimizing sub-components in isolation
+ - Evaluating & fine tuning the composite ensemble
+ - Facilitating independent workflows for each piece
+
+ ----
+
+ In the interest of keeping the code simple and runtimes quick, we will use dummy data, models, and training functions so that the reader can focus on the mechanics / workflow.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ import os
+ PROJECT = "ensemble_example_dev_3"
+ return (PROJECT,)
+
+
+@app.cell
+def _():
+ ## Dummy Functions:
+ import random
+ import json
+ import wandb
+
+ def _dummy_return():
+ return {'val': random.random()}
+
+ def make_artifact(obj, name, typename, file, run, log=True):
+ _art = wandb.Artifact(name, typename)
+ with _art.new_file(file) as file:
+ file.write(json.dumps(obj))
+ if log:
+ _run.log_artifact(_art)
+ return _art
+
+ def load_obj_from_artifact(run, name, file):
+ _art = _run.use_artifact(name)
+ ref = _art.get_path(file)
+ with open(ref.download(), 'r') as file:
+ obj = json.load(file)
+ return (obj, ref)
+
+ def generate_data():
+ return _dummy_return()
+
+ def train_m1(data):
+ return _dummy_return()
+
+ def m1_predict(m1, data):
+ return _dummy_return()
+
+ def train_m2(data):
+ return _dummy_return()
+
+ def m2_predict(m2, data):
+ return _dummy_return()
+
+ def prepare_m3_data(m1_data, m2_data):
+ return _dummy_return()
+
+ def train_m3(m3_data):
+ return _dummy_return()
+
+ def m3_predict(m3, m3_data):
+ return _dummy_return()
+
+ def make_ensemble_model(m1, m2, m3):
+ return _dummy_return()
+
+ def ensemble_predict(ensemble, data):
+ return _dummy_return()
+
+ return (
+ generate_data,
+ load_obj_from_artifact,
+ m1_predict,
+ m2_predict,
+ make_artifact,
+ make_ensemble_model,
+ prepare_m3_data,
+ train_m1,
+ train_m2,
+ train_m3,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 0: Create Registered Models
+ Before we train our models, we will want to create Registered Models for **each of our sub models as well as one for the entire ensemble.**
+ - Navigate to the model registry and click "Create Registered Model".
+ - Set the entity to the team you want this registered model to be visible to and the project to "model-registry". Do this for:
+ - ensemble_m1
+ - ensemble_m2
+ - ensemble_m3
+ - ensemble_all
+
+ See screenshots below:
+ 
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 2: Train each piece of the ensemble and link them to the Registry
+ - This can happen independently. The benefit of the registry is tracking the state of independent workstreams asynchronously.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Generate Data
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, generate_data, make_artifact, wandb):
+ ## Run 1: Generate Data:
+ with wandb.init(project=PROJECT, job_type='generate_data') as _run:
+ _data = generate_data()
+ make_artifact(_data, 'example_data', 'dataset', 'data.json', _run)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train M1
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, load_obj_from_artifact, make_artifact, train_m1, wandb):
+ ## Run 2: Train M1
+ with wandb.init(project=PROJECT, job_type='train_m1') as _run:
+ _data, _ = load_obj_from_artifact(_run, 'example_data:latest', 'data.json')
+ for _ in range(3): # Let's train 3 epochs
+ _m1 = train_m1(_data)
+ _art = make_artifact(_m1, f'm1-{_run.id}', 'model', 'model.json', _run)
+ _run.link_artifact(_art, 'ensemble_m1', ['latest']) # And we will link our last epoch to the m1_collection:
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ After the above two cells, W&B will have artifacts for both a dataset and a model, and the implicit DAG will begin to grow:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train M2
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, load_obj_from_artifact, make_artifact, train_m2, wandb):
+ with wandb.init(project=PROJECT, job_type='train_m2') as _run:
+ _data, _ = load_obj_from_artifact(_run, 'example_data:latest', 'data.json')
+ for _ in range(3): # Let's train 3 epochs
+ _m2 = train_m2(_data)
+ _art = make_artifact(_m2, f'm2-{_run.id}', 'model', 'model.json', _run)
+ _run.link_artifact(_art, 'ensemble_m2', ['latest']) # And we will link our last epoch to the m2_collection
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Train M3
+ """)
+ return
+
+
+@app.cell
+def _(
+ PROJECT,
+ load_obj_from_artifact,
+ m1_predict,
+ m2_predict,
+ make_artifact,
+ prepare_m3_data,
+ train_m3,
+ wandb,
+):
+ ## Run 4: Train M3
+ with wandb.init(project=PROJECT, job_type='train_m3') as _run:
+ _data, _ = load_obj_from_artifact(_run, 'example_data:latest', 'data.json')
+ _m1, _ = load_obj_from_artifact(_run, 'ensemble_m1:latest', 'model.json')
+ _m2, _ = load_obj_from_artifact(_run, 'ensemble_m2:latest', 'model.json')
+ m3_data = prepare_m3_data(m1_predict(_m1, _data), m2_predict(_m1, _data)) # here, we reference the collection
+ make_artifact(m3_data, 'intermedia_m3_data', 'dataset', 'data.json', _run)
+ for _ in range(3): # here, we reference the collection
+ _m3 = train_m3(m3_data)
+ _art = make_artifact(_m3, f'm3-{_run.id}', 'model', 'model.json', _run)
+ _run.link_artifact(_art, 'ensemble_m3', ['latest']) # let's make the input data for the model (and save it for good measure) # Let's train 3 epochs # And we will link our last epoch to the m3_collection:
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Step 3: Assemble the Ensemble and link to Registry
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Assemble the Ensemble
+ You can create composite artifacts by adding references to other W&B artifacts into a single artifact. This allows you to create complex dependency structures among model collections.
+
+ `ref = wandb.Artifact.get_path(file)` retrieves an internal reference uri to the W&B artifact. We can pass this reference around to other collections that might be using this model in some way.
+
+ In this example, `load_obj_from_artifact` gives us the model and its W&B reference uri, which we then add to the composite ensemble artifact.
+ ```
+ def load_obj_from_artifact(run, name, file):
+ art = run.use_artifact(name)
+ ref = art.get_path(file)
+ with open(ref.download(), 'r') as file:
+ obj = json.load(file)
+ return obj, ref
+ ```
+ """)
+ return
+
+
+@app.cell
+def _(
+ PROJECT,
+ load_obj_from_artifact,
+ make_artifact,
+ make_ensemble_model,
+ wandb,
+):
+ ## Run 5: Assemble the Ensemble
+ with wandb.init(project=PROJECT, job_type='ensemble_assembly') as _run:
+ _m1, m1_ref = load_obj_from_artifact(_run, 'ensemble_m1:latest', 'model.json')
+ _m2, m2_ref = load_obj_from_artifact(_run, 'ensemble_m2:latest', 'model.json') # here, we reference the collection
+ _m3, m3_ref = load_obj_from_artifact(_run, 'ensemble_m3:latest', 'model.json')
+ ensemble = make_ensemble_model(_m1, _m2, _m3) # here, we reference the collection
+ _art = make_artifact(ensemble, f'all-{_run.id}', 'model', 'model.json', _run, False)
+ _art.add_reference(m1_ref, 'm1.model.json') # here, we reference the collection
+ _art.add_reference(m2_ref, 'm2.model.json')
+ _art.add_reference(m3_ref, 'm3.model.json')
+ _run.log_artifact(_art)
+ _run.link_artifact(_art, 'ensemble_all', ['latest'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Accessing Ensemble Pieces
+ Specific sub models from runs can be accessed with:
+ - `run.use("sub_model_collection_name-RUN_ID:ALIAS")`
+
+ Specific sub models can be accessed with:
+ - `run.use("sub_model_collection_name:ALIAS")`
+
+ Specific composite ensembles can be accessed with:
+ - `run.use("ensemble_name:ALIAS")`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Finally, W&B reflects the underlying structure of your models:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-model-registry-model-registry-e2e/wandb_model_registry_model_registry_e2e.py b/marimo/convert/wandb-model-registry-model-registry-e2e/wandb_model_registry_model_registry_e2e.py
new file mode 100644
index 00000000..3fdc1a98
--- /dev/null
+++ b/marimo/convert/wandb-model-registry-model-registry-e2e/wandb_model_registry_model_registry_e2e.py
@@ -0,0 +1,592 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Model Registry Tutorial
+ The model registry is a central place to house and organize all the model tasks and their associated artifacts being worked on across an org:
+ - Model checkpoint management
+ - Document your models with rich model cards
+ - Maintain a history of all the models being used/deployed
+ - Facilitate clean hand-offs and stage management of models
+ - Tag and organize various model tasks
+ - Set up automatic notifications when models progress
+
+ This tutorial will walkthrough how to track the model development lifecycle for a simple image classification task.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## 🛠️ Install `wandb`
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install -q wandb
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Login to W&B
+ - You can explicitly login using `wandb login` or `wandb.login()` (See below)
+ - Alternatively you can set environment variables. There are several env variables which you can set to change the behavior of W&B logging. The most important are:
+ - `WANDB_API_KEY` - create a new API key in your "Settings" section under your profile at [wandb.ai/settings](https://wandb.ai/settings)
+ - `WANDB_BASE_URL` - this is the url of the W&B server
+ - Create a new API key in "Profile" -> "Settings" in the W&B App. Store your API key securely. It can only be viewed once when created.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ # Login to W&B
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log Data and Model Checkpoints as Artifacts
+ W&B Artifacts allows you to track and version arbitrary serialized data (e.g. datasets, model checkpoints, evaluation results). When you create an artifact, you give it a name and a type, and that artifact is forever linked to the experimental system of record. If the underlying data changes, and you log that data asset again, W&B will automatically create new versions through checksummming its contents. W&B Artifacts can be thought of as a lightweight abstraction layer on top of shared unstructured file systems.
+
+ ### Anatomy of an artifact
+
+ The `Artifact` class will correspond to an entry in the W&B Artifact registry. The artifact has
+ * a name
+ * a type
+ * metadata
+ * description
+ * files, directory of files, or references
+
+ Example usage:
+ ```
+ run = wandb.init(project = "my-project")
+ artifact = wandb.Artifact(name = "my_artifact", type = "data")
+ artifact.add_file("/path/to/my/file.txt")
+ run.log_artifact(artifact)
+ run.finish()
+ ```
+
+ In this tutorial, the first thing we will do is download a training dataset and log it as an artifact to be used downstream in the training job.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ import sys
+ from pathlib import Path
+
+ # FORM VARIABLES
+ PROJECT_NAME = "model-registry-tutorial"
+ ENTITY = wandb.api.default_entity # replace with your Team name or username
+
+ # Dataset constants
+ DATASET_NAME = "nature_100"
+ DATA_DIR = (Path(sys.path[0]) / "data")
+ DATA_DIR.mkdir(parents=True, exist_ok=True)
+ DATA_SRC = DATA_DIR / DATASET_NAME
+ IMAGES_PER_LABEL = 10
+ BALANCED_SPLITS = {"train" : 8, "val" : 1, "test": 1}
+ MODEL_TYPE = "squeezenet"
+ return DATASET_NAME, DATA_DIR, DATA_SRC, ENTITY, PROJECT_NAME, Path
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Let's grab a version of our Dataset
+ """)
+ return
+
+
+@app.cell
+def _(DATASET_NAME, DATA_DIR):
+ import requests, zipfile, io
+
+ # Download the dataset from a bucket
+ src_url = f"https://storage.googleapis.com/wandb_datasets/{DATASET_NAME}.zip"
+ src_zip = f"{DATASET_NAME}.zip"
+
+ # Download the zip file from W&B
+ r = requests.get(src_url)
+
+ # Create a file object using the string data
+ z = zipfile.ZipFile(io.BytesIO(r.content))
+ z.extractall(path=DATA_DIR)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We are going to generate a file containing the image
+ """)
+ return
+
+
+@app.cell
+def _(DATASET_NAME, DATA_SRC, ENTITY, PROJECT_NAME, wandb):
+ with wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type='log_datasets') as run:
+ train_art = wandb.Artifact(name=DATASET_NAME,
+ type='raw_images',
+ description='nature image dataset with 10 classes, 10 images per class')
+ train_art.add_dir(DATA_SRC)
+ wandb.log_artifact(train_art)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Using Artifact names and aliases to easily hand-off and abstract data assets
+ - By simply referring to the `name:alias` combination of a dataset or model, we can better standardize components of a workflow
+ - For instance, you can build PyTorch `Dataset`'s or `DataModule`'s which take as arguments W&B Artifact names and aliases to load appropriately
+
+ You can now see all the metadata associated with this dataset, the W&B runs consuming it, and the whole lineage of upstream and downstream artifacts!
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _(DATA_SRC, Path, wandb):
+ import math
+
+ from PIL import Image
+ import torch
+ from torchvision import transforms, models
+ from torch.utils.data import Dataset, DataLoader, random_split
+
+ class NatureDataset(Dataset):
+ def __init__(self, artifact_name_alias: str, transform=None):
+ self.transform = transform
+
+ # Pull down the artifact locally to load it into memory
+ art = wandb.use_artifact(artifact_name_alias)
+ self.path_at = Path(art.download())
+
+ self.img_paths = list(DATA_SRC.rglob("*.jpg"))
+ labels = [image_path.parent.name for image_path in self.img_paths]
+ self.class_names = sorted(set(labels))
+ self.idx_to_class = {k: v for k, v in enumerate(self.class_names)}
+ self.class_to_idx = {v: k for k, v in enumerate(self.class_names)}
+
+ def __len__(self):
+ return len(self.img_paths)
+
+ def __getitem__(self, idx):
+ if torch.is_tensor(idx):
+ idx = idx.tolist()
+
+ image_path = Path(self.path_at) / self.img_paths[idx]
+
+ image = Image.open(image_path)
+ label = image_path.parent.name
+ label = torch.tensor(self.class_to_idx[label], dtype=torch.long)
+
+ if self.transform:
+ image = self.transform(image)
+
+ return image, label
+
+ class Dataloaders:
+ def __init__(self,
+ artifact_name_alias: str,
+ batch_size: int,
+ input_size: int,
+ seed: int = 42):
+ self.artifact_name_alias = artifact_name_alias
+ self.batch_size = batch_size
+ self.input_size = input_size
+ self.seed = seed
+
+ tfms = transforms.Compose([transforms.ToTensor(),
+ transforms.CenterCrop(self.input_size),
+ transforms.Normalize((0.485, 0.456, 0.406),
+ (0.229, 0.224, 0.225))])
+
+ print(f"Setting up data from artifact: {self.artifact_name_alias}")
+ self.dataset = NatureDataset(artifact_name_alias=self.artifact_name_alias,
+ transform=tfms)
+
+ nature_length = len(self.dataset)
+ train_size = math.floor(0.8 * nature_length)
+ val_size = math.floor(0.2 * nature_length)
+ print(f"Splitting dataset into {train_size} training samples and {val_size} validation samples")
+ self.ds_train, self.ds_valid = random_split(
+ self.dataset,
+ [train_size, val_size],
+ generator=torch.Generator().manual_seed(self.seed))
+
+ self.train = DataLoader(self.ds_train, batch_size=self.batch_size)
+ self.valid = DataLoader(self.ds_valid, batch_size=self.batch_size)
+
+ return models, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Model Training
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Writing the Model Class and Validation Function
+ """)
+ return
+
+
+@app.cell
+def _(models, torch, wandb):
+ import torch.nn.functional as F
+ import torch.optim as optim
+ from torch.optim.lr_scheduler import StepLR
+
+ def set_parameter_requires_grad(model, feature_extracting):
+ if feature_extracting:
+ for param in model.parameters():
+ param.requires_grad = False
+
+ def initialize_model(num_classes, feature_extract, use_pretrained=True):
+ "Create a model from torchvision.models"
+ model_ft = None
+
+ # SqueezeNet
+ model_ft = models.squeezenet1_0(pretrained=use_pretrained)
+ set_parameter_requires_grad(model_ft, feature_extract)
+ model_ft.classifier[1] = torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))
+ model_ft.num_classes = num_classes
+
+ return model_ft, 224
+
+ class NaturePyTorchModule(torch.nn.Module):
+ def __init__(self,
+ model_name,
+ num_classes=10,
+ feature_extract=True,
+ lr=0.01):
+ '''method used to define our model parameters'''
+ super().__init__()
+
+ self.model_name = model_name
+ self.num_classes = num_classes
+ self.feature_extract = feature_extract
+ self.lr = lr
+ self.model, self.input_size = initialize_model(num_classes=self.num_classes,
+ feature_extract=True)
+
+ def forward(self, x):
+ '''method used for inference input -> output'''
+ return self.model(x)
+
+ def evaluate_model(model, val_dl, idx_to_class, class_names):
+ device = torch.device("cpu")
+ model.eval()
+ test_loss = 0
+ correct = 0
+ preds = []
+ actual = []
+
+ val_table = wandb.Table(columns=['pred', 'actual', 'image'])
+
+ with torch.no_grad():
+ for data, target in val_dl:
+ data, target = data.to(device), target.to(device)
+ output = model(data)
+ test_loss += F.cross_entropy(
+ output, target, reduction="sum"
+ ).item() # sum up batch loss
+ pred = output.argmax(
+ dim=1, keepdim=True
+ ) # get the index of the max log-probability
+ preds += list(pred.flatten().tolist())
+ actual += target.numpy().tolist()
+ correct += pred.eq(target.view_as(pred)).sum().item()
+
+ for idx, img in enumerate(data):
+ img = img.numpy().transpose(1, 2, 0)
+ pred_class = idx_to_class[pred.numpy()[idx][0]]
+ target_class = idx_to_class[target.numpy()[idx]]
+ val_table.add_data(pred_class, target_class, wandb.Image(img))
+
+ test_loss /= len(val_dl.dataset)
+ accuracy = 100.0 * correct / len(val_dl.dataset)
+ conf_mat = wandb.plot.confusion_matrix(y_true=actual, preds=preds, class_names=class_names)
+ return test_loss, accuracy, preds, val_table, conf_mat
+
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Tracking the Training Loop
+ During training, it is a best practice to checkpoint your models overtime, so if training gets interrupted or your instance crashes you can resume from where you left off. With artifact logging, we can track all our checkpoints with W&B and attach any metadata we want (like format of serialization, class labels, etc.). That way, when someone needs to consume a checkpoint they know how to use it. When logging models of any form as artifacts, ensure to set the `type` of the artifact to `model`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%wandb -h 600
+ #
+ # run = wandb.init(project=PROJECT_NAME,
+ # entity=ENTITY,
+ # job_type='training',
+ # config={'model_type': MODEL_TYPE,
+ # 'lr': 1.0,
+ # 'gamma': 0.75,
+ # 'batch_size': 16,
+ # 'epochs': 5})
+ #
+ # model = NaturePyTorchModule(wandb.config['model_type'])
+ #
+ # wandb.config['input_size'] = 224
+ #
+ # dls = Dataloaders(artifact_name_alias=f"{DATASET_NAME}:latest",
+ # batch_size=wandb.config['batch_size'],
+ # input_size=wandb.config['input_size'])
+ #
+ # # Train the model
+ # learning_rate = wandb.config["lr"]
+ # gamma = wandb.config["gamma"]
+ # epochs = wandb.config["epochs"]
+ #
+ # device = torch.device("cpu")
+ # optimizer = optim.Adadelta(model.parameters(), lr=wandb.config['lr'])
+ # scheduler = StepLR(optimizer, step_size=1, gamma=wandb.config['gamma'])
+ #
+ # best_loss = float("inf")
+ # best_model = None
+ #
+ # for epoch_ndx in range(epochs):
+ # model.train()
+ # for batch_ndx, batch in enumerate(dls.train):
+ # data, target = batch[0].to("cpu"), batch[1].to("cpu")
+ # optimizer.zero_grad()
+ # preds = model(data)
+ # loss = F.cross_entropy(preds, target)
+ # loss.backward()
+ # optimizer.step()
+ # scheduler.step()
+ #
+ # ### Log your metrics ###
+ # wandb.log({
+ # "train/epoch_ndx": epoch_ndx,
+ # "train/batch_ndx": batch_ndx,
+ # "train/train_loss": loss,
+ # "train/learning_rate": optimizer.param_groups[0]["lr"]
+ # })
+ # print(f"Epoch: {epoch_ndx}, Batch: {batch_ndx}, Loss: {loss}")
+ #
+ # ### Evaluation at the end of each epoch ###
+ # test_loss, accuracy, preds, val_table, conf_mat = evaluate_model(
+ # model,
+ # dls.valid,
+ # dls.dataset.idx_to_class,
+ # dls.dataset.class_names,
+ # )
+ #
+ # is_best = test_loss < best_loss
+ #
+ # wandb.log({
+ # 'eval/test_loss': test_loss,
+ # 'eval/accuracy': accuracy,
+ # 'eval/conf_mat': conf_mat,
+ # 'eval/val_table': val_table})
+ #
+ # ### Checkpoing your model weights ###
+ # torch.save(model.state_dict(), "model.pth")
+ # art = wandb.Artifact(f"nature-{wandb.run.id}",
+ # type="model",
+ # metadata={'format': 'onnx',
+ # 'num_classes': len(dls.dataset.class_names),
+ # 'model_type': wandb.config['model_type'],
+ # 'model_input_size': wandb.config['input_size'],
+ # 'index_to_class': dls.dataset.idx_to_class})
+ #
+ # art.add_file("model.pth")
+ #
+ # ### Add aliases to keep track of your best checkpoints over time
+ # wandb.log_artifact(art, aliases=["best", "latest"] if is_best else None)
+ # if is_best:
+ # best_model = art
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Manage all your model checkpoints for a project under one roof.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Model Registry
+ After logging a bunch of checkpoints across multiple runs during experimentation, now comes time to hand-off the best checkpoint to the next stage of the workflow (e.g. testing, deployment).
+
+ The Model Registry is a central page that lives above individual W&B projects. It houses **Registered Models**, portfolios that store "links" to the valuable checkpoints living in individual W&B Projects.
+
+ The model registry offers a centralized place to house the best checkpoints for all your model tasks. Any `model` artifact you log can be "linked" to a Registered Model.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Creating **Registered Models** and Linking through the UI
+ #### 1. Access your team's model registry by going the team page and selecting `Model Registry`
+
+ 
+
+ #### 2. Create a new Registered Model.
+
+ 
+
+ #### 3. Go to the artifacts tab of the project that holds all your model checkpoints
+
+ 
+
+ #### 4. Click "Link to Registry" for the model artifact version you want.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Creating Registered Models and Linking through the **API**
+ You can [link a model via api](https://docs.wandb.ai/guides/models) with `wandb.run.link_artifact` passing in the artifact object, and the name of the **Registered Model**, along with aliases you want to append to it. **Registered Models** are entity (team) scoped in W&B so only members of a team can see and access the **Registered Models** there. You indicate a registered model name via api with `/model-registry/`. If a Registered Model doesn't exist, one will be created automatically.
+ """)
+ return
+
+
+@app.cell
+def _(ENTITY, best_model, wandb):
+ if ENTITY is not None:
+ wandb.run.link_artifact(best_model, f'{ENTITY}/model-registry/Model Registry Tutorial', aliases=['staging'])
+ else:
+ print('Must indicate entity where Registered Model will exist')
+ wandb.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### What is "Linking"?
+ When you link to the registry, this creates a new version of that Registered Model, which is just a pointer to the artifact version living in that project. There's a reason W&B segregates the versioning of artifacts in a project from the versioning of a Registered Model. The process of linking a model artifact version is equivalent to "bookmarking" that artifact version under a Registered Model task.
+
+ Typically during R&D/experimentation, researchers generate 100s, if not 1000s of model checkpoint artifacts, but only one or two of them actually "see the light of day." This process of linking those checkpoints to a separate, versioned registry helps delineate the model development side from the model deployment/consumption side of the workflow. The globally understood version/alias of a model should be unpolluted from all the experimental versions being generated in R&D and thus the versioning of a Registered Model increments according to new "bookmarked" models as opposed to model checkpoint logging.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Create a Centralized Hub for all your models
+ - Add a model card, tags, slack notifactions to your Registered Model
+ - Change aliases to reflect when models move through different phases
+ - Embed the model registry in reports for model documentation and regression reports. See this report as an [example](https://api.wandb.ai/links/wandb-smle/r82bj9at)
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Set up Slack Notifications when new models get linked to the registry
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Consuming a Registered Model
+ You now can consume any registered model via API by referring the corresponding `name:alias`. Model consumers, whether they are engineers, researchers, or CI/CD processes, can go to the model registry as the central hub for all models that should "see the light of day": those that need to go through testing or move to production.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%wandb -h 600
+ #
+ # run = wandb.init(project=PROJECT_NAME, entity=ENTITY, job_type='inference')
+ # artifact = run.use_artifact(f'{ENTITY}/model-registry/Model Registry Tutorial:staging', type='model')
+ # artifact_dir = artifact.download()
+ # wandb.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-model-registry-models-quickstart/wandb_model_registry_models_quickstart.py b/marimo/convert/wandb-model-registry-models-quickstart/wandb_model_registry_models_quickstart.py
new file mode 100644
index 00000000..394f4dfb
--- /dev/null
+++ b/marimo/convert/wandb-model-registry-models-quickstart/wandb_model_registry_models_quickstart.py
@@ -0,0 +1,85 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Models Quickstart
+
+ Quickly see the mechanics for logging and linking a model to the Weights & Biases model registry:
+ 1. `run = wandb.init()`: Start a run to track training
+ 2. `run.log_artifact()`: Track your trained model weights as an artifact
+ 3. `run.link_artifact()`: Link a specific model version it to the registry
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ import wandb
+ import random
+
+ # Start a new W&B run
+ with wandb.init(project="models_quickstart") as run:
+
+ # Simulate logging model metrics
+ run.log({"acc": random.random()})
+
+ # Create a simulated model file
+ with open("my_model.h5", "w") as f: f.write("Model: " + str(random.random()))
+
+ # Save the dummy model to W&B
+ best_model = wandb.Artifact(f"model_{run.id}", type='model')
+ best_model.add_file('my_model.h5')
+ run.log_artifact(best_model)
+
+ # Link the model to the Model Registry
+ run.link_artifact(best_model, 'model-registry/My Registered Model')
+
+ run.finish()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## How do you use Models in a real project?
+ This example keeps it simple. We're not training a real model, just focusing on the model mechanics of `log_artifact()` and `link_artifact()`.
+
+ In the real world, you don't want to link _every_ model version to the registry. Instead, use the model registry as a place to bookmark and organize your best models.
+
+ Learn more in the [Models docs](https://docs.wandb.ai/guides/models).
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-model-registry-new-model-logging-in-w-b/wandb_model_registry_new_model_logging_in_w_b.py b/marimo/convert/wandb-model-registry-new-model-logging-in-w-b/wandb_model_registry_new_model_logging_in_w_b.py
new file mode 100644
index 00000000..af80b5cb
--- /dev/null
+++ b/marimo/convert/wandb-model-registry-new-model-logging-in-w-b/wandb_model_registry_new_model_logging_in_w_b.py
@@ -0,0 +1,388 @@
+# /// script
+# dependencies = ["einops", "wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Logging and Registering Models in W&B
+ It's never been easier to log your model checkpoints, keep track of the best ones, and maintain lineage of runs and results!
+
+ W&B is introducing a few convenience methods to make logging models and linking them to the registry simple:
+ - `log_model`
+ - `use_model`
+ - `link_model`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Imports
+ """)
+ return
+
+
+@app.cell
+def _():
+ # packages added via marimo's package management: wandb einops !pip install -qqq wandb einops
+ return
+
+
+@app.cell
+def _():
+ import torch
+ from torch import nn
+ from einops import rearrange, repeat
+ from einops.layers.torch import Rearrange
+ from torch.utils.data import Dataset
+ from torchvision import transforms
+ from torch.optim import Adam
+ from torch.utils.data import DataLoader
+ from torchvision import datasets
+ import wandb
+
+ return (
+ Adam,
+ DataLoader,
+ Dataset,
+ Rearrange,
+ nn,
+ rearrange,
+ repeat,
+ torch,
+ wandb,
+ )
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Log in to W&B
+ - You can explicitly login using `wandb login` or `wandb.login()` (See below)
+ - Alternatively you can set environment variables. There are several env variables which you can set to change the behavior of W&B logging. The most important are:
+ - `WANDB_API_KEY` - create a new API key in your "Settings" section under your profile at [wandb.ai/settings](https://wandb.ai/settings)
+ - `WANDB_BASE_URL` - this is the url of the W&B server
+ - Create a new API key in "Profile" -> "Settings" in the W&B App. Store your API key securely. It can only be viewed once when created.
+
+ 
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Define the Model and Dataset
+ This is a simple implementation of a Vision Transformer (ViT) and utilizes a random dataset for training.
+ - Credit to https://github.com/lucidrains/vit-pytorch
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Define some config for the model and dataset
+ """)
+ return
+
+
+@app.cell
+def _():
+ # Define the number of samples, classes, and image size
+ num_samples = 100
+ num_classes = 10
+ image_size = 256
+ batch_size = 32
+ return batch_size, image_size, num_classes, num_samples
+
+
+@app.cell
+def _(
+ DataLoader,
+ Dataset,
+ Rearrange,
+ batch_size,
+ image_size,
+ nn,
+ num_classes,
+ num_samples,
+ rearrange,
+ repeat,
+ torch,
+):
+ # helpers
+ def pair(t):
+ return t if isinstance(t, tuple) else (t, t)
+
+ # classes
+ class FeedForward(nn.Module):
+ def __init__(self, dim, hidden_dim, dropout = 0.):
+ super().__init__()
+ self.net = nn.Sequential(
+ nn.LayerNorm(dim),
+ nn.Linear(dim, hidden_dim),
+ nn.GELU(),
+ nn.Dropout(dropout),
+ nn.Linear(hidden_dim, dim),
+ nn.Dropout(dropout)
+ )
+
+ def forward(self, x):
+ return self.net(x)
+
+ class Attention(nn.Module):
+ def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+ super().__init__()
+ inner_dim = dim_head * heads
+ project_out = not (heads == 1 and dim_head == dim)
+
+ self.heads = heads
+ self.scale = dim_head ** -0.5
+
+ self.norm = nn.LayerNorm(dim)
+
+ self.attend = nn.Softmax(dim = -1)
+ self.dropout = nn.Dropout(dropout)
+
+ self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+ self.to_out = nn.Sequential(
+ nn.Linear(inner_dim, dim),
+ nn.Dropout(dropout)
+ ) if project_out else nn.Identity()
+
+ def forward(self, x):
+ x = self.norm(x)
+
+ qkv = self.to_qkv(x).chunk(3, dim = -1)
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+ dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+ attn = self.attend(dots)
+ attn = self.dropout(attn)
+
+ out = torch.matmul(attn, v)
+ out = rearrange(out, 'b h n d -> b n (h d)')
+ return self.to_out(out)
+
+ class Transformer(nn.Module):
+ def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+ super().__init__()
+ self.norm = nn.LayerNorm(dim)
+ self.layers = nn.ModuleList([])
+ for _ in range(depth):
+ self.layers.append(nn.ModuleList([
+ Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+ FeedForward(dim, mlp_dim, dropout = dropout)
+ ]))
+
+ def forward(self, x):
+ for attn, ff in self.layers:
+ x = attn(x) + x
+ x = ff(x) + x
+
+ return self.norm(x)
+
+ class ViT(nn.Module):
+ def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
+ super().__init__()
+ image_height, image_width = pair(image_size)
+ patch_height, patch_width = pair(patch_size)
+
+ assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+ num_patches = (image_height // patch_height) * (image_width // patch_width)
+ patch_dim = channels * patch_height * patch_width
+ assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+ self.to_patch_embedding = nn.Sequential(
+ Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+ nn.LayerNorm(patch_dim),
+ nn.Linear(patch_dim, dim),
+ nn.LayerNorm(dim),
+ )
+
+ self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+ self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+ self.dropout = nn.Dropout(emb_dropout)
+
+ self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+ self.pool = pool
+ self.to_latent = nn.Identity()
+
+ self.mlp_head = nn.Linear(dim, num_classes)
+
+ def forward(self, img):
+ x = self.to_patch_embedding(img)
+ b, n, _ = x.shape
+
+ cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)
+ x = torch.cat((cls_tokens, x), dim=1)
+ x += self.pos_embedding[:, :(n + 1)]
+ x = self.dropout(x)
+
+ x = self.transformer(x)
+
+ x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+ x = self.to_latent(x)
+ return self.mlp_head(x)
+
+
+ # Define a custom dataset
+ class RandomImageDataset(Dataset):
+ def __init__(self, num_samples, num_classes, image_size):
+ self.num_samples = num_samples
+ self.num_classes = num_classes
+ self.image_size = image_size
+
+ def __len__(self):
+ return self.num_samples
+
+ def __getitem__(self, idx):
+ # Generate a random image tensor
+ image = torch.randn(3, self.image_size, self.image_size) # 3 channels, image_size x image_size
+ # Generate a random label
+ label = torch.randint(0, self.num_classes, (1,)).item()
+ return image, label
+
+
+
+ # Create the dataset
+ dataset = RandomImageDataset(num_samples=num_samples, num_classes=num_classes, image_size=image_size)
+
+ # Create a DataLoader
+ dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+ return ViT, dataloader
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Log Model Checkpoints to W&B with 1 Line!
+
+ Use the `log_model` method to log a model artifact containing the contents inside the ‘path’ to an a run. It also marks it as an output to the run. You can see the full lineage graph of the model artifact by accessing the [lineage](https://docs.wandb.ai/guides/artifacts/explore-and-traverse-an-artifact-graph#docusaurus_skipToContent_fallback) tab inside the Artifacts view.
+
+ `log_model()` takes as input:
+
+ - `path`: A path to the model file(s), which can be a local file (of the form `/local/directory/file.txt`), directory (of the form `/local/directory`), or reference path to S3 (`s3://bucket/path`).
+ - `name`: An optional name for the model artifact the files will be logged to. Note that if no name is specified, This will default to the basename of the input path prepended with the run ID.
+ - `aliases`: An optional list of aliases, which can be thought of as semantic ‘nicknames’ or identifiers for a model version. For example, if this model yielded the best accuracy, you might add the alias ‘highest-accuracy’ or ‘best’.
+ """)
+ return
+
+
+@app.cell
+def _(Adam, ViT, dataloader, image_size, nn, num_classes, torch, wandb):
+ run = wandb.init(project="new_model_logging",
+ job_type="training")
+
+ v = ViT(
+ image_size = image_size,
+ patch_size = 32,
+ num_classes = num_classes,
+ dim = 128,
+ depth = 3,
+ heads = 2,
+ mlp_dim = 256,
+ dropout = 0.1,
+ emb_dropout = 0.1
+ )
+
+ # Define the loss function and optimizer
+ criterion = nn.CrossEntropyLoss()
+ optimizer = Adam(v.parameters(), lr=0.003)
+
+ # Training loop
+ best_accuracy = 0
+ for epoch in range(5): # number of epochs
+ for images, labels in dataloader:
+ # Forward pass
+ preds = v(images)
+ loss = criterion(preds, labels)
+
+ # Backward pass and optimization
+ optimizer.zero_grad()
+ loss.backward()
+ optimizer.step()
+
+ wandb.log({"train/loss": loss})
+
+ # Model evaluation after each epoch (using a validation set)
+ # Here you would write your validation loop and calculate accuracy
+ val_accuracy = 0.5 # Assume this is the validation accuracy you compute
+ model_path = 'model_vit.pth'
+ torch.save(v.state_dict(), model_path)
+
+ # Check if this is the best model so far
+ if val_accuracy > best_accuracy:
+ best_accuracy = val_accuracy
+ # Log the model to your W&B run
+ wandb.log_model(name=f"model_vit-{wandb.run.id}", path=model_path, aliases=["best", f"epoch_{epoch}"])
+ else:
+ wandb.log_model(name=f"model_vit-{wandb.run.id}", path=model_path, aliases=[f"epoch_{epoch}"])
+ return (run,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Link Your Best Models to the Model Registry
+ You can bookmark your best model checkpoints and centralize them across your team. The Model Registry allows you can organize your best models by task, manage model lifecycle, facilitate easy tracking and auditing throughout the ML lifecyle, and automate downstream actions with webhooks or jobs. You can this via api through `link_model()`, which takes as input:
+
+ - `path`: A path to the model file(s), which can be a local file (of the form `/local/directory/file.txt`), directory (of the form `/local/directory`), or reference path to S3 (`s3://bucket/path`).
+ - `registered_model_name`: the name of the Registered Model - a collection of linked model versions in the Model Registry, typically representing a team’s ML task - that the model should be linked to. If no Registered Model with the given name exists, a new one will be created with this name.
+ - `name`: An **optional** name for the model artifact the files will be logged to. Note that if no name is specified, This will default to the basename of the input path prepended with the run ID.
+ - `aliases`: An **optional** list of aliases, which can be thought of as semantic ‘nicknames’ or identifiers for a linked model version. For example, since this model is being linked, or published, to the Model Registry, you might add an alias “staging” or “QA”.
+ """)
+ return
+
+
+@app.cell
+def _(run, wandb):
+ # Link the best model to the W&B Model Registry (after all epochs are finished)
+ artifact_name = f"model_vit-{wandb.run.id}:best"
+ best_model_path = wandb.use_model(artifact_name)
+
+ # Link the best model to the registry
+ wandb.link_model(path=best_model_path,
+ registered_model_name="Industrial ViT",
+ aliases=["staging"])
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/wandb-model-registry-w-b-model-registry-quickstart/wandb_model_registry_w_b_model_registry_quickstart.py b/marimo/convert/wandb-model-registry-w-b-model-registry-quickstart/wandb_model_registry_w_b_model_registry_quickstart.py
new file mode 100644
index 00000000..1aadc8fb
--- /dev/null
+++ b/marimo/convert/wandb-model-registry-w-b-model-registry-quickstart/wandb_model_registry_w_b_model_registry_quickstart.py
@@ -0,0 +1,123 @@
+# /// script
+# dependencies = ["wandb"]
+# ///
+
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # W&B Model Registry Quickstart
+ """)
+ return
+
+
+@app.cell
+def _():
+ #@title 1) Run this cell to set up `wandb` and define helper functions
+
+
+ # INSTALL W&B LIBRARY
+ # packages added via marimo's package management: wandb !pip install wandb -qqq
+
+ import wandb
+ import os
+ import math
+ import random
+
+ # FORM VARIABLES
+ PROJECT = "Model_Registry_Quickstart" #@param {type:"string"}
+ RUN_COUNT = 3 #@param {type:"integer"}
+
+
+ # HELPER FUNCTIONS
+ # Create fake data to simulate training a model.
+
+ # Simulate setting up hyperparameters
+ # Return: A dict of params to log as config to W&B
+ def set_config():
+ config={
+ "learning_rate": 0.01 + 0.1 * random.random(),
+ "batch_size": 128,
+ "architecture": "CNN",
+ }
+ return config
+
+ # Simulate training a model
+ # Return: A model file to log as an artifact to W&B
+ def get_model():
+ file_name = "demo_model.h5"
+ model_file = open(file_name, 'w')
+ model_file.write('Imagine this is a big model file! ' + str(random.random()))
+ model_file.close()
+ return file_name
+
+ # Simulate logging metrics from model training
+ # Return: A dictionary of metrics to log to W&B
+ def get_metrics(epoch):
+ metrics = {
+ "acc": .8 + 0.04 * (math.log(1 + epoch + random.random()) + (0.3 * random.random())),
+ "val_acc": .75 + 0.04 * (math.log(1 + epoch + random.random()) - (0.3 * random.random())),
+ "loss": .1 + 0.1 * (4 - math.log(1 + epoch + random.random()) + (0.3 * random.random())),
+ "val_loss": .1 + 0.16 * (5 - math.log(1 + epoch + random.random()) - (0.3 * random.random())),
+ }
+ return metrics
+
+ return PROJECT, RUN_COUNT, get_metrics, get_model, set_config, wandb
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### 2) Log a model
+ """)
+ return
+
+
+@app.cell
+def _(PROJECT, RUN_COUNT, get_metrics, get_model, set_config, wandb):
+ for _ in range(RUN_COUNT):
+
+ # 1️⃣ Initialize a new W&B run to track this job
+ run = wandb.init(project=PROJECT, config=set_config())
+
+ for epoch in range(5):
+ # 2️⃣ Log metrics to W&B for each epoch of training
+ run.log(get_metrics(epoch))
+
+ # 3️⃣ At the end of training, save the model artifact
+ # Name this artifact after the current run
+ model_artifact_name = "demo_model_" + run.id
+ # Create a new artifact
+ model = wandb.Artifact(model_artifact_name, type='model')
+ # Add files to the artifact, in this case a simple text file
+ model.add_file(get_model())
+ # Log the model to W&B
+ run.log_artifact(model)
+
+ # Call finish if you're in a notebook, to mark the run as done
+ run.finish()
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/yolo-logging-yolov5-experiments-with-w-b/yolo_logging_yolov5_experiments_with_w_b.py b/marimo/convert/yolo-logging-yolov5-experiments-with-w-b/yolo_logging_yolov5_experiments_with_w_b.py
new file mode 100644
index 00000000..107074f5
--- /dev/null
+++ b/marimo/convert/yolo-logging-yolov5-experiments-with-w-b/yolo_logging_yolov5_experiments_with_w_b.py
@@ -0,0 +1,208 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ # You Always Log Everything (YALE)
+
+ ### Logging YOLOv5 Experiments with W&B
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [YOLO](https://github.com/ultralytics/yolov5) ("You Only Look Once") provides tools for real-time object detection with convolutional neural networks.
+
+ YOLO now works with [Weights & Biases](http://wandb.com),
+ an experiment tracking toolkit, so you can
+ keep track of all the hyperparameters you've tried,
+ view real-time updates on system and model metrics,
+ version and store datasets and models,
+ [and more](http://github.com/wandb/examples)!
+
+ In this colab, we'll show you how to use YOLO and W&B together.
+ **It's as easy as running a single `pip install` before you run your YOLO experiments!**
+
+ Follow along with a video tutorial on YouTube.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 0. Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ First, let's get ourselves organized: clone the repo, install our dependencies, and confirm we've got PyTorch and a GPU.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! git clone --depth 1 https://github.com/ultralytics/yolov5
+ subprocess.call(['git', 'clone', '--depth', '1', 'https://github.com/ultralytics/yolov5'])
+ # clone repo
+ import os
+ os.chdir('yolov5')
+ # '%pip install -qr requirements.txt # install dependencies' command supported automatically in marimo
+
+ import torch
+ from IPython.display import Image, clear_output # to display images
+
+ clear_output()
+ print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))
+ return Image, torch
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 1. Inference
+
+ Now, let's apply a pre-trained, already existing object detection network.
+
+ `detect.py` runs inference on a variety of sources, downloading models automatically from the [latest YOLOv5 release](https://github.com/ultralytics/yolov5/releases).
+
+ Here, we'll just run a sample image through that network to make sure everything is working.
+ """)
+ return
+
+
+@app.cell
+def _(Image, subprocess):
+ #! python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images/
+ subprocess.call(['python', 'detect.py', '--weights', 'yolov5s.pt', '--img', '640', '--conf', '0.25', '--source', 'data/images/'])
+ Image(filename='runs/detect/exp/bus.jpg', width=600)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Results are saved to `runs/detect`. A full list of available inference sources:
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # 2. Training to Fine-Tune
+
+ For applications, it's often important to take a pre-trained model
+ and fine-tune it to work on a specific dataset --
+ [for example, in a construction safety application](https://wandb.ai/authors/artifact-workplace-safety/reports/Organize-Your-Machine-Learning-Pipelines-with-Artifacts--VmlldzoxODQwNTY), we might use fine-tuning to specialize our network in detecting the presence/absence of protective equipment.
+
+ We'll mimic this process on the
+ [COCO128](https://www.kaggle.com/ultralytics/coco128) image tutorial dataset.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess, torch):
+ 2# Download COCO128
+ torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip', 'tmp.zip')
+ #! unzip -q tmp.zip -d ../ && rm tmp.zip
+ subprocess.call(['unzip', '-q', 'tmp.zip', '-d', '../', '&&', 'rm', 'tmp.zip'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Weights & Biases Logging (🚀 NEW)
+
+ [Weights & Biases](https://www.wandb.com/) (W&B) is now integrated with YOLOv5 for real-time visualization and cloud logging of training runs. This allows for better run comparison and introspection, as well improved visibility and collaboration among team members. To enable W&B logging install `wandb`, and then train normally (you will be guided through setting up `wandb` account during your first use).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !pip install "wandb==0.12.10"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ And that's it! So long as W&B is installed, you'll get rich, detailed metrics in a live dashboard accessible from a browser on any device.
+
+ Just click the link that appears below next to `wandb` and the 🚀 emoji.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ # Train YOLOv5s on COCO128 for 5 epochs
+ #! python train.py --img 640 --batch 64 --epochs 5 --data coco128.yaml --weights yolov5s.pt
+ subprocess.call(['python', 'train.py', '--img', '640', '--batch', '64', '--epochs', '5', '--data', 'coco128.yaml', '--weights', 'yolov5s.pt'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ With W&B, during training you will see live updates on the dashboard at [wandb.ai](https://www.wandb.ai/), including interactive bounding box visualizations (look for a panel called "Images" in the Media panel section), and you can create and share detailed [Reports](https://wandb.ai/glenn-jocher/yolov5_tutorial/reports/YOLOv5-COCO128-Tutorial-Results--VmlldzozMDI5OTY) of your results. For more information see the [YOLOv5 Weights & Biases Tutorial](https://github.com/ultralytics/yolov5/issues/1289)
+ or check out the [video tutorial for this notebook](http://wandb.me/yolov5-video).
+
+ 
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/yolo-train-and-debug-yolov5-models-with-weights-biases/yolo_train_and_debug_yolov5_models_with_weights_biases.py b/marimo/convert/yolo-train-and-debug-yolov5-models-with-weights-biases/yolo_train_and_debug_yolov5_models_with_weights_biases.py
new file mode 100644
index 00000000..b4c1e6fc
--- /dev/null
+++ b/marimo/convert/yolo-train-and-debug-yolov5-models-with-weights-biases/yolo_train_and_debug_yolov5_models_with_weights_biases.py
@@ -0,0 +1,326 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train and Debug YOLOv5 Models with Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this colab,
+ we'll demonstrate how to use the W&B integration with
+ version 5 of the "You Only Look Once"
+ (aka [YOLOv5](https://github.com/ultralytics/yolov5))
+ real-time object detection framework
+ to track model metrics,
+ inspect model outputs,
+ and restart interrupted runs.
+
+ ### Follow along with a [video tutorial →](https://wandb.me/yolo-video)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We begin by downloading the
+ [YOLOv5 GitHub repo](https://github.com/ultralytics/yolov5)
+ and a
+ [dataset of chessboard images with labeled bounding boxes around the pieces](https://public.roboflow.com/object-detection/chess-full).
+ Below, we'll use this dataset to train a model to detect chess pieces in images.
+
+ We also install all the requirements for YOLOv5 and `wandb`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !git clone --depth 1 https://github.com/ultralytics/yolov5.git
+ # !curl -L "https://public.roboflow.com/ds/1BpjFZe9ST?key=KXD7eDvwTa" > roboflow.zip; unzip -o roboflow.zip; rm roboflow.zip
+ # %cd /content/yolov5
+ # !pip install -r requirements.txt
+ # !pip install "wandb==0.12.10"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Detect
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ YOLOv5 provides highly-accurate, fast models that are pretrained on the
+ [Common Objects in COntext (COCO) dataset](https://cocodataset.org/#home).
+
+ If your object detection application involves only
+ [classes from the COCO dataset](https://gist.github.com/AruniRC/7b3dadd004da04c80198557db5da4bda),
+ like "Stop Sign" and "Pizza",
+ then these pretrained models may be all you need!
+
+ The cell below runs a pretrained model on an example image
+ using `detect.py` from the YOLOv5 toolkit.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ from IPython.display import Image
+
+ #! python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images/bus.jpg
+ subprocess.call(['python', 'detect.py', '--weights', 'yolov5s.pt', '--img', '640', '--conf', '0.25', '--source', 'data/images/bus.jpg'])
+ Image(filename='runs/detect/exp/bus.jpg', width=600)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Chess pieces are not among the objects in COCO,
+ so our pretrained models don't know how to detect them
+ and we can't just use `detect.py` with one of those models.
+
+ Instead, we need to train the models to detect chess pieces,
+ using YOLOv5's `train.py`.
+ We don't have to start our models from scratch though!
+ We can [finetune](https://morioh.com/p/4cd5996c0e64)
+ the pretrained models on our chess piece dataset.
+ This substantially speeds up training.
+
+ Model training is a complex process,
+ so we'll want to track the inputs and outputs,
+ log information about model behavior during training,
+ and record system state and metrics.
+ Additionally, we might
+
+ That's where [Weights & Biases](https://wandb.ai/site)
+ comes in:
+ the `wandb` library provides all the tools you need to thoroughly
+ and effectively log model training experiments.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ YOLOv5 comes with `wandb` already integrated,
+ so all you need to do is configure the logging
+ with command line arguments.
+
+ * `--project` sets the W&B project to which we're logging
+ (akin to a GitHub repo).
+ * `--upload_dataset` tells `wandb`
+ to upload the dataset as [a dataset-visualization Table](https://docs.wandb.ai/guides/datasets-and-predictions).
+ At regular intervals set by `--bbox_interval`,
+ the model's outputs on the validation set will also be logged to W&B.
+ * `--save-period` sets the number of epochs to wait
+ in between logging the model checkpoints.
+ If not set, only the final trained model is logged.
+
+ Even without these arguments,
+ basic model metrics and some model outputs will still be saved to W&B.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_: to use this same training and logging setup on a different dataset,
+ just [create a `data.yaml` for that dataset](https://github.com/ultralytics/yolov5/issues/12)
+ and provide it to the `--data` argument.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python train.py --data ../data.yaml --epochs 3 --project yolo-wandb-demo --bbox_interval 1 --save-period 1
+ subprocess.call(['python', 'train.py', '--data', '../data.yaml', '--epochs', '3', '--project', 'yolo-wandb-demo', '--bbox_interval', '1', '--save-period', '1'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's where you can find the uploaded evaluation results in the W&B UI:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Resume Crashed Runs
+
+ In addition to making it easier to debug our models,
+ the W&B integration can help rescue crash or interrupted runs.
+
+ Two steps above helped set us up for this:
+ 1. By setting a `--save-period`, we regularly logged the model to W&B,
+ which means we can recreate our model and then resume the run on any
+ device with the dataset available.
+ 2. By using `--upload_dataset`, we logged the data to W&B,
+ which means we can recreate the data as well and so
+ resume runs on any device, whether the dataset is present on disk or not
+
+ To resume a crashed or interrupted run:
+ * Go to that run's overview section on W&B dashboard
+ * Copy the run path
+ * Pass the run path as the `--resume` argument, plus the prefix
+ `wandb-artifact://`.
+ This prefix tells YOLO that the files are located on wandb, rather than locally.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ```python
+ crashed_run_path = "entity/project/run-id" # your path here
+ !python train.py --resume wandb-artifact://{crashed_run_path}
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # End Notes
+
+ ### Distributed Data-Parallel Training
+
+ All YOLO+W&B features are
+ DDP-aware and compatible.
+ Train on as many GPUs as you can muster,
+ and we'll keep logging!
+
+ ### Logging Large Datasets
+
+ For very large datasets,
+ the initial dataset upload triggered by `--log_dataset`
+ might be prohibitively expensive.
+
+ In that case,
+ check out the
+ [`log_dataset.py` script](https://github.com/ultralytics/yolov5/blob/master/utils/wandb_logging/log_dataset.py)
+ included in YOLOv5.
+
+ ### `stripped` Models
+
+ At the end of training,
+ a "stripped" version of the model is saved
+ to W&B.
+ This version of the model file is much smaller,
+ but is missing
+ accumulated data required for resuming training.
+ It's intended for use in downstream inference.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/yolo-train-yolov5-model-on-a-custom-dataset-with-weights-biases/yolo_train_yolov5_model_on_a_custom_dataset_with_weights_biases.py b/marimo/convert/yolo-train-yolov5-model-on-a-custom-dataset-with-weights-biases/yolo_train_yolov5_model_on_a_custom_dataset_with_weights_biases.py
new file mode 100644
index 00000000..f3bf4e6d
--- /dev/null
+++ b/marimo/convert/yolo-train-yolov5-model-on-a-custom-dataset-with-weights-biases/yolo_train_yolov5_model_on_a_custom_dataset_with_weights_biases.py
@@ -0,0 +1,324 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+ 
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train YOLOv5 model on a Custom Dataset with Weights & Biases (as a part of the YOLOv5 Series)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ This is a Colab for training a custom [YOLOv5](https://github.com/ultralytics/yolov5) model and using Weights & Biases to track training metrics, checkpoint weights and datasets. This Colab is featured in part 3 of the YOLOv5 Series.
+
+ ### Follow along with [YOLOv5 Series →](https://www.youtube.com/playlist?list=PLD80i8An1OEHEpJVjtujEb0lQWc0GhX_4)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We begin by downloading the
+ [YOLOv5 GitHub repo](https://github.com/ultralytics/yolov5) and installing all the requirements for YOLOv5 and `wandb`.
+
+ Here's an example of a [wandb dashboard](https://wandb.ai/ivangoncharov/custom_yolov5?workspace=user-ivangoncharov).
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !git clone --depth 1 https://github.com/ultralytics/yolov5.git
+ # %cd /content/yolov5
+ # !pip install -r requirements.txt
+ # !pip install wandb==0.12.10
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Detect
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ YOLOv5 provides highly-accurate, fast models that are pretrained on the
+ [Common Objects in COntext (COCO) dataset](https://cocodataset.org/#home).
+
+ If your object detection application involves only
+ [classes from the COCO dataset](https://gist.github.com/AruniRC/7b3dadd004da04c80198557db5da4bda),
+ like "Stop Sign" and "Pizza",
+ then these pretrained models may be all you need!
+
+ The cell below runs a pretrained model on an example image
+ using `detect.py` from the YOLOv5 toolkit.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ from IPython.display import Image
+
+ #! python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images/bus.jpg
+ subprocess.call(['python', 'detect.py', '--weights', 'yolov5s.pt', '--img', '640', '--conf', '0.25', '--source', 'data/images/bus.jpg'])
+ Image(filename='runs/detect/exp/bus.jpg', width=600)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Generating a .yaml file for training on the bus dataset that's featured in the YOLOv5 Series. You can skip this step when using your own custom dataset.
+ """)
+ return
+
+
+@app.cell
+def _():
+ import json
+ import yaml
+ data ={
+ 'names':['closed_door', 'opened_door', 'bus' ,'number'],
+ 'nc': 4,
+ 'train': "wandb-artifact://wandb/custom_yolov5/train",
+ 'val': "wandb-artifact://wandb/custom_yolov5/val",
+ }
+ with open('bus_dataset.yaml', 'w') as outfile:
+ yaml.dump(data, outfile)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Your custom classes (probably) are not among the objects in COCO,
+ so our pretrained models don't know how to detect them
+ and we can't just use `detect.py` with one of those models.
+
+ Instead, we need to train the models to detect our custom classes,
+ using YOLOv5's `train.py`.
+ We don't have to start our models from scratch though!
+ We can finetune the pretrained models on our custom dataset.
+ This substantially speeds up training.
+
+ Model training is a complex process,
+ so we'll want to track the inputs and outputs,
+ log information about model behavior during training,
+ and record system state and metrics.
+
+ That's where [Weights & Biases](https://wandb.ai/site)
+ comes in:
+ the `wandb` library provides all the tools you need to thoroughly
+ and effectively log model training experiments.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ YOLOv5 comes with `wandb` already integrated,
+ so all you need to do is configure the logging
+ with command line arguments.
+
+ * `--project` sets the W&B project to which we're logging
+ (akin to a GitHub repo).
+ * `--upload_dataset` tells `wandb`
+ to upload the dataset as [a dataset-visualization Table](https://docs.wandb.ai/guides/datasets-and-predictions).
+ At regular intervals set by `--bbox_interval`,
+ the model's outputs on the validation set will also be logged to W&B.
+ * `--save_period` sets the number of epochs to wait
+ in between logging the model checkpoints.
+ If not set, only the final trained model is logged.
+
+ Even without these arguments,
+ basic model metrics and some model outputs will still be saved to W&B.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### To train on your custom dataset you'll need a special .yaml file. In the YOLOv5 Series we use Weights & Biases to upload our custom dataset to the cloud and generate the required .yaml file.
+
+ #### To learn more you can watch [part 2 of the series →](https://youtu.be/a9Bre0YJ8L8)
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python train.py --data bus_dataset.yaml --epochs 30 --project custom_yolov5 --bbox_interval 1 --save_period 1 --weights yolov5s.pt
+ subprocess.call(['python', 'train.py', '--data', 'bus_dataset.yaml', '--epochs', '30', '--project', 'custom_yolov5', '--bbox_interval', '1', '--save_period', '1', '--weights', 'yolov5s.pt'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's where you can find the uploaded evaluation results in the W&B UI:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Resume Crashed Runs
+
+ In addition to making it easier to debug our models,
+ the W&B integration can help rescue crash or interrupted runs.
+
+ Two steps above helped set us up for this:
+ 1. By setting a `--save_period`, we regularly logged the model to W&B,
+ which means we can recreate our model and then resume the run on any
+ device with the dataset available.
+ 2. By using `--upload_dataset`, we logged the data to W&B,
+ which means we can recreate the data as well and so
+ resume runs on any device, whether the dataset is present on disk or not
+
+ To resume a crashed or interrupted run:
+ * Go to that run's overview section on W&B dashboard
+ * Copy the run path
+ * Pass the run path as the `--resume` argument, plus the prefix
+ `wandb-artifact://`.
+ This prefix tells YOLO that the files are located on wandb, rather than locally.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ crashed_run_path = "ivangoncharov/custom_yolov5/1xnnwg15" # your path here
+ #! python train.py --resume wandb-artifact://{crashed_run_path}
+ subprocess.call(['python', 'train.py', '--resume', 'wandb-artifact://{crashed_run_path}'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # End Notes
+
+ ### Distributed Data-Parallel Training
+
+ All YOLO+W&B features are
+ DDP-aware and compatible.
+ Train on as many GPUs as you can muster,
+ and we'll keep logging!
+
+ ### Logging Large Datasets
+
+ For very large datasets,
+ the initial dataset upload triggered by `--log_dataset`
+ might be prohibitively expensive.
+
+ In that case,
+ check out the
+ [`log_dataset.py` script](https://github.com/ultralytics/yolov5/blob/master/utils/wandb_logging/log_dataset.py)
+ included in YOLOv5.
+
+ ### `stripped` Models
+
+ At the end of training,
+ a "stripped" version of the model is saved
+ to W&B.
+ This version of the model file is much smaller,
+ but is missing
+ accumulated data required for resuming training.
+ It's intended for use in downstream inference.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/yolo-wildfire-smoke-detection-with-yolov5-roboflow-and-weights-biases-1/yolo_wildfire_smoke_detection_with_yolov5_roboflow_and_weights_biases_1.py b/marimo/convert/yolo-wildfire-smoke-detection-with-yolov5-roboflow-and-weights-biases-1/yolo_wildfire_smoke_detection_with_yolov5_roboflow_and_weights_biases_1.py
new file mode 100644
index 00000000..7b7b1069
--- /dev/null
+++ b/marimo/convert/yolo-wildfire-smoke-detection-with-yolov5-roboflow-and-weights-biases-1/yolo_wildfire_smoke_detection_with_yolov5_roboflow_and_weights_biases_1.py
@@ -0,0 +1,283 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+
+
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train and Debug YOLOv5 Models with Weights & Biases
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this colab,
+ we'll demonstrate how to use the W&B integration with
+ version 5 of the "You Only Look Once"
+ (aka [YOLOv5](https://github.com/ultralytics/yolov5))
+ real-time object detection framework
+ to track model metrics,
+ inspect model outputs,
+ and restart interrupted runs. We'll also make use of Roboflow's functionality for preprocessing and annotating our computer vision datasets.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%capture
+ # !git clone --depth 1 https://github.com/ultralytics/yolov5.git
+ # !curl -L "https://app.roboflow.com/ds/5F7B42TQI7?key=VhVxco4lnb" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip
+ # %cd /content/yolov5
+ # !pip install -r requirements.txt
+ # !pip install "wandb==0.12.10"
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Detect
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ YOLOv5 provides highly-accurate, fast models that are pretrained on the
+ [Common Objects in COntext (COCO) dataset](https://cocodataset.org/#home).
+
+ If your object detection application involves only
+ [classes from the COCO dataset](https://gist.github.com/AruniRC/7b3dadd004da04c80198557db5da4bda),
+ like "Stop Sign" and "Pizza",
+ then these pretrained models may be all you need!
+
+ The cell below runs a pretrained model on an example image
+ using `detect.py` from the YOLOv5 toolkit.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ from IPython.display import Image
+
+ #! python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images/bus.jpg
+ subprocess.call(['python', 'detect.py', '--weights', 'yolov5s.pt', '--img', '640', '--conf', '0.25', '--source', 'data/images/bus.jpg'])
+ Image(filename='runs/detect/exp/bus.jpg', width=600)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ YOLOv5 comes with `wandb` already integrated,
+ so all you need to do is configure the logging
+ with command line arguments.
+
+ * `--project` sets the W&B project to which we're logging
+ (akin to a GitHub repo).
+ * `--upload_dataset` tells `wandb`
+ to upload the dataset as [a dataset-visualization Table](https://docs.wandb.ai/guides/datasets-and-predictions).
+ At regular intervals set by `--bbox_interval`,
+ the model's outputs on the validation set will also be logged to W&B.
+ * `--save-period` sets the number of epochs to wait
+ in between logging the model checkpoints.
+ If not set, only the final trained model is logged.
+
+ Even without these arguments,
+ basic model metrics and some model outputs will still be saved to W&B.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ > _Note_: to use this same training and logging setup on a different dataset,
+ just [create a `data.yaml` for that dataset](https://github.com/ultralytics/yolov5/issues/12)
+ and provide it to the `--data` argument.
+ """)
+ return
+
+
+@app.cell
+def _(subprocess):
+ #! python train.py --data ../data.yaml --epochs 10 --project yolov5-roboflow-wandb --upload_dataset --bbox_interval 1 --save-period 1
+ subprocess.call(['python', 'train.py', '--data', '../data.yaml', '--epochs', '10', '--project', 'yolov5-roboflow-wandb', '--upload_dataset', '--bbox_interval', '1', '--save-period', '1'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Here's where you can find the uploaded evaluation results in the W&B UI:
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Resume Crashed Runs
+
+ In addition to making it easier to debug our models,
+ the W&B integration can help rescue crash or interrupted runs.
+
+ Two steps above helped set us up for this:
+ 1. By setting a `--save-period`, we regularly logged the model to W&B,
+ which means we can recreate our model and then resume the run on any
+ device with the dataset available.
+ 2. By using `--upload_dataset`, we logged the data to W&B,
+ which means we can recreate the data as well and so
+ resume runs on any device, whether the dataset is present on disk or not
+
+ To resume a crashed or interrupted run:
+ * Go to that run's overview section on W&B dashboard
+ * Copy the run path
+ * Pass the run path as the `--resume` argument, plus the prefix
+ `wandb-artifact://`.
+ This prefix tells YOLO that the files are located on wandb, rather than locally.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ```python
+ crashed_run_path = "entity/project/run-id" # your path here
+ !python train.py --resume wandb-artifact://{crashed_run_path}
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # End Notes
+
+ ### Distributed Data-Parallel Training
+
+ All YOLO+W&B features are
+ DDP-aware and compatible.
+ Train on as many GPUs as you can muster,
+ and we'll keep logging!
+
+ ### Logging Large Datasets
+
+ For very large datasets,
+ the initial dataset upload triggered by `--log_dataset`
+ might be prohibitively expensive.
+
+ In that case,
+ check out the
+ [`log_dataset.py` script](https://github.com/ultralytics/yolov5/blob/master/utils/wandb_logging/log_dataset.py)
+ included in YOLOv5.
+
+ ### `stripped` Models
+
+ At the end of training,
+ a "stripped" version of the model is saved
+ to W&B.
+ This version of the model file is much smaller,
+ but is missing
+ accumulated data required for resuming training.
+ It's intended for use in downstream inference.
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/marimo/convert/yolox-train-and-debug-yolox-models-with-weights-biases/yolox_train_and_debug_yolox_models_with_weights_biases.py b/marimo/convert/yolox-train-and-debug-yolox-models-with-weights-biases/yolox_train_and_debug_yolox_models_with_weights_biases.py
new file mode 100644
index 00000000..c374b2ed
--- /dev/null
+++ b/marimo/convert/yolox-train-and-debug-yolox-models-with-weights-biases/yolox_train_and_debug_yolox_models_with_weights_biases.py
@@ -0,0 +1,367 @@
+import marimo
+
+__generated_with = "0.24.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell
+def _():
+ import subprocess
+
+ return (subprocess,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+
+
+
+
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Train and Debug YOLOX Models with Weights & Biases 🪄🐝
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ In this colab, we'll demonstrate how to use the [W&B integration](https://docs.wandb.ai/guides/integrations/other/yolox) with [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX) for real-time object detection framework to track model metrics, log checkpoints and visualize predictions.
+
+ It can be done with just **1** added argument to your command!
+
+ ```
+ python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache] --logger wandb
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ To log the metrics and checkpoints to W&B during training, the wandb client now has a direct integration into YOLOX. Using wandb for logging automatically adds all the metrics to your W&B dashboard, saves the models at every evaluation step , tags the model with the best average precision and shows you visualizations of the predicted bounding boxes along with the confidence score!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Setup 🖥
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We begin by downloading the [YOLOX GitHub repository](https://github.com/Megvii-BaseDetection/YOLOX) and a subset of the [COCO dataset](https://cocodataset.org/#overview) for object detection.
+
+ Below, we'll use this dataset to train a model to detect objects in images.
+
+ We also install all the requirements for YOLOX and `wandb`.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%shell
+ # git clone --depth 1 https://github.com/manangoel99/YOLOX
+ # cd YOLOX
+ # git config --global user.name "Manan Goel"
+ # git config --global user.email "manangoel1999@gmail.com"
+ # git checkout -b WandbTables; git pull origin WandbTables; pip install -e .
+ # pip install wandb -qqq
+ return
+
+
+@app.cell
+def _():
+ import wandb
+
+ return (wandb,)
+
+
+@app.cell
+def _(wandb):
+ wandb.login()
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Downloading the dataset
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _api = wandb.Api()
+ _artifact = _api.artifact('manan-goel/YOLOX-coco/coco128:latest', type='dataset')
+ _artifact_dir = _artifact.download(root='/content')
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%shell
+ # unzip coco128.zip
+ # mv /content/coco128 /content/YOLOX/datasets/COCO
+ # cd /content/YOLOX
+ return
+
+
+@app.cell
+def _():
+ import os
+ os.chdir('YOLOX')
+ return (os,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Training 🏋️
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Using wandb just requires configuring the command line argument `--logger wandb`. This automatically turns on the wandb logger for your experiment and further arguments can be added -
+
+ 1. `wandb-project`: To specify the project in which experiment is being run.
+ 2. `wandb-run`: The name of the wandb run
+ 3. `wandb-entity`: Entity which is starting the run
+ 4. `wandb-log_checkpoints`: True/False to log model checkpoints to the wandb dashboard
+ 5. `wandb-num_eval_images`: Number of images from the validation set to be logged to wandb. Predictions corresponding to these can be visualized on the dashboard. No images are logged if the value is 0 and all are logged if the value is -1.
+
+ and more!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Reproduce the results
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The different YOLOX models can be trained from scratch with the entire process being logged to W&B. In this case we train on a much smaller subset of the COCO dataset.
+ """)
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%shell
+ # python -m yolox.tools.train -n yolox-nano -d 1 -b 64 --fp16 --logger wandb \
+ # wandb-project yolox-colab \
+ # wandb-log_checkpoints True \
+ # wandb-num_eval_images 3 \
+ # eval_interval 1 \
+ # print_interval 1 \
+ # max_epoch 10
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ [This W&B dashboard](https://wandb.ai/manan-goel/yolox-nano) shows the visualization of how all the metrics vary over time. Average precision is logged against epoch and losses are logged against the step.
+
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The checkpoints are logged to the wandb dashboard and tagged with epoch and if it is the best model. Along with that, metadata is also provided which consists of the optimizer state and the average precision on the validation set from that model.
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The first `num_eval_images` from the validation set are logged to the dasboard and the corresponding predictions are logged to the dashboard for visualization along with the confidence scores!
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Finetuning a pretrained model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ You can also finetune a pretrained model on a [custom dataset.](https://github.com/manangoel99/YOLOX/blob/WandbLogger/docs/train_custom_data.md) In this case we continue working on the subset of COCO.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _api = wandb.Api()
+ _artifact = _api.artifact('manan-goel/YOLOX-coco/coco128:latest', type='dataset')
+ _artifact_dir = _artifact.download(root='/content')
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%shell
+ # cd /content
+ # unzip /content/coco128.zip
+ # mv /content/coco128 /content/YOLOX/datasets/COCO
+ # cd /content/YOLOX
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ We will finetune the trained model from the previous step. To do that we download the logged artifact using the wandb API.
+ """)
+ return
+
+
+@app.cell
+def _(wandb):
+ _api = wandb.Api()
+ _artifact = _api.artifact('manan-goel/yolox-nano/run_3ntph3ki_model:best')
+ _artifact.download()
+ return
+
+
+@app.cell
+def _():
+ # magic command not supported in marimo; please file an issue to add support
+ # %%shell
+ # python tools/train.py -f exps/example/custom/nano.py -d 1 -b 64 --fp16 -o -c ./artifacts/run_3ntph3ki_model:v40/model_ckpt.pth --logger wandb \
+ # wandb-project yolox-nano-finetune \
+ # max_epoch 5 \
+ # print_interval 1 \
+ # eval_interval 1
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Using the trained model
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ The cell below runs detection using a pretrained model on a given image.
+ """)
+ return
+
+
+@app.cell
+def _(os, subprocess):
+ os.chdir('YOLOX')
+ subprocess.call(['wget', 'https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth'])
+ #! wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth
+ #! python tools/demo.py image -n yolox-s -c yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --device gpu --save_result
+ subprocess.call(['python', 'tools/demo.py', 'image', '-n', 'yolox-s', '-c', 'yolox_s.pth', '--path', 'assets/dog.jpg', '--conf', '0.25', '--nms', '0.45', '--tsize', '640', '--device', 'gpu', '--save_result'])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Input | Annotated Image
+ :-------------------------:|:-------------------------:
+  | 
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Resources 📚
+
+ * [W&B and YOLOX Documentation](https://docs.wandb.ai/guides/integrations/other/yolox) contains a few tips for taking most advantage of W&B.
+ * More YOLOX documentation is available [here](https://yolox.readthedocs.io/en/latest/)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Questions about W&B❓
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ If you have any questions about using W&B to track your model performance and predictions, please contact support@wandb.com
+ """)
+ return
+
+
+if __name__ == "__main__":
+ app.run()