fixed config

requirements.txt

@@ -0,0 +1,9 @@
+torch~=2.6.0
+numpy~=2.2.3
+torchvision~=0.21.0
+pillow~=11.1.0
+streamlit~=1.42.0
+wandb~=0.19.6
+tqdm~=4.67.1
+typing_extensions~=4.12.2
+matplotlib~=3.10.0

train_byol.py

@@ -261,41 +261,28 @@ def main(config: dict):
 
 
 if __name__ == "__main__":
-    # parser = argparse.ArgumentParser(description="Train BYOL model")
-    # parser.add_argument("--batch_size", type=int, default=512)
-    # parser.add_argument("--lr", type=float, default=5e-4)
-    # parser.add_argument("--num_epochs", type=int, default=15)
-    # parser.add_argument("--num_train_samples", type=int, default=100000)
-    # parser.add_argument("--num_val_samples", type=int, default=10000)
-    # parser.add_argument("--random_intensity", type=int, default=1)
-    # parser.add_argument("--early_stopping_patience", type=int, default=3)
-    # parser.add_argument("--save_path", type=str, default="best_byol.pth")
-    # args = parser.parse_args()
-
-    # config = {
-    #     "batch_size": args.batch_size,
-    #     "lr": args.lr,
-    #     "num_epochs": args.num_epochs,
-    #     "num_train_samples": args.num_train_samples,
-    #     "num_val_samples": args.num_val_samples,
-    #     "shape_params": {
-    #         "random_intensity": bool(args.random_intensity)
-    #     },
-    #     "early_stopping_patience": args.early_stopping_patience,
-    #     "save_path": args.save_path
-    # }
+    parser = argparse.ArgumentParser(description="Train BYOL model")
+    parser.add_argument("--batch_size", type=int, default=512)
+    parser.add_argument("--lr", type=float, default=5e-4)
+    parser.add_argument("--num_epochs", type=int, default=15)
+    parser.add_argument("--num_train_samples", type=int, default=100000)
+    parser.add_argument("--num_val_samples", type=int, default=10000)
+    parser.add_argument("--random_intensity", type=int, default=1)
+    parser.add_argument("--early_stopping_patience", type=int, default=3)
+    parser.add_argument("--save_path", type=str, default="best_byol.pth")
+    args = parser.parse_args()
 
     config = {
-        "batch_size": 1024,
-        "lr": 1e-3,
-        "num_epochs": 15,
-        "num_train_samples": 100000,
-        "num_val_samples": 10000,
+        "batch_size": args.batch_size,
+        "lr": args.lr,
+        "num_epochs": args.num_epochs,
+        "num_train_samples": args.num_train_samples,
+        "num_val_samples": args.num_val_samples,
         "shape_params": {
-            "random_intensity": True
+            "random_intensity": bool(args.random_intensity)
         },
-        "early_stopping_patience": 3,
-        "save_path": "best_byol.pth"
+        "early_stopping_patience": args.early_stopping_patience,
+        "save_path": args.save_path
     }
 
     main(config)

train_cross_classifier.py

@@ -250,43 +250,30 @@ def main(config):
 
 if __name__ == "__main__":
 
-    # parser = argparse.ArgumentParser(description="Train classifier model")
-    # parser.add_argument("--path_to_encoder", type=str, default="best_byol.pth")
-    # parser.add_argument("--batch_size", type=int, default=256)
-    # parser.add_argument("--lr", type=float, default=8e-5)
-    # parser.add_argument("--weight_decay", type=float, default=1e-4)
-    # parser.add_argument("--step_size", type=int, default=10)
-    # parser.add_argument("--gamma", type=float, default=0.1)
-    # parser.add_argument("--num_epochs", type=int, default=10)
-    # parser.add_argument("--num_train_samples", type=int, default=10000)
-    # parser.add_argument("--num_val_samples", type=int, default=2000)
-    # parser.add_argument("--save_path", type=str, default="best_attention_classifier.pth")
-    # args = parser.parse_args()
-
-    # config = {
-    #     "path_to_encoder": args.path_to_encoder,
-    #     "batch_size": args.batch_size,
-    #     "lr": args.lr,
-    #     "weight_decay": args.weight_decay,
-    #     "step_size": args.step_size,
-    #     "gamma": args.gamma,
-    #     "num_epochs": args.num_epochs,
-    #     "num_train_samples": args.num_train_samples,
-    #     "num_val_samples": args.num_val_samples,
-    #     "save_path": args.save_path,
-    # }
+    parser = argparse.ArgumentParser(description="Train classifier model")
+    parser.add_argument("--path_to_encoder", type=str, default="best_byol.pth")
+    parser.add_argument("--batch_size", type=int, default=256)
+    parser.add_argument("--lr", type=float, default=8e-5)
+    parser.add_argument("--weight_decay", type=float, default=1e-4)
+    parser.add_argument("--step_size", type=int, default=10)
+    parser.add_argument("--gamma", type=float, default=0.1)
+    parser.add_argument("--num_epochs", type=int, default=10)
+    parser.add_argument("--num_train_samples", type=int, default=10000)
+    parser.add_argument("--num_val_samples", type=int, default=2000)
+    parser.add_argument("--save_path", type=str, default="best_attention_classifier.pth")
+    args = parser.parse_args()
 
     config = {
-        "path_to_encoder": "best_byol.pth",
-        "batch_size": 256,
-        "lr": 8e-5,
-        "weight_decay": 1e-4,
-        "step_size": 10,
-        "gamma": 0.1,
-        "num_epochs": 10,
-        "num_train_samples": 10000,
-        "num_val_samples": 2000,
-        "save_path": "best_attention_classifier.pth",
+        "path_to_encoder": args.path_to_encoder,
+        "batch_size": args.batch_size,
+        "lr": args.lr,
+        "weight_decay": args.weight_decay,
+        "step_size": args.step_size,
+        "gamma": args.gamma,
+        "num_epochs": args.num_epochs,
+        "num_train_samples": args.num_train_samples,
+        "num_val_samples": args.num_val_samples,
+        "save_path": args.save_path,
     }
 
     if "shape_params" not in config: