Merge pull request #32 from enesozeren/bugfix_vertex

Bugfix vertex

Author: enesozeren

cloudbuild/config_gpu.yaml

@@ -9,4 +9,4 @@ workerPoolSpecs:
     env:
       - name: WANDB_API_KEY
         value: $WANDB_API_KEY
-    args: ["--config", "config/config-defaults-sweep.yaml"]
+    args: ["--config", "config/config-defaults.yaml"]

mlops_project/config/config-defaults.yaml

@@ -1,5 +1,5 @@
 program: train_model.py
-name: sweep_XXX # Change sweep name
+name: sweep_best # Change sweep name
 method: grid
 metric:
   goal: minimize
@@ -8,6 +8,6 @@ parameters:
   BATCH_SIZE:
     values: [16]
   EPOCHS:
-    values: [5]
+    values: [7]
   LEARNING_RATE:
     values: [0.000005]

mlops_project/models/.gitkeep

@@ -27,6 +27,12 @@
 train_set = TensorDataset(train_token_ids, train_attention_masks, train_labels)
 val_set = TensorDataset(val_token_ids, val_attention_masks, val_labels)
 
+CLOUD_BUCKET = "data_bucket_lmu"
+checkpoint_path = (
+    os.path.join("/gcs", CLOUD_BUCKET, "checkpoints")
+    if os.path.exists("/gcs/data_bucket_lmu/")
+    else "mlops_project/checkpoints"
+)
 
 # Reproducibility
 seed_everything(47, workers=True)
@@ -62,13 +68,15 @@ def main():
     train_dataloader = DataLoader(
         train_set,
         worker_init_fn=seed_worker,
+        num_workers=7,
         generator=g,
         sampler=RandomSampler(train_set),
         batch_size=wandb.config.BATCH_SIZE,
     )
     validation_dataloader = DataLoader(
         val_set,
         worker_init_fn=seed_worker,
+        num_workers=7,
         generator=g,
         sampler=SequentialSampler(val_set),
         batch_size=wandb.config.BATCH_SIZE,
@@ -78,7 +86,7 @@ def main():
     model = HatespeechModel(wandb.config.LEARNING_RATE)
 
     checkpoint_callback = ModelCheckpoint(
-        monitor="val_loss", dirpath="mlops_project/checkpoints", filename="best-checkpoint", save_top_k=1, mode="min"
+        monitor="val_loss", dirpath=checkpoint_path, filename="best-checkpoint", save_top_k=1, mode="min"
     )
     # early_stopping_callback = EarlyStopping(monitor="val_loss", patience=3, verbose=True, mode="min")
 
@@ -97,12 +105,7 @@ def main():
     # Train the model
     trainer.fit(model, train_dataloader, validation_dataloader)
     # save best model as model weights
-    CLOUD_BUCKET = "data_bucket_lmu"
-    checkpoint_path = (
-        os.path.join("/gcs", CLOUD_BUCKET, "checkpoints")
-        if os.path.exists("/gcs/data_bucket_lmu/")
-        else "mlops_project/checkpoints"
-    )
+
     checkpoint = torch.load(os.path.join(checkpoint_path, "best-checkpoint.ckpt"))
     state = {key[6:]: value for key, value in checkpoint["state_dict"].items()}
     weight_path = os.path.join(checkpoint_path, "best-checkpoint.pth")