diff --git a/colabs/wandb_registry/zoo_wandb.ipynb b/colabs/wandb_registry/zoo_wandb.ipynb index 4f61a044..08f74174 100644 --- a/colabs/wandb_registry/zoo_wandb.ipynb +++ b/colabs/wandb_registry/zoo_wandb.ipynb @@ -2,6 +2,7 @@ "cells": [ { "cell_type": "markdown", + "id": "acc18434", "metadata": {}, "source": [ "\"Open\n", @@ -183,9 +184,21 @@ "id": "44c2b1fb-a1af-4311-aae9-86497d00ffda", "metadata": {}, "source": [ + "## Create a registry\n", + "Let's save our dataset to the W&B Registry as an artifact so that future teammates can reproduce our experiments, analyze the data, etc. at a future date. To do this, let's first create a *registry*:\n", + "\n", + "1. Navigate to your organization's Registry App\n", + "2. Click on the **Create registry** button\n", + "3. Provide a name for your registry in the **Name** field. For this notebook, we use \"Zoo\".\n", + "4. Select \"Organization\" from the **Registry visibility** dropdown.\n", + "5. Select \"All types\" from the **Accepted artifact types** dropdown. \n", + "6. Click **Create registry**.\n", + "\n", + "Now that you have a registry created, let's add our dataset to our registry.\n", + "\n", "## Track and publish dataset \n", "\n", - "Within the Dataset registry we will create a collection called \"zoo-dataset-tensors\". A *collection* is a set of linked artifact versions in a registry. \n", + "Within the \"Zoo\" registry we will create a collection called \"zoo-dataset-tensors\". A *collection* is a set of linked artifact versions in a registry. \n", "\n", "To create a collection we need to do two things:\n", "1. Specify the collection and registry we want to link our artifact version to. To do this, we specify a \"target path\" for our artifact version.\n", @@ -210,7 +223,7 @@ "source": [ "### Publish dataset to registry\n", "\n", - "Let's publish our dataset to the Dataset registry in a collection called \"zoo-dataset-tensors\". To do this, we will \n", + "Let's publish our dataset to the \"Zoo\" registry in a collection called \"zoo-dataset-tensors\". To do this, we will \n", "\n", "1. Get or create the target path. (For this notebook we will programmatically create the target path).\n", "1. Initialize a run\n", @@ -228,7 +241,7 @@ "metadata": {}, "outputs": [], "source": [ - "REGISTRY_NAME = \"Dataset\"\n", + "REGISTRY_NAME = \"Zoo\"\n", "COLLECTION_NAME = \"zoo-dataset-tensors\"\n", "\n", "# Path to link the artifact to a collection\n", @@ -240,7 +253,7 @@ "id": "66d03210-1de4-44a4-9773-67b4f5d71dfc", "metadata": {}, "source": [ - "Now that we have the target path, let's publish the dataset to the \"Dataset\" registry. In the following code cell, ensure to replace the values enclosed in `<>` with your team's entity:" + "Now that we have the target path, let's publish the dataset to the \"Zoo\" Registry. In the following code cell, ensure to replace the values enclosed in `<>` with your team's entity:" ] }, { @@ -253,24 +266,22 @@ "TEAM_ENTITY = \"\"\n", "PROJECT = \"zoo_experiment\"\n", "\n", - "run = wandb.init(\n", + "with wandb.init(\n", " entity=TEAM_ENTITY,\n", " project=PROJECT,\n", " job_type=\"publish_dataset\"\n", - ")\n", + ") as run:\n", "\n", - "artifact = wandb.Artifact(\n", - " name=\"zoo_dataset\",\n", - " type=\"dataset\", \n", - " description=\"Processed dataset and labels.\"\n", - ")\n", - "\n", - "artifact.add_file(local_path=\"zoo_dataset.pt\", name=\"zoo_dataset\")\n", - "artifact.add_file(local_path=\"zoo_labels.pt\", name=\"zoo_labels\")\n", + " artifact = wandb.Artifact(\n", + " name=\"zoo_dataset\",\n", + " type=\"dataset\", \n", + " description=\"Processed dataset and labels.\"\n", + " )\n", "\n", - "run.link_artifact(artifact=artifact, target_path=dataset_target_path)\n", + " artifact.add_file(local_path=\"zoo_dataset.pt\", name=\"zoo_dataset\")\n", + " artifact.add_file(local_path=\"zoo_labels.pt\", name=\"zoo_labels\")\n", "\n", - "run.finish()" + " run.link_artifact(artifact=artifact, target_path=dataset_target_path)" ] }, { @@ -325,7 +336,7 @@ "id": "028d0992-15a9-4687-8c3b-4218bf70adb9", "metadata": {}, "source": [ - "Next, let's publish this dataset into a different collection within the Dataset registry called \"zoo-dataset-tensors-split\":" + "Next, let's publish this dataset into a different collection within the registry called \"zoo-dataset-tensors-split\":" ] }, { @@ -335,35 +346,32 @@ "metadata": {}, "outputs": [], "source": [ - "run = wandb.init(\n", + "with wandb.init(\n", " entity=TEAM_ENTITY,\n", " project=PROJECT,\n", " job_type=\"publish_split_dataset\", \n", " config=config\n", - ")\n", + ") as run:\n", "\n", - "# Let's add a description to let others know which file to use in future experiments\n", - "artifact = wandb.Artifact(\n", - " name=\"split_zoo_dataset\",\n", - " type=\"dataset\", \n", - " description=\"Artifact contains `zoo_dataset` split into 4 datasets. \\\n", - " For training, use `zoo_dataset_X_train` and `zoo_labels_y_train`. \\\n", - " For testing, use `zoo_dataset_X_test` and `zoo_labels_y_test`.\"\n", - ")\n", - "\n", - "artifact.add_file(local_path=\"zoo_dataset_X_train.pt\", name=\"zoo_dataset_X_train\")\n", - "artifact.add_file(local_path=\"zoo_labels_y_train.pt\", name=\"zoo_labels_y_train\")\n", - "artifact.add_file(local_path=\"zoo_dataset_X_test.pt\", name=\"zoo_dataset_X_test\")\n", - "artifact.add_file(local_path=\"zoo_labels_y_test.pt\", name=\"zoo_labels_y_test\")\n", + " # Let's add a description to let others know which file to use in future experiments\n", + " artifact = wandb.Artifact(\n", + " name=\"split_zoo_dataset\",\n", + " type=\"dataset\", \n", + " description=\"Artifact contains `zoo_dataset` split into 4 datasets. \\\n", + " For training, use `zoo_dataset_X_train` and `zoo_labels_y_train`. \\\n", + " For testing, use `zoo_dataset_X_test` and `zoo_labels_y_test`.\"\n", + " )\n", "\n", - "# Create a target path for our artifact in the registry\n", - "REGISTRY_NAME = \"Dataset\"\n", - "COLLECTION_NAME = \"zoo-dataset-tensors-split\"\n", - "target_dataset_path=f\"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}\"\n", + " artifact.add_file(local_path=\"zoo_dataset_X_train.pt\", name=\"zoo_dataset_X_train\")\n", + " artifact.add_file(local_path=\"zoo_labels_y_train.pt\", name=\"zoo_labels_y_train\")\n", + " artifact.add_file(local_path=\"zoo_dataset_X_test.pt\", name=\"zoo_dataset_X_test\")\n", + " artifact.add_file(local_path=\"zoo_labels_y_test.pt\", name=\"zoo_labels_y_test\")\n", "\n", - "run.link_artifact(artifact=artifact, target_path=target_dataset_path)\n", + " # Create a target path for our artifact in the registry\n", + " COLLECTION_NAME = \"zoo-dataset-tensors-split\"\n", + " target_dataset_path=f\"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}\"\n", "\n", - "run.finish()" + " run.link_artifact(artifact=artifact, target_path=target_dataset_path)" ] }, { @@ -374,7 +382,7 @@ "We can verify we correctly linked our artifact to our desired collection and registry with W&B App UI: \n", "\n", "1. Navigate to the Registry App\n", - "2. Select on the Dataset registry\n", + "2. Select the \"Zoo\" Registry\n", "3. Click **View details** \"zoo-dataset-tensors-split\" collection\n", "4. Click the **View** button next to the artifact version\n", "5. Select the **Files** tab\n", @@ -458,7 +466,7 @@ "\n", "Next, let's train a model using the training data we published to the registry earlier in this notebook. After we train the model, we will publish that model to W&B.\n", "\n", - "To do this, let's first get the artifact we published to the \"Dataset\" registry. To retrieve an artifact from a registry, we need to know the name of that artifact. The name of an artifact in a registry consists of the prefix `wandb-registry-`, the name of the registry, the name of the collection, and the artifact version:\n", + "To do this, let's first get the artifact we published to the \"Zoo\" registry. To retrieve an artifact from a registry, we need to know the name of that artifact. The name of an artifact in a registry consists of the prefix `wandb-registry-`, the name of the registry, the name of the collection, and the artifact version:\n", "\n", "```python\n", "# Artifact name/filepath for downloading and using artifacts published to a registry\n", @@ -482,66 +490,64 @@ "metadata": {}, "outputs": [], "source": [ - "run = wandb.init(entity = TEAM_ENTITY, project = PROJECT, job_type = \"training\", config = hyperparameter_config)\n", + "with wandb.init(entity = TEAM_ENTITY, project = PROJECT, job_type = \"training\", config = hyperparameter_config) as run:\n", "\n", - "# Get dataset artifacts from registry\n", - "VERSION = 0\n", - "artifact_name = f\"wandb-registry-{REGISTRY_NAME.lower()}/{COLLECTION_NAME}:v{VERSION}\"\n", - "dataset_artifact = run.use_artifact(artifact_or_name=artifact_name)\n", - "\n", - "# Download only the training data\n", - "X_train_path = dataset_artifact.download(path_prefix=\"zoo_dataset_X_train\")\n", - "y_train_path = dataset_artifact.download(path_prefix=\"zoo_labels_y_train\")\n", - "\n", - "# Load data as tensors \n", - "X_train = torch.load(f=X_train_path+\"/zoo_dataset_X_train\")\n", - "y_train = torch.load(f=y_train_path+\"/zoo_labels_y_train\")\n", - "\n", - "# Set initial dummy loss value to compare to in training loop\n", - "prev_best_loss = 1e10 \n", - "\n", - "# Training loop\n", - "for e in range(hyperparameter_config[\"epochs\"] + 1):\n", - " pred = model(X_train)\n", - " loss = loss_fn(pred, y_train.squeeze(1))\n", - " \n", - " loss.backward()\n", - " optimizer.step()\n", - " optimizer.zero_grad()\n", - "\n", - " wandb.log({\n", - " \"train/epoch_ndx\": e,\n", - " \"train/train_loss\": loss\n", - " })\n", - "\n", - " # Checkpoint/save model if loss improves\n", - " if (e % 100 == 0) and (loss <= prev_best_loss):\n", - " print(\"epoch: \", e, \"loss:\", loss.item())\n", - " \n", - " PATH = 'zoo_wandb.pth' \n", - " torch.save(model.state_dict(), PATH)\n", - "\n", - " model_artifact_name = f\"zoo-{wandb.run.id}\"\n", - " artifact = wandb.Artifact(\n", - " name=model_artifact_name,\n", - " type=\"model\",\n", - " metadata={\n", - " \"num_classes\": 7,\n", - " \"model_type\": wandb.config[\"model_type\"]\n", - " }\n", - " )\n", + " # Get dataset artifacts from registry\n", + " VERSION = 0\n", + " artifact_name = f\"wandb-registry-{REGISTRY_NAME.lower()}/{COLLECTION_NAME}:v{VERSION}\"\n", + " dataset_artifact = run.use_artifact(artifact_or_name=artifact_name)\n", + "\n", + " # Download only the training data\n", + " X_train_path = dataset_artifact.download(path_prefix=\"zoo_dataset_X_train\")\n", + " y_train_path = dataset_artifact.download(path_prefix=\"zoo_labels_y_train\")\n", + "\n", + " # Load data as tensors \n", + " X_train = torch.load(f=X_train_path+\"/zoo_dataset_X_train\")\n", + " y_train = torch.load(f=y_train_path+\"/zoo_labels_y_train\")\n", "\n", + " # Set initial dummy loss value to compare to in training loop\n", + " prev_best_loss = 1e10 \n", "\n", - " # Store new best loss\n", - " prev_best_loss = loss\n", + " # Training loop\n", + " for e in range(hyperparameter_config[\"epochs\"] + 1):\n", + " pred = model(X_train)\n", + " loss = loss_fn(pred, y_train.squeeze(1))\n", + " \n", + " loss.backward()\n", + " optimizer.step()\n", + " optimizer.zero_grad()\n", "\n", - "print(f'Saving model artifact {model_artifact_name}')\n", + " run.log({\n", + " \"train/epoch_ndx\": e,\n", + " \"train/train_loss\": loss\n", + " })\n", "\n", - "# Add saved model to artifact\n", - "artifact.add_file(PATH)\n", - "artifact.save()\n", + " # Checkpoint/save model if loss improves\n", + " if (e % 100 == 0) and (loss <= prev_best_loss):\n", + " print(\"epoch: \", e, \"loss:\", loss.item())\n", + " \n", + " PATH = 'zoo_wandb.pth' \n", + " torch.save(model.state_dict(), PATH)\n", "\n", - "run.finish()" + " model_artifact_name = f\"zoo-{wandb.run.id}\"\n", + " artifact = wandb.Artifact(\n", + " name=model_artifact_name,\n", + " type=\"model\",\n", + " metadata={\n", + " \"num_classes\": 7,\n", + " \"model_type\": wandb.config[\"model_type\"]\n", + " }\n", + " )\n", + "\n", + "\n", + " # Store new best loss\n", + " prev_best_loss = loss\n", + "\n", + " print(f'Saving model artifact {model_artifact_name}')\n", + "\n", + " # Add saved model to artifact\n", + " artifact.add_file(PATH)\n", + " artifact.save()" ] }, { @@ -551,7 +557,7 @@ "source": [ "The preceding cell might look intimidating. Let's break it down:\n", "\n", - "* First, we download the dataset from the Dataset registry and load it as a tensor\n", + "* First, we download the dataset from the registry and load it as a tensor\n", "* Next, we create a simple training loop\n", " * Within the training loop we log the loss for each step\n", " * We checkpoint(save) the model every time the remainder of the epoch divided by 100 is 0 and the loss is lower than the previously recorded loss.\n", @@ -568,7 +574,7 @@ "metadata": {}, "source": [ "## Publish model to registry\n", - "Let's make this model artifact available to other users in our organization. To do this, we will create a collection within the Model registry.\n", + "Let's make this model artifact available to other users in our organization. To do this, we will create a collection within our registry.\n", "\n", "To create a collection within a registry, we need to know the full name of the artifact. The full name of the artifact consists of the name we provided to it when we created the Artifact object and its location within our team's project.\n", "\n", @@ -606,9 +612,9 @@ "id": "6aaa9221-f6ac-454b-9a48-47597f47a572", "metadata": {}, "source": [ - "Now that we have the full name of our model artifact. Let's publish it to the model registry.\n", + "Now that we have the full name of our model artifact. Let's publish it to the registry.\n", "\n", - "Similar to how we created a target path when we published our dataset artifact to the Dataset registry, let's create the target path for our model artifact. The target path tells W&B the collection and registry (Model registry) to link our artifact version to. \n", + "Similar to how we created a target path when we published our dataset artifact to the registry, let's create the target path for our model artifact. The target path tells W&B the collection and registry (Zoo registry) to link our artifact version to. \n", "\n", "As a reminder, the target path to link an artifact to a registry consists of:\n", "\n", @@ -625,7 +631,6 @@ "metadata": {}, "outputs": [], "source": [ - "REGISTRY_NAME = \"Model\"\n", "COLLECTION_NAME = \"Zoo_Classifier_Models\"\n", "\n", "target_path = f\"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}\"\n", @@ -647,10 +652,9 @@ "metadata": {}, "outputs": [], "source": [ - "run = wandb.init(entity=TEAM_ENTITY, project=PROJECT)\n", - "model_artifact = run.use_artifact(artifact_or_name=artifact_name, type=\"model\")\n", - "run.link_artifact(artifact=model_artifact, target_path=target_path)\n", - "run.finish()" + "with wandb.init(entity=TEAM_ENTITY, project=PROJECT) as run:\n", + " model_artifact = run.use_artifact(artifact_or_name=artifact_name, type=\"model\")\n", + " run.link_artifact(artifact=model_artifact, target_path=target_path)" ] }, { @@ -731,7 +735,6 @@ "outputs": [], "source": [ "# Create model artifact name\n", - "REGISTRY_NAME = \"model\"\n", "COLLECTION_NAME = \"Zoo_Classifier_Models\"\n", "VERSION = 0\n", "\n", @@ -760,9 +763,9 @@ "DIFFERENT_TEAM_ENTITY = \"\"\n", "DIFFERENT_PROJECT = \"Check_Zoo_Model\"\n", "\n", - "run = wandb.init(entity=DIFFERENT_TEAM_ENTITY, project=DIFFERENT_PROJECT)\n", - "registry_model = run.use_artifact(artifact_or_name=model_artifact_name)\n", - "local_model_path = registry_model.download()" + "with wandb.init(entity=DIFFERENT_TEAM_ENTITY, project=DIFFERENT_PROJECT) as run:\n", + " registry_model = run.use_artifact(artifact_or_name=model_artifact_name)\n", + " local_model_path = registry_model.download()" ] }, { @@ -802,9 +805,9 @@ "id": "93895213-e75e-46e6-b3a8-8017ccaff9e2", "metadata": {}, "source": [ - "### Get test dataset from Dataset registry\n", + "### Get test dataset from the registry\n", "\n", - "Let's get the test dataset from our registry. Similar to the above code block, we will specify the full name of the artifact version we want from our Dataset registry." + "Let's get the test dataset from our registry. Similar to the above code block, we will specify the full name of the artifact version we want from our registry." ] }, { @@ -815,7 +818,6 @@ "outputs": [], "source": [ "# Create dataset artifact name\n", - "REGISTRY_NAME = \"dataset\"\n", "COLLECTION_NAME = \"zoo-dataset-tensors-split\"\n", "VERSION = 0\n", "\n", @@ -830,27 +832,17 @@ "metadata": {}, "outputs": [], "source": [ - "run = wandb.init(entity=DIFFERENT_TEAM_ENTITY, project=DIFFERENT_PROJECT)\n", - "dataset_artifact = run.use_artifact(artifact_or_name=data_artifact_name, type=\"dataset\")\n", - "local_dataset_path = dataset_artifact.download()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "77089d70-5248-4946-be78-ce4a87bb58ad", - "metadata": {}, - "outputs": [], - "source": [ - "# Test data and label filenames\n", - "test_data_filename = \"zoo_dataset_X_test\"\n", - "test_labels_filename = \"zoo_labels_y_test\" \n", + "with wandb.init(entity=DIFFERENT_TEAM_ENTITY, project=DIFFERENT_PROJECT) as run:\n", + " dataset_artifact = run.use_artifact(artifact_or_name=data_artifact_name, type=\"dataset\")\n", + " local_dataset_path = dataset_artifact.download()\n", "\n", - "# Load dataset and labels into notebook\n", - "loaded_data = torch.load(f\"{local_dataset_path}/{test_data_filename}\")\n", - "loaded_labels = torch.load(f\"{local_dataset_path}/{test_labels_filename}\")\n", + " # Test data and label filenames\n", + " test_data_filename = \"zoo_dataset_X_test\"\n", + " test_labels_filename = \"zoo_labels_y_test\" \n", "\n", - "run.finish()" + " # Load dataset and labels into notebook\n", + " loaded_data = torch.load(f\"{local_dataset_path}/{test_data_filename}\")\n", + " loaded_labels = torch.load(f\"{local_dataset_path}/{test_labels_filename}\")" ] }, { @@ -1023,18 +1015,7 @@ ] } ], - "metadata": { - "accelerator": "GPU", - "colab": { - "include_colab_link": true, - "provenance": [], - "toc_visible": true - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - } - }, + "metadata": {}, "nbformat": 4, "nbformat_minor": 5 }