main

+# -*- coding: utf-8 -*-
+"""main.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/1X0vDvK01A8_3JbSAu1whlPD-rBx0JlM-
+"""
+
+import argparse
+import importlib
+import json
+import logging
+import os
+import pprint
+import sys
+
+import dill
+import torch
+import wandb
+from box import Box
+from torch.utils.data import DataLoader
+
+from src.common.dataset import get_dataset
+from lib.utils import logging as logging_utils, os as os_utils, optimizer as optimizer_utils
+from lib.base_trainer import Trainer
+import easydict
+
+def parser_setup():
+    # define argparsers
+
+    str2bool = os_utils.str2bool
+    listorstr = os_utils.listorstr
+    
+    parser = easydict.EasyDict({
+        "debug":False,
+        "config":None,
+        "seed":0,
+        "wandb_use":False,
+        "wandb_run_id":None,
+        "wandb.watch":False,
+        "project":"brain-age",
+        "exp_name":None,
+        "device":"cuda",
+        "result_folder":"a",
+        "mode":["test", "train"],
+        "statefile":None,
+
+        "data" : {
+            "name":"brain_age",
+            "root_path":"**root path",
+            "train_csv":"**train csv",
+            "valid_csv":"**valid csv",
+            "test_csv":"**test csv",
+            "feat_csv":None,
+            "train_num_sample":-1,
+            "frame_dim":1,
+            "frame_keep_style":"random",
+            "frame_keep_fraction":1,
+            "impute":"drop",
+        },
+
+        "model" : {
+            "name":"regression",
+            "arch": {
+                "file":"src/arch/brain_age_3d.py",
+                "lstm_feat_dim":2,
+                "lstm_latent_dim":128,
+                "attn_num_heads":1,
+                "attn_dim":32,
+                "attn_drop":False,
+                "agg_fn":"attention"
+            }
+        },
+        
+        "train":{
+            "batch_size":8,
+
+            "patience":100,
+            "max_epoch":100,
+            "optimizer":"adam",
+            "lr":1e-3,
+            "weight_decay":1e-4,
+            "gradient_norm_clip":-1,
+
+            "save_strategy":["best", "last"],
+            "log_every":100,
+
+            "stopping_criteria":"loss",
+            "stopping_criteria_direction":"lower",
+            "evaluations":None,
+
+            "optimizer_momentum":None,
+
+            "scheduler":None,
+            "scheduler_gamma":None,
+            "scheduler_milestones":None,
+            "scheduler_patience":None,
+            "scheduler_step_size":None,
+            "scheduler_load_on_reduce":None,
+        },
+        
+        "test":{
+            "batch_size":8,
+            "evaluations":None,
+            "eval_model":"best",
+        },
+
+        "_actions":None,
+        "_defaults":None
+    })
+   
+    print(parser.seed)
+    return parser
+
+if __name__ == "__main__":
+    # set seeds etc here
+    torch.backends.cudnn.benchmark = True
+
+    # define logger etc
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
+    logger = logging.getLogger()
+
+    parser = parser_setup()
+    #config = os_utils.parse_args(parser)
+    config = parser
+
+    logger.info("Config:")
+    logger.info(pprint.pformat(config, indent=4))
+
+    os_utils.safe_makedirs(config.result_folder)
+    statefile, run_id, result_folder = os_utils.get_state_params(
+        config.wandb_use, config.wandb_run_id, config.result_folder, config.statefile
+    )
+    config.statefile = statefile
+    config.wandb_run_id = run_id
+    config.result_folder = result_folder
+
+    if statefile is not None:
+        data = torch.load(open(statefile, "rb"), pickle_module=dill)
+        epoch = data["epoch"]
+        if epoch >= config.train.max_epoch:
+            logger.error("Aleady trained upto max epoch; exiting")
+            sys.exit()
+
+    if config.wandb_use:
+        wandb.init(
+            name=config.exp_name if config.exp_name is not None else config.result_folder,
+            config=config.to_dict(),
+            project=config.project,
+            dir=config.result_folder,
+            resume=config.wandb_run_id,
+            id=config.wandb_run_id,
+            sync_tensorboard=True,
+        )
+        logger.info(f"Starting wandb with id {wandb_run_id}")
+
+    # NOTE: WANDB creates git patch so we probably can get rid of this in future
+    os_utils.copy_code("src", config.result_folder, replace=True,)
+    json.dump(
+        config,
+        open(f"{wandb_run.dir if config.wandb_use else config.result_folder}/config.json", "w")
+    )
+
+    logger.info("Getting data and dataloaders")
+    data, meta = get_dataset(**config.data, device=config.device, replace=True, frac=2000)
+
+    # num_workers = max(min(os.cpu_count(), 8), 1)
+    num_workers = os.cpu_count()
+    logger.info(f"Using {num_workers} workers")
+    train_loader = DataLoader(data["train"], shuffle=False, batch_size=config.train.batch_size,
+                              num_workers=num_workers)
+    valid_loader = DataLoader(data["valid"], shuffle=False, batch_size=config.test.batch_size,
+                              num_workers=num_workers)
+    test_loader = DataLoader(data["test"], shuffle=False, batch_size=config.test.batch_size,
+                             num_workers=num_workers)
+    
+    logger.info("Getting model")
+    
+    # load arch module
+    arch_module = importlib.import_module(config.model.arch.file.replace("/", ".")[:-3])
+    model_arch = arch_module.get_arch(
+        input_shape=meta.get("input_shape"), output_size=meta.get("num_class"), 
+        **config.model.arch,
+        slice_dim=config.data.frame_dim
+    )
+
+    # declaring models
+    if config.model.name in "regression":
+        from src.models.regression import Regression
+
+        model = Regression(**model_arch)
+    else:
+        raise Exception("Unknown model")
+
+    model.to(config.device)
+    model.stats()
+
+    if config.wandb_use and config.wandb_watch:
+        wandb_watch(model, log="all")
+
+    # declaring trainer
+    optimizer, scheduler = optimizer_utils.get_optimizer_scheduler(
+        model,
+        lr=config.train.lr,
+        optimizer=config.train.optimizer,
+        opt_params={
+                "weight_decay": config.train.get("weight_decay", 1e-4),
+                "momentum"    : config.train.get("optimizer_momentum", 0.9)
+        },
+        scheduler=config.train.get("scheduler", None),
+        scheduler_params={
+                "gamma"         : config.train.get("scheduler_gamma", 0.1),
+                "milestones"    : config.train.get("scheduler_milestones", [100, 200, 300]),
+                "patience"      : config.train.get("scheduler_patience", 100),
+                "step_size"     : config.train.get("scheduler_step_size", 100),
+                "load_on_reduce": config.train.get("scheduler_load_on_reduce"),
+                "mode"          : "max" if config.train.get(
+                    "stopping_criteria_direction") == "bigger" else "min"
+        },
+    )
+    trainer = Trainer(model, optimizer, scheduler=scheduler, statefile=config.statefile,
+                      result_dir=config.result_folder, log_every=config.train.log_every,
+                      save_strategy=config.train.save_strategy,
+                      patience=config.train.patience,
+                      max_epoch=config.train.max_epoch,
+                      stopping_criteria=config.train.stopping_criteria,
+                      gradient_norm_clip=config.train.gradient_norm_clip,
+                      stopping_criteria_direction=config.train.stopping_criteria_direction,
+                      evaluations=Box({"train": config.train.evaluations,
+                                       "test" : config.test.evaluations}))
+
+    if "train" in config.mode:
+        logger.info("starting training")
+        print(train_loader.dataset)
+        trainer.train(train_loader, valid_loader)
+        logger.info("Training done;")
+
+        # copy current step and write test results to
+        step_to_write = trainer.step
+        step_to_write += 1
+
+        if "test" in config.mode and config.test.eval_model == "best":
+            if os.path.exists(f"{trainer.result_dir}/best_model.pt"):
+                logger.info("Loading best model")
+                trainer.load(f"{trainer.result_dir}/best_model.pt")
+            else:
+                logger.info("eval_model is best, but best model not found ::: evaling last model")
+        else:
+            logger.info("eval model is not best, so skipping loading at end of training")
+
+    if "test" in config.mode:
+        logger.info("evaluating model on test set")
+        logger.info(f"Model was trained upto {trainer.epoch}")
+        # copy current step and write test results to
+        step_to_write = trainer.step
+        step_to_write += 1
+
+        print("<<<<<test>>>>>")
+        loss, aux_loss = trainer.test(train_loader, test_loader)
+        logging_utils.loss_logger_helper(loss, aux_loss, writer=trainer.summary_writer,
+                                         force_print=True, step=step_to_write,
+                                         epoch=trainer.epoch,
+                                         log_every=trainer.log_every, string="test",
+                                         new_line=True)
+
+        print("<<<<<training>>>>>")
+        loss, aux_loss = trainer.test(train_loader, train_loader)
+        logging_utils.loss_logger_helper(loss, aux_loss, writer=trainer.summary_writer,
+                                         force_print=True, step=step_to_write,
+                                         epoch=trainer.epoch,
+                                         log_every=trainer.log_every, string="train_eval",
+                                         new_line=True)
+
+        print("<<<<<validation>>>>>")
+        loss, aux_loss = trainer.test(train_loader, valid_loader)
+        logging_utils.loss_logger_helper(loss, aux_loss, writer=trainer.summary_writer,
+                                         force_print=True, step=step_to_write,
+                                         epoch=trainer.epoch,
+                                         log_every=trainer.log_every, string="valid_eval",
+                                         new_line=True)
\ No newline at end of file