Enhance Deepedit bundle to enable finetune and early stopping

models/spleen_deepedit_annotation/configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json",
-    "version": "0.4.9",
+    "version": "0.5.0",
     "changelog": {
+        "0.5.0": "enable finetune and early stop",
         "0.4.9": "fix orientation issue on clicks",
         "0.4.8": "Add infer transforms to manage clicks from viewer",
         "0.4.7": "fix the wrong GPU index issue of multi-node",

models/spleen_deepedit_annotation/configs/multi_gpu_train.json

@@ -23,14 +23,45 @@
         "shuffle": false
     },
     "validate#dataloader#sampler": "@validate#sampler",
-    "validate#evaluator#val_handlers": "$None if dist.get_rank() > 0 else @validate#handlers",
+    "validate#handlers": [
+        {
+            "_target_": "EarlyStopHandler",
+            "_disabled_": "$not @early_stop",
+            "trainer": null,
+            "patience": 1,
+            "score_function": "$scripts.score_function",
+            "min_delta": 0.01
+        },
+        {
+            "_target_": "StatsHandler",
+            "_disabled_": "$dist.get_rank() > 0",
+            "iteration_log": false
+        },
+        {
+            "_target_": "TensorBoardStatsHandler",
+            "_disabled_": "$dist.get_rank() > 0",
+            "log_dir": "@output_dir",
+            "iteration_log": false
+        },
+        {
+            "_target_": "CheckpointSaver",
+            "_disabled_": "$dist.get_rank() > 0",
+            "save_dir": "@ckpt_dir",
+            "save_dict": {
+                "model": "@network"
+            },
+            "save_key_metric": true,
+            "key_metric_filename": "model.pt"
+        }
+    ],
     "initialize": [
         "$import torch.distributed as dist",
         "$dist.is_initialized() or dist.init_process_group(backend='nccl')",
         "$torch.cuda.set_device(@device)",
         "$monai.utils.set_determinism(seed=123)"
     ],
     "run": [
+        "$@validate#handlers#0.set_trainer(trainer=@train#trainer) if @early_stop else None",
         "$@train#trainer.run()"
     ],
     "finalize": [

models/spleen_deepedit_annotation/configs/train.json

@@ -2,7 +2,8 @@
     "imports": [
         "$import glob",
         "$import os",
-        "$import ignite"
+        "$import ignite",
+        "$import scripts"
     ],
     "bundle_root": ".",
     "ckpt_dir": "$@bundle_root + '/models'",
@@ -14,6 +15,10 @@
         "spleen": 1,
         "background": 0
     },
+    "finetune": false,
+    "finetune_model_path": "$@bundle_root + '/models/model.pt'",
+    "early_stop": false,
+    "epochs": 500,
     "spatial_size": [
         128,
         128,
@@ -302,6 +307,14 @@
             ]
         },
         "handlers": [
+            {
+                "_target_": "CheckpointLoader",
+                "_disabled_": "$not @finetune",
+                "load_path": "@finetune_model_path",
+                "load_dict": {
+                    "model": "@network"
+                }
+            },
             {
                 "_target_": "LrScheduleHandler",
                 "lr_scheduler": "@lr_scheduler",
@@ -342,7 +355,7 @@
         "trainer": {
             "_target_": "SupervisedTrainer",
             "device": "@device",
-            "max_epochs": 500,
+            "max_epochs": "@epochs",
             "train_data_loader": "@train#dataloader",
             "network": "@network",
             "optimizer": "@optimizer",
@@ -379,6 +392,14 @@
         },
         "postprocessing": "%train#postprocessing",
         "handlers": [
+            {
+                "_target_": "EarlyStopHandler",
+                "_disabled_": "$not @early_stop",
+                "trainer": null,
+                "patience": 1,
+                "score_function": "$scripts.score_function",
+                "min_delta": 0.01
+            },
             {
                 "_target_": "StatsHandler",
                 "iteration_log": false
@@ -429,6 +450,7 @@
         "$monai.utils.set_determinism(seed=123)"
     ],
     "run": [
+        "$@validate#handlers#0.set_trainer(trainer=@train#trainer) if @early_stop else None",
         "$@train#trainer.run()"
     ]
 }

models/spleen_deepedit_annotation/scripts/__init__.py

@@ -0,0 +1 @@
+from .early_stop_score_function import score_function

models/spleen_deepedit_annotation/scripts/early_stop_score_function.py

@@ -0,0 +1,15 @@
+import os
+
+import torch
+import torch.distributed as dist
+
+
+def score_function(engine):
+    val_metric = engine.state.metrics["val_mean_dice"]
+    if dist.is_initialized():
+        device = torch.device("cuda:" + os.environ["LOCAL_RANK"])
+        val_metric = torch.tensor([val_metric]).to(device)
+        dist.all_reduce(val_metric, op=dist.ReduceOp.SUM)
+        val_metric /= dist.get_world_size()
+        return val_metric.item()
+    return val_metric