Add hooks building to classification task (facebookresearch#402)

Summary:
Pull Request resolved: facebookresearch#402

Pull Request resolved: fairinternal/ClassyVision#62

Add configurable hooks to classification task, had to remove the typehints from ClassyHook to avoid a circular dependency.

Reviewed By: mannatsingh, vreis

Differential Revision: D19770583

fbshipit-source-id: 0011c3519bf0af5c5e317319e80ede70127e754c

Author: Aaron Adcock

classy_vision/hooks/classy_hook.py

@@ -7,8 +7,6 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict
 
-from classy_vision import tasks
-
 
 class ClassyHookState:
     """Class to store state within instances of ClassyHook.
@@ -69,27 +67,27 @@ def name(cls) -> str:
         return cls.__name__
 
     @abstractmethod
-    def on_start(self, task: "tasks.ClassyTask") -> None:
+    def on_start(self, task) -> None:
         """Called at the start of training."""
         pass
 
     @abstractmethod
-    def on_phase_start(self, task: "tasks.ClassyTask") -> None:
+    def on_phase_start(self, task) -> None:
         """Called at the start of each phase."""
         pass
 
     @abstractmethod
-    def on_step(self, task: "tasks.ClassyTask") -> None:
+    def on_step(self, task) -> None:
         """Called each time after parameters have been updated by the optimizer."""
         pass
 
     @abstractmethod
-    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
+    def on_phase_end(self, task) -> None:
         """Called at the end of each phase (epoch)."""
         pass
 
     @abstractmethod
-    def on_end(self, task: "tasks.ClassyTask") -> None:
+    def on_end(self, task) -> None:
         """Called at the end of training."""
         pass
 

classy_vision/tasks/classification_task.py

@@ -24,6 +24,7 @@
     recursive_copy_to_gpu,
     update_classy_state,
 )
+from classy_vision.hooks import build_hooks
 from classy_vision.losses import ClassyLoss, build_loss
 from classy_vision.meters import build_meters
 from classy_vision.models import ClassyModel, build_model
@@ -328,6 +329,10 @@ def from_config(cls, config: Dict[str, Any]) -> "ClassificationTask":
         amp_args = config.get("amp_args")
         meters = build_meters(config.get("meters", {}))
         model = build_model(config["model"])
+
+        # hooks config is optional
+        hooks_config = config.get("hooks", [])
+        hooks = build_hooks(hooks_config)
         optimizer = build_optimizer(optimizer_config)
 
         task = (
@@ -348,6 +353,7 @@ def from_config(cls, config: Dict[str, Any]) -> "ClassificationTask":
                     config.get("batch_norm_sync_mode", "disabled").upper()
                 ],
             )
+            .set_hooks(hooks)
         )
         for phase_type in phase_types:
             task.set_dataset(datasets[phase_type], phase_type)

test/generic/config_utils.py

@@ -12,6 +12,7 @@ def get_test_task_config(head_num_classes=1000):
         "name": "classification_task",
         "num_epochs": 12,
         "loss": {"name": "CrossEntropyLoss"},
+        "hooks": [{"name": "loss_lr_meter_logging"}],
         "dataset": {
             "train": {
                 "name": "synthetic_image",

test/tasks_classification_task_test.py

@@ -44,6 +44,11 @@ def test_build_task(self):
         task = build_task(config)
         self.assertTrue(isinstance(task, ClassificationTask))
 
+        # Hooks configuration is optional
+        del config["hooks"]
+        task = build_task(config)
+        self.assertTrue(isinstance(task, ClassificationTask))
+
     def test_get_state(self):
         config = get_test_task_config()
         loss = build_loss(config["loss"])