comment wandb

main.py

@@ -9,7 +9,7 @@ from PIL import Image
 import time
 import numpy as np
 import math
-import wandb
+# import wandb
 from datasets import load_dataset
 import os
 
@@ -98,7 +98,7 @@ if __name__ == "__main__":
         "BS": 32,
         "HIDDEN": 128,
         "DROPOUT": 0.01,
-        "WANDB": True,
+        # "WANDB": True,
         "REPO_ID": "TristanKE/RemainingLifespanPredictionFaces",
         # "REPO_ID": "TristanKE/RemainingLifespanPredictionWholeImgs",
         "DINO_MODEL": "dinov2_vitl14_reg",
@@ -107,8 +107,8 @@ if __name__ == "__main__":
         # "DINO_DIM": 1536, #for the larger model
     }
 
-    if cfg["WANDB"]:
-        wandb.init(project="mortpred", config=cfg)
+    # if cfg["WANDB"]:
+    #     wandb.init(project="mortpred", config=cfg)
 
     torch.manual_seed(1)
     ds = LifespanDataset(repo_id=cfg["REPO_ID"],transform=imgtransform)
@@ -147,12 +147,12 @@ if __name__ == "__main__":
 
             tr_nll += loss.item() * imgs.size(0)
             tr_mae += torch.abs(mu.detach() - tgt).sum().item()
-            if cfg["WANDB"]:
-                wandb.log({
-                    "train_nll": loss.item(),
-                    "train_mae": torch.abs(mu.detach() - tgt).mean().item() * ds.lifespan_std,
-                    "train_std": torch.exp(0.5 * logvar).mean().item() * ds.lifespan_std,
-                })
+            # if cfg["WANDB"]:
+            #     wandb.log({
+            #         "train_nll": loss.item(),
+            #         "train_mae": torch.abs(mu.detach() - tgt).mean().item() * ds.lifespan_std,
+            #         "train_std": torch.exp(0.5 * logvar).mean().item() * ds.lifespan_std,
+            #     })
 
         tr_nll /= train_sz
         tr_mae = tr_mae / train_sz * ds.lifespan_std
@@ -172,13 +172,13 @@ if __name__ == "__main__":
 
         print(f"Epoch {epoch+1}/{cfg['N_EPOCHS']} | {time.time()-t0:.1f}s | NLL tr {tr_nll:.3f} / te {te_nll:.3f} | MAE(te) {te_mae:.2f} yrs")
 
-        if cfg["WANDB"]:
-            wandb.log({
-                "train_nll": tr_nll,
-                "test_nll": te_nll,
-                "test_mae_yrs": te_mae,
-                "lr": scheduler.get_last_lr()[0],
-            })
+        # if cfg["WANDB"]:
+        #     wandb.log({
+        #         "train_nll": tr_nll,
+        #         "test_nll": te_nll,
+        #         "test_mae_yrs": te_mae,
+        #         "lr": scheduler.get_last_lr()[0],
+        #     })
 
         scheduler.step()
 
@@ -190,5 +190,5 @@ if __name__ == "__main__":
             torch.save(model.state_dict(), f"savedmodels/dino_finetuned_faces_l1_{cfg['DINO_DIM']}_best.pth")
             print(f"\tNew best model saved (test MAE {te_mae:.3f})")
 
-    if cfg["WANDB"]:
-        wandb.finish()
+    # if cfg["WANDB"]:
+    #     wandb.finish()