From f078c2a593cc73db2ec2b521013daef71c2b89e5 Mon Sep 17 00:00:00 2001 From: Abhishek Aggarwal Date: Thu, 17 Apr 2025 11:35:21 +0200 Subject: [PATCH 1/4] Update image classifier example to track artifacts using W&B - Track the trained model via W&B - Done to test the new W&B Team created in https://mozilla-hub.atlassian.net/browse/DENG-8307 --- examples/image_classifier/image_classifier_flow.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/examples/image_classifier/image_classifier_flow.py b/examples/image_classifier/image_classifier_flow.py index c77ee09..cfcb5f0 100644 --- a/examples/image_classifier/image_classifier_flow.py +++ b/examples/image_classifier/image_classifier_flow.py @@ -83,6 +83,7 @@ def train(self): import wandb import os + tracking_run = {} if not self.offline_wandb: tracking_run = wandb.init(project=os.getenv("WANDB_PROJECT")) wandb_url = tracking_run.get_url() @@ -144,6 +145,16 @@ def train(self): buffer = BytesIO() torch.save(image_classifier_model.state_dict(), buffer) self.model_state_dict_bytes = buffer.getvalue() + + if not self.offline_wandb: + # Save trained model locally and then track it in W&B + torch.save(image_classifier_model.state_dict(), "./trained_model.pt") + model_artifact = wandb.Artifact( + name="trained_image_classifier", type="model" + ) + model_artifact.add_file(local_path="./trained_model.pt") + tracking_run.log_artifact(model_artifact) + self.next(self.evaluate) # Test the model on the test data From 477d76269f873649fe9a75477804d87e985387d0 Mon Sep 17 00:00:00 2001 From: Abhishek Aggarwal Date: Tue, 22 Apr 2025 21:31:37 +0200 Subject: [PATCH 2/4] Updated torchvision and torch dependency versions to latest --- examples/image_classifier/image_classifier_flow.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/examples/image_classifier/image_classifier_flow.py b/examples/image_classifier/image_classifier_flow.py index cfcb5f0..6d4cb9d 100644 --- a/examples/image_classifier/image_classifier_flow.py +++ b/examples/image_classifier/image_classifier_flow.py @@ -33,7 +33,7 @@ class ImageClassifierFlow(FlowSpec): default=True, ) - @pypi(python="3.11.9", packages={"torchvision": "0.19.1"}) + @pypi(python="3.11.9", packages={"torchvision": "0.21.0"}) @card(type="default") @kubernetes @step @@ -62,7 +62,7 @@ def start(self): # Keep @nvidia decorator before @step decorator else the flow fails @pypi( python="3.11.9", - packages={"torch": "2.4.1", "torchvision": "0.19.1", "mozmlops": "0.1.4"}, + packages={"torch": "2.6.0", "torchvision": "0.21.0", "mozmlops": "0.1.4"}, ) @nvidia # @kubernetes @@ -137,8 +137,9 @@ def train(self): running_loss += loss.item() if i % 2000 == 1999: # print every 2000 mini-batches print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}") - # log metrics to wandb - wandb.log({"mini-batches": {i + 1}, "loss": {running_loss / 2000}}) + if not self.offline_wandb: + # log metrics to wandb + wandb.log({"mini-batches": {i + 1}, "loss": {running_loss / 2000}}) running_loss = 0.0 print("Finished Training") @@ -161,8 +162,8 @@ def train(self): @pypi( python="3.11.9", packages={ - "torch": "2.4.1", - "torchvision": "0.19.1", + "torch": "2.6.0", + "torchvision": "0.21.0", }, ) # Check https://docs.metaflow.org/api/step-decorators/kubernetes for details on @kubernetes decorator From b27e1e8cd1df7f7cb469652d90d707bc0a2835c0 Mon Sep 17 00:00:00 2001 From: Abhishek Aggarwal Date: Wed, 23 Apr 2025 11:02:28 +0200 Subject: [PATCH 3/4] Revert changes to image classifier flow example --- .../image_classifier/image_classifier_flow.py | 24 +++++-------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/examples/image_classifier/image_classifier_flow.py b/examples/image_classifier/image_classifier_flow.py index 6d4cb9d..c77ee09 100644 --- a/examples/image_classifier/image_classifier_flow.py +++ b/examples/image_classifier/image_classifier_flow.py @@ -33,7 +33,7 @@ class ImageClassifierFlow(FlowSpec): default=True, ) - @pypi(python="3.11.9", packages={"torchvision": "0.21.0"}) + @pypi(python="3.11.9", packages={"torchvision": "0.19.1"}) @card(type="default") @kubernetes @step @@ -62,7 +62,7 @@ def start(self): # Keep @nvidia decorator before @step decorator else the flow fails @pypi( python="3.11.9", - packages={"torch": "2.6.0", "torchvision": "0.21.0", "mozmlops": "0.1.4"}, + packages={"torch": "2.4.1", "torchvision": "0.19.1", "mozmlops": "0.1.4"}, ) @nvidia # @kubernetes @@ -83,7 +83,6 @@ def train(self): import wandb import os - tracking_run = {} if not self.offline_wandb: tracking_run = wandb.init(project=os.getenv("WANDB_PROJECT")) wandb_url = tracking_run.get_url() @@ -137,33 +136,22 @@ def train(self): running_loss += loss.item() if i % 2000 == 1999: # print every 2000 mini-batches print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}") - if not self.offline_wandb: - # log metrics to wandb - wandb.log({"mini-batches": {i + 1}, "loss": {running_loss / 2000}}) + # log metrics to wandb + wandb.log({"mini-batches": {i + 1}, "loss": {running_loss / 2000}}) running_loss = 0.0 print("Finished Training") buffer = BytesIO() torch.save(image_classifier_model.state_dict(), buffer) self.model_state_dict_bytes = buffer.getvalue() - - if not self.offline_wandb: - # Save trained model locally and then track it in W&B - torch.save(image_classifier_model.state_dict(), "./trained_model.pt") - model_artifact = wandb.Artifact( - name="trained_image_classifier", type="model" - ) - model_artifact.add_file(local_path="./trained_model.pt") - tracking_run.log_artifact(model_artifact) - self.next(self.evaluate) # Test the model on the test data @pypi( python="3.11.9", packages={ - "torch": "2.6.0", - "torchvision": "0.21.0", + "torch": "2.4.1", + "torchvision": "0.19.1", }, ) # Check https://docs.metaflow.org/api/step-decorators/kubernetes for details on @kubernetes decorator From edb06f049aec1241434c94cf28b874823c7ec993 Mon Sep 17 00:00:00 2001 From: Abhishek Aggarwal Date: Tue, 22 Apr 2025 21:33:40 +0200 Subject: [PATCH 4/4] Test a new W&B Team configured to use Mozilla-managed storage - Add a simple image classifer example - Track trained model using W&B --- .../requirements.example_image_classifier.txt | 2 + .../test_wandb_integration.py | 144 ++++++++++++++++++ 2 files changed, 146 insertions(+) create mode 100644 examples/image_classifier/test_wandb_integration.py diff --git a/examples/image_classifier/requirements.example_image_classifier.txt b/examples/image_classifier/requirements.example_image_classifier.txt index 3a0da7e..f448931 100644 --- a/examples/image_classifier/requirements.example_image_classifier.txt +++ b/examples/image_classifier/requirements.example_image_classifier.txt @@ -1 +1,3 @@ mozmlops==0.1.4 +torch==2.6.0 +torchvision==0.21.0 diff --git a/examples/image_classifier/test_wandb_integration.py b/examples/image_classifier/test_wandb_integration.py new file mode 100644 index 0000000..14a2640 --- /dev/null +++ b/examples/image_classifier/test_wandb_integration.py @@ -0,0 +1,144 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at https://mozilla.org/MPL/2.0/. + +import os +import torch +import torch.nn as nn +import torch.optim as optim +from image_classifier_model import ImageClassifierModel +from io import BytesIO +import wandb +import torchvision +import torchvision.transforms as transforms + + +class ImageClassifier(): + # This is an example of a parameter. You can toggle this when you call the flow + # with python template_flow.py run --offline False + offline_wandb = False + + def start(self): + # Download and normalize CIFAR10 + print("downloading and normalizing dataset") + + transform = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), + ] + ) + + self.trainset = torchvision.datasets.CIFAR10( + root="./data", train=True, download=True, transform=transform + ) + self.testset = torchvision.datasets.CIFAR10( + root="./data", train=False, download=True, transform=transform + ) + + tracking_run = {} + if not self.offline_wandb: + wandb_project = os.getenv("WANDB_PROJECT") + tracking_run = wandb.init(project=wandb_project) + wandb_url = tracking_run.get_url() + print(f"Your training run is tracked [here]({wandb_url}).") + + device = torch.device("cpu") + # Check if GPU is available + if torch.cuda.is_available(): + print(os.system("nvidia-smi")) + device = torch.device("cuda") + + print(f"Training on: {device}") + + image_classifier_model = ImageClassifierModel().to(device) + + # Define a Loss function and optimizer + criterion = nn.CrossEntropyLoss() + optimizer = optim.SGD( + image_classifier_model.parameters(), lr=0.001, momentum=0.9 + ) + + # Load train data + batch_size = 4 + trainloader = torch.utils.data.DataLoader( + self.trainset, batch_size=batch_size, shuffle=True, num_workers=2 + ) + + # Start training + num_epochs = 2 + for epoch in range(num_epochs): # loop over the dataset multiple times + running_loss = 0.0 + for i, data in enumerate(trainloader, 0): + # get the inputs; data is a list of [inputs, labels] + inputs, labels = data[0].to(device), data[1].to(device) + + # zero the parameter gradients + optimizer.zero_grad() + + # forward + backward + optimize + outputs = image_classifier_model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + # print statistics + running_loss += loss.item() + if i % 2000 == 1999: # print every 2000 mini-batches + print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}") + if not self.offline_wandb: + # log metrics to wandb + wandb.log({"mini-batches": {i + 1}, "loss": {running_loss / 2000}}) + running_loss = 0.0 + + print("Finished Training") + buffer = BytesIO() + torch.save(image_classifier_model.state_dict(), buffer) + self.model_state_dict_bytes = buffer.getvalue() + + if not self.offline_wandb: + # Save trained model locally and then track it in W&B + print("Tracking trained model via W&B") + torch.save(image_classifier_model.state_dict(), "./trained_model.pt") + model_artifact = wandb.Artifact( + name="trained_image_classifier", type="model" + ) + model_artifact.add_file(local_path="./trained_model.pt") + tracking_run.log_artifact(model_artifact) + + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Evaluating on: {device}") + + image_classifier_model = ImageClassifierModel().to(device) + buffer = BytesIO(self.model_state_dict_bytes) + image_classifier_model.load_state_dict( + torch.load(buffer, map_location=device, weights_only=True) + ) + + correct = 0 + total = 0 + + # load test data + batch_size = 4 + testloader = torch.utils.data.DataLoader( + self.testset, batch_size=batch_size, shuffle=False, num_workers=2 + ) + # since we're not training, we don't need to calculate the gradients for our outputs + with torch.no_grad(): + for data in testloader: + images, labels = data[0].to(device), data[1].to(device) + # calculate outputs by running images through the network + outputs = image_classifier_model(images) + # the class with the highest energy is what we choose as prediction + _, predicted = torch.max(outputs.data, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + + print( + f"Accuracy of the network on the 10000 test images: {100 * correct // total} %" + ) + +if __name__ == "__main__": + image_classifier = ImageClassifier() + image_classifier.start()