detectron2_earlystopping.py

logger.info("Starting training from iteration {}".format(start_iter))
    with EventStorage(start_iter) as storage:
        previous_val_map = 0  # ADDED NEW
        no_improvement_for_epochs = 0  # ADDED NEW
        for data, iteration in zip(data_loader, range(start_iter, max_iter)):
            storage.iter = iteration
            ...
            if (
                    cfg.TEST.EVAL_PERIOD > 0
                    and (iteration + 1) % cfg.TEST.EVAL_PERIOD == 0
                    and iteration != max_iter - 1
            ):
                results = do_test(cfg, model) # ADDED NEW: TAKES RETURN VALUES 
                # results = OrderedDict([('segm', {'AP': 0.0, 'AP50': 0.0})])

                ##### ADDED NEW ####################################################################
                logger.info(f"Previous allAP50 is: {previous_val_map}")
                if results['segm']['AP50'] > previous_val_map:
                    previous_val_map = results['segm']['AP50']
                    no_improvement_for_epochs = 0

                    # Save best model
                    checkpointer.save("best_model")
                else:
                    no_improvement_for_epochs += 1
                    logger.info(f"There is no improvement for {no_improvement_for_epochs} epochs")
                #####################################################################################

                # Compared to "train_net.py", the test results are not dumped to EventStorage
                comm.synchronize()

            if iteration - start_iter > 5 and (
                    (iteration + 1) % 20 == 0 or iteration == max_iter - 1
            ):
                for writer in writers:
                    writer.write()
            periodic_checkpointer.step(iteration)

            ##### ADDED NEW ####################################################################
            # STOP IF THERE IS NO IMPROVEMENT FOR 5 EPOCHS
            if no_improvement_for_epochs >= EARLY_STOPPING_ITERATION:
                logger.info(f"Quiting from the training...")
                try:
                    for writer in writers:
                        writer.write()
                except Exception as e:
                    logger.info(f"before termination latest logs cannot be sent to tensorboard:{e}")
                break
            #####################################################################################