add callbacks for inference_on_dataset

Differential Revision: D51540498

d2go/runner/default_runner.py

@@ -386,6 +386,16 @@ def _create_evaluators(
                 )
         return evaluator
 
+    # experimental API
+    @classmethod
+    def _get_inference_callbacks(cls):
+        return {
+            "on_start": lambda: None,
+            "on_end": lambda: None,
+            "before_inference": lambda: None,
+            "after_inference": lambda: None,
+        }
+
     def _do_test(self, cfg, model, train_iter=None, model_tag="default"):
         """train_iter: Current iteration of the model, None means final iteration"""
         assert len(cfg.DATASETS.TEST)
@@ -430,7 +440,10 @@ def _get_inference_dir_name(base_dir, inference_type, dataset_name):
                 else model,
             )
 
-            results_per_dataset = inference_on_dataset(model, data_loader, evaluator)
+            inference_callbacks = self._get_inference_callbacks()
+            results_per_dataset = inference_on_dataset(
+                model, data_loader, evaluator, callbacks=inference_callbacks
+            )
 
             if comm.is_main_process():
                 results[model_tag][dataset_name] = results_per_dataset