Validation loss smaller than training loss - Detectron2

[misabellerv]

Hello! I am running a project using Detectron2 (object detection with faster_rcnn_R_50_FPN_1x), where I have training and validation data. When I check the metrics JSON file, the training loss is always higher than the validation loss! Validation also converges very fast while training loss is still very high.

I thought maybe my data was badly distributed, but now I suspect that it's something about my code in LossEvalHook.py and train.py.

There is this thread on GitHub about printing validation loss, and I tried to use the codes from there. Although it prints validation loss at each iteration, it's always smaller than the training loss!

Could someone help me, please? =)

Here is LossEvalHook.py:

import os
import torch
import detectron2.utils.comm as comm
from detectron2.data import build_detection_train_loader
from detectron2.engine import DefaultTrainer
from detectron2.engine.hooks import HookBase
from detectron2.evaluation import COCOEvaluator
import nni
from detectron2.engine.hooks import HookBase, PeriodicWriter

class ValidationLoss(HookBase):
    def __init__(self, cfg, model, checkpointer, min_delta=0.001, patience=2000) -> None:
        super().__init__()
        self.cfg = cfg.clone()
        self.cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS = False
        self.cfg.DATASETS.TRAIN = self.cfg.DATASETS.VAL
        self._loader = iter(build_detection_train_loader(self.cfg))
        self.model = model
        self._checkpointer = checkpointer
        self.best_loss = float('inf')
        self.min_delta = min_delta
        self.patience = patience
        self.counter = 0
        self.cfg.OUTPUT_DIR = cfg.OUTPUT_DIR

    def after_step(self):
        data = next(self._loader)
        with torch.no_grad():
            loss_dict = self.model(data)

            losses = sum(loss_dict.values())
            assert torch.isfinite(losses).all(), loss_dict

            loss_dict_reduced = {
                "val_" + k: v.item()
                for k, v in comm.reduce_dict(loss_dict).items()
            }
            losses_reduced = sum(loss for loss in loss_dict_reduced.values())
            if comm.is_main_process():
                self.trainer.storage.put_scalars(total_val_loss=losses_reduced, **loss_dict_reduced)

            nni.report_intermediate_result(losses_reduced)

            # Check if the current model has the best validation loss so far
            if losses_reduced < self.best_loss - self.min_delta:
                self.best_loss = losses_reduced
                self.counter = 0
            else:
                self.counter += 1
                print(f'COUNTER: {self.counter}')

                if self.counter >= self.patience:
                    print(f"Early stopping at iteration {self.trainer.iter} due to no improvement in validation loss.")
                    self._checkpointer.save("model_early_stop")
                    self.trainer.storage.put_scalar('early_stop', self.trainer.iter)
                    raise StopIteration

    def after_train(self):
        nni.report_final_result(self.trainer.storage.latest()["total_val_loss"][0])

class MyTrainer(DefaultTrainer):

    def build_hooks(self):
        hooks = super().build_hooks()

        if self.cfg.DATASETS.VAL[0] is not None:
            hooks.insert(-1, ValidationLoss(self.cfg, self.model, self.checkpointer))
            hooks[-1] = PeriodicWriter(self.build_writers(), period=self.cfg.VAL_PERIOD)
        return hooks

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):

        if output_folder is None:
            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
        coco_evaluator = COCOEvaluator(dataset_name, cfg, True, output_folder)

        return [coco_evaluator]

    @classmethod
    def build_train_loader(cls, cfg):

        print("USING DEFAULT AUGMENTATIONS!")
        mapper = None
        return build_detection_train_loader(cfg, mapper=mapper)

Here is my train.py:

import os
from detectron2.config import get_cfg
from LossEvalHookEarlyStop import MyTrainer2 as MyTrainer
#from LossEvalHookTest import MyTrainer2 as MyTrainer
#from LossEvalHookTest import ValidationLoss
from detectron2.utils.logger import setup_logger
import argparse
from register_dataset import register_mosquitoes
import logging
import warnings
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.engine.hooks import HookBase, PeriodicWriter
warnings.filterwarnings("ignore")

# ARGPARSE CUSTOMIZE COMMAND LINE
logger = logging.getLogger("lm")
parser = argparse.ArgumentParser(description="MBG Training")
parser.add_argument("--config-file", default=None, metavar="FILE", help="path to config file")
parser.add_argument("--data-train", default="mbg_train0_tire", metavar="FILE", help="path to data")
parser.add_argument("--data-test", default="mbg_val0_tire", metavar="FILE", help="path to data")
parser.add_argument("--data-val", default="mbg_val0_tire", metavar="FILE", help="path to data")
parser.add_argument("--iters", default=None, help="Max number of iterations")
args = parser.parse_args()

# EXTRACT FOLD INDEX
fold_index = args.data_train.split('_')[1][-1]
obj_index = args.data_train.split('_')[-1]
fold_index = f"fold{fold_index}_{obj_index}"

# GET CONFIG FILE AND ADDITIONAL ARGS
cfg = get_cfg()
config_file = args.config_file
cfg.merge_from_file(config_file)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

cfg.DATASETS.TRAIN = (args.data_train,)
cfg.DATASETS.TEST = (args.data_val,)
cfg.DATASETS.VAL = (args.data_val,)
cfg.TEST.EVAL_PERIOD = 0
cfg.VAL_PERIOD = 1
cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 32
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.min_delta = 0.019
cfg.patience = 350
#cfg.min_delta = 0.1
#cfg.patience = 5

logger.info(args)

cfg.SOLVER.MAX_ITER = int(args.iters)
cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, fold_index)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
setup_logger(cfg.OUTPUT_DIR)

# REGISTER DATASET
register_mosquitoes()

# SAVE TRAINING CONFIGURATIONS
with open(os.path.join(cfg.OUTPUT_DIR, "training_config.yaml"), "w") as tc:
    cfg.dump(stream=tc)

# TRAINING
trainer = MyTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

You can see here the training/validation loss curve:

Thanks!