eval written

main.py

@@ -3,18 +3,21 @@
 import os
 from distutils.util import strtobool
 from pathlib import Path
-from argparse import ArgumentParser
+from argparse import ArgumentParser, Namespace
 
 import warnings
 
 import torch
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from torch.utils.data import DataLoader
 
 from lib.modules.utils import LightningBaseModule
 from lib.utils.config import Config
 from lib.utils.logging import Logger
 
+from lib.evaluation.classification import ROCEvaluation
+
 warnings.filterwarnings('ignore', category=FutureWarning)
 warnings.filterwarnings('ignore', category=UserWarning)
 
@@ -31,7 +34,7 @@ main_arg_parser.add_argument("--main_eval", type=strtobool, default=False, help=
 main_arg_parser.add_argument("--main_seed", type=int, default=69, help="")
 
 # Data Parameters
-main_arg_parser.add_argument("--data_worker", type=int, default=0, help="")
+main_arg_parser.add_argument("--data_worker", type=int, default=10, help="")
 main_arg_parser.add_argument("--data_batchsize", type=int, default=100, help="")
 main_arg_parser.add_argument("--data_root", type=str, default='/data/', help="")
 main_arg_parser.add_argument("--data_map_root", type=str, default='/res/maps', help="")
@@ -61,16 +64,15 @@ main_arg_parser.add_argument("--project_owner", type=str, default='si11ium', hel
 main_arg_parser.add_argument("--project_neptune_key", type=str, default=os.getenv('NEPTUNE_KEY'), help="")
 
 # Parse it
-args = main_arg_parser.parse_args()
-config = Config.read_namespace(args)
+args: Namespace = main_arg_parser.parse_args()
 
 
-if __name__ == "__main__":
+def run_lightning_loop(config_obj):
 
     # Logging
-    # =============================================================================
+    # ================================================================================
     # Logger
-    with Logger(config) as logger:
+    with Logger(config_obj) as logger:
         # Callbacks
         # =============================================================================
         # Checkpoint Saving
@@ -90,12 +92,12 @@ if __name__ == "__main__":
         # Model
         # =============================================================================
         # Init
-        model: LightningBaseModule = config.model_class(config.model_paramters)
+        model: LightningBaseModule = config_obj.model_class(config_obj.model_paramters)
         model.init_weights()
 
         # Trainer
         # =============================================================================
-        trainer = Trainer(max_epochs=config.train.epochs,
+        trainer = Trainer(max_epochs=config_obj.train.epochs,
                           show_progress_bar=True,
                           weights_save_path=logger.log_dir,
                           gpus=[0] if torch.cuda.is_available() else None,
@@ -103,15 +105,23 @@ if __name__ == "__main__":
                           log_save_interval=(model.data_len * 0.04),  # TODO: Better Value / Setting
                           checkpoint_callback=checkpoint_callback,
                           logger=logger,
-                          fast_dev_run=config.main.debug,
+                          fast_dev_run=config_obj.main.debug,
                           early_stop_callback=None
                           )
 
-        # Train it
-        trainer.fit(model)
+        # Train It
+        trainer.fit(model,)
 
         # Save the last state & all parameters
         trainer.save_checkpoint(logger.log_dir / 'weights.ckpt')
         model.save_to_disk(logger.log_dir)
-    pass
-    # TODO: Eval here!
+
+        # Evaluate It
+        trainer.test()
+    return model
+
+
+if __name__ == "__main__":
+
+    config = Config.read_namespace(args)
+    trained_model = run_lightning_loop(config)