Custom Trainer with COCO eval for val

Author: litluc

src/data_utils.py

@@ -6,7 +6,7 @@
 from config import ProjConfig
 
 
-def register_isaid_truck_data(extra_meta={}, register_val=True, register_test):
+def register_isaid_truck_data(extra_meta={}, register_val=True, register_test=False):
     """
     register project data with name isaid_truck_train/val
     """

src/model_train.py

@@ -1,30 +1,35 @@
 # detectron
+import torch
 import detectron2
 from detectron2.utils.logger import setup_logger
 setup_logger()
-from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, launch
+from detectron2.engine import DefaultTrainer
 from detectron2.config import get_cfg
+from detectron2 import model_zoo
+from detectron2.evaluation import COCOEvaluator
 
 # common imports
 import os
+from datetime import datetime
 
 # custom utilities
 from config import ProjConfig
 from data_utils import register_isaid_truck_data
 
 
-def setup_train_config():
+def setup_train_config(train_data_name, val_data_name=None, output_dir=None):
     """
     Specify the model training configuration
     """
     cfg = get_cfg()
     cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
-    cfg.DATASETS.TRAIN = ("isaid_truck_train",)
-    cfg.DATASETS.TEST = ("isaid_truck_val", )
-    cfg.TEST.EVAL_PERIOD = 1      # how often to eval val
+    cfg.DATASETS.TRAIN = (train_data_name,)
+    if val_data_name:
+        cfg.DATASETS.TEST = (val_data_name, )
+        cfg.TEST.EVAL_PERIOD = 1      # how often to eval val
     cfg.MODEL.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
     cfg.DATALOADER.NUM_WORKERS = 1
-    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  
+    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
     cfg.MODEL.BACKBONE.FREEZE_AT = 2   # freeze the first X stages of backbone
     cfg.SOLVER.IMS_PER_BATCH = 2
     cfg.SOLVER.BASE_LR = 0.00025
@@ -34,18 +39,40 @@ def setup_train_config():
     cfg.SOLVER.CHECKPOINT_PERIOD = 5 # Save a checkpoint after every this number of iterations
     cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # default 512, smaller numbers are faster
     cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class
+    if output_dir:
+        # specify an output with a few key hyper params
+        cfg.OUTPUT_DIR = os.path.join(output_dir, \
+            f'detectron_{datetime.now().strftime("%Y%m%d%H%M%S")}_freeze{cfg.MODEL.BACKBONE.FREEZE_AT}_batchsize{cfg.SOLVER.IMS_PER_BATCH}_lr{cfg.SOLVER.BASE_LR}')
+    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
     return cfg
 
 
+class TrainerWithVal(DefaultTrainer):
+    """
+    Build the appropriate evaluate is needed with train with a validation set
+    """
+
+    @classmethod
+    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
+        """class method for evaluating the validation set"""
+        if output_folder is None:
+            output_folder = os.path.join(cfg.OUTPUT_DIR,"inference")
+        return COCOEvaluator(dataset_name, output_dir=output_folder)
+
+
 def main():
     # configure the data
     proj_config = ProjConfig()
     _ = register_isaid_truck_data(extra_meta={}, register_val=True)
 
-    cfg = setup_train_config()
+    cfg = setup_train_config(
+        proj_config.train_data_name,
+        proj_config.val_data_name,
+        proj_config.model_dir
+    )
     # set up the trainer: wrapper for model training with config
-    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
-    trainer = DefaultTrainer(cfg) 
+
+    trainer = TrainerWithVal(cfg)
     trainer.resume_or_load(resume=False)
     trainer.train()
 
@@ -56,4 +83,4 @@ def main():
 
 # # Look at training curves in tensorboard:
 # %load_ext tensorboard
-# %tensorboard --logdir output
+# %tensorboard --logdir models