add callbacks for inference_on_dataset

Summary: Pull Request resolved: #637

Reviewed By: tglik

Differential Revision: D51540498

d2go/evaluation/evaluator.py

@@ -36,22 +36,23 @@ def inference_on_dataset(
     model: torch.nn.Module,
     data_loader: Iterable,
     evaluator: Union[DatasetEvaluator, List[DatasetEvaluator], None],
+    **kwargs,
 ):
     """
     A drop-in replacement for d2's inference_on_dataset to run inference on datasets,
     supports customization for checkpointing
     * has_finished_process(self) -> bool: return True if `self.process()` could be skipped
     """
     if evaluator is None:
-        return inference_on_dataset_d2(model, data_loader, evaluator)
+        return inference_on_dataset_d2(model, data_loader, evaluator, **kwargs)
 
     if isinstance(evaluator, abc.MutableSequence):
         evaluator = DatasetEvaluators(evaluator)
 
     if not (
         hasattr(evaluator, "has_finished_process") and evaluator.has_finished_process()
     ):
-        return inference_on_dataset_d2(model, data_loader, evaluator)
+        return inference_on_dataset_d2(model, data_loader, evaluator, **kwargs)
 
     evaluator.reset()
     results = evaluator.evaluate()

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