Added model registry and moved training files into src/ folder

.gitignore

@@ -1,12 +1,15 @@
 # Data and Model Checkpoints
 data/
 models/
-frontend/.streamlit/secrets.toml
+
+# W&B log files  
+wandb/
 
 # Cloud authentication keys
 keys/
 gha-creds-*.json
 .env
+frontend/.streamlit/secrets.toml
 
 # IntelliJ IDEs
 .idea/

cloudbuild.yaml

@@ -64,6 +64,7 @@ steps:
         "push",
         "europe-west1-docker.pkg.dev/$PROJECT_ID/training-images/pokemon-trainer",
       ]
+  # Run Lightning Executor
   - name: "gcr.io/cloud-builders/docker"
     id: RunTrainingPipeline
     waitFor:

training/configurations/Configuration.py → training/src/configurations/Configuration.py

@@ -15,7 +15,8 @@ class Configuration:
     local_dataset_path = Path("./data/")
     training_bucket_name = "zhaw_algorithmic_quartet_training_images"
 
-    output_dir = "./models/ddpm-pokemon-128"
+    model_name = "pokemon-generator"
+    output_dir = f"./models/{model_name}"
 
     push_to_hub = False  # whether to upload the saved model to the HF Hub
     hub_private_repo = False

training/configurations/WandBConfig.py → training/src/configurations/WandBConfig.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import wandb
 import yaml
 
@@ -31,7 +33,7 @@ def create_run(self, experiment_name=None):
         """
         # Append mode to experiment name if provided
         if experiment_name:
-            experiment_name = f"{self.mode.upper()}_{experiment_name}"
+            experiment_name = experiment_name
         else:
             experiment_name = self.mode.upper()
 
@@ -71,6 +73,18 @@ def log_image(self, image, caption):
         if self.run and self.mode != "production":
             self.run.log({"image": [wandb.Image(image, caption=caption)]})
 
+    def link_model(self, model_path: str, model_name: str):
+        """
+        Uploads a model to the W&B model registry. Only works if the given path actually
+        exists and the W&B run was initalized.
+
+        Args:
+        model_path: Path to the file that stores the trained model.
+        model_name: How the model should be called in the registry.
+        """
+        if self.run and Path(model_path).exists():
+            self.run.link_model(path=model_path, registered_model_name=model_name)
+
     def finish_run(self):
         """
         Finishes the active W&B run.

training/pipeline.py → training/src/pipeline.py

@@ -1,5 +1,3 @@
-import argparse
-
 import torch
 from configurations.Configuration import Configuration
 from configurations.create_config import create_config_from_arguments
@@ -25,9 +23,7 @@ def prepare_data(config: Configuration):
                 config.training_bucket_name, max_results=config.num_images
             )
 
-        return load_dataset(
-            str(config.local_dataset_path.resolve()), split="train[0:6]"
-        )
+        return load_dataset(str(config.local_dataset_path.resolve()), split="train")
 
     if not config.local_dataset_path.exists():
         download_bucket_with_transfer_manager(config.training_bucket_name)
@@ -113,6 +109,11 @@ def transform(examples):
         wandb_config,
     )
 
+    wandb_config.link_model(
+        config.output_dir,
+        config.model_name,
+    )
+
 
 if __name__ == "__main__":
     config = create_config_from_arguments()