Remove local_variables from on_start (#413)

vreis · facebook-github-bot · commit fd79c1c113b3 · 2020-03-02T19:28:36.000-08:00
Summary: Pull Request resolved: #413 This is part of a series of diffs to eliminate local_variables (see D20171981). Proceed removing local_variables from on_start Reviewed By: mannatsingh Differential Revision: D20178165 fbshipit-source-id: 069a05dd418668d3c5122e9e13177efec05b2125
diff --git a/classy_vision/hooks/checkpoint_hook.py b/classy_vision/hooks/checkpoint_hook.py
@@ -78,9 +78,7 @@ def _save_checkpoint(self, task, filename):
         if checkpoint_file:
             PathManager.copy(checkpoint_file, f"{self.checkpoint_folder}/{filename}")
 
-    def on_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_start(self, task: "tasks.ClassyTask") -> None:
         if not is_master() or getattr(task, "test_only", False):
             return
         if not PathManager.exists(self.checkpoint_folder):
diff --git a/classy_vision/hooks/classy_hook.py b/classy_vision/hooks/classy_hook.py
@@ -65,9 +65,7 @@ def name(cls) -> str:
         return cls.__name__
 
     @abstractmethod
-    def on_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_start(self, task: "tasks.ClassyTask") -> None:
         """Called at the start of training."""
         pass
 
diff --git a/classy_vision/hooks/exponential_moving_average_model_hook.py b/classy_vision/hooks/exponential_moving_average_model_hook.py
@@ -78,7 +78,7 @@ def _save_current_model_state(self, model: nn.Module, model_state: Dict[str, Any
         for name, param in self.get_model_state_iterator(model):
             model_state[name] = param.detach().clone().to(device=self.device)
 
-    def on_start(self, task: ClassyTask, local_variables: Dict[str, Any]) -> None:
+    def on_start(self, task: ClassyTask) -> None:
         if self.state.model_state:
             # loaded state from checkpoint, do not re-initialize, only move the state
             # to the right device
diff --git a/classy_vision/hooks/model_complexity_hook.py b/classy_vision/hooks/model_complexity_hook.py
@@ -28,9 +28,7 @@ class ModelComplexityHook(ClassyHook):
     on_phase_end = ClassyHook._noop
     on_end = ClassyHook._noop
 
-    def on_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_start(self, task: "tasks.ClassyTask") -> None:
         """Measure number of parameters, FLOPs and activations."""
         self.num_flops = 0
         self.num_activations = 0
diff --git a/classy_vision/hooks/model_tensorboard_hook.py b/classy_vision/hooks/model_tensorboard_hook.py
@@ -51,9 +51,7 @@ def __init__(self, tb_writer) -> None:
 
         self.tb_writer = tb_writer
 
-    def on_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_start(self, task: "tasks.ClassyTask") -> None:
         """
         Plot the model on Tensorboard.
         """
diff --git a/classy_vision/hooks/profiler_hook.py b/classy_vision/hooks/profiler_hook.py
@@ -25,9 +25,7 @@ class ProfilerHook(ClassyHook):
     on_phase_end = ClassyHook._noop
     on_end = ClassyHook._noop
 
-    def on_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_start(self, task: "tasks.ClassyTask") -> None:
         """Profile the forward pass."""
         logging.info("Profiling forward pass...")
         batchsize_per_replica = getattr(
diff --git a/classy_vision/tasks/classification_task.py b/classy_vision/tasks/classification_task.py
@@ -854,8 +854,9 @@ def get_global_batchsize(self):
         """
         return self.dataloaders[self.phase_type].dataset.get_global_batchsize()
 
-    def on_start(self, local_variables):
-        self.run_hooks(local_variables, ClassyHookFunctions.on_start.name)
+    def on_start(self):
+        for hook in self.hooks:
+            hook.on_start(self)
 
     def on_phase_start(self, local_variables):
         self.phase_start_time_total = time.perf_counter()
diff --git a/classy_vision/tasks/classy_task.py b/classy_vision/tasks/classy_task.py
@@ -119,7 +119,7 @@ def train_step(self, use_gpu) -> None:
         pass
 
     @abstractmethod
-    def on_start(self, local_variables):
+    def on_start(self):
         """
         Start training.
 
diff --git a/classy_vision/trainer/classy_trainer.py b/classy_vision/trainer/classy_trainer.py
@@ -72,7 +72,7 @@ def train(self, task: ClassyTask):
 
         local_variables = {}
 
-        task.on_start(local_variables)
+        task.on_start()
         while not task.done_training():
             task.on_phase_start(local_variables)
             while True:
diff --git a/classy_vision/trainer/elastic_trainer.py b/classy_vision/trainer/elastic_trainer.py
@@ -78,7 +78,7 @@ def elastic_train_step(orig_state):
             if state.run_start_hooks:
                 # need this to ensure we don't run the on_start hooks every time
                 # a trainer starts
-                state.task.on_start(local_variables)
+                state.task.on_start()
                 state.run_start_hooks = False
                 return state, self._ClassyWorkerStats(None)
 
diff --git a/test/hooks_checkpoint_hook_test.py b/test/hooks_checkpoint_hook_test.py
@@ -45,7 +45,7 @@ def test_state_checkpointing(self) -> None:
         # checkpoint directory doesn't exist
         # call the on start function
         with self.assertRaises(FileNotFoundError):
-            checkpoint_hook.on_start(task, local_variables)
+            checkpoint_hook.on_start(task)
         # call the on end phase function
         with self.assertRaises(AssertionError):
             checkpoint_hook.on_phase_end(task, local_variables)
@@ -55,7 +55,7 @@ def test_state_checkpointing(self) -> None:
 
         # create checkpoint dir, verify on_start hook runs
         os.mkdir(checkpoint_folder)
-        checkpoint_hook.on_start(task, local_variables)
+        checkpoint_hook.on_start(task)
 
         # Phase_type is test, expect no checkpoint
         task.train = False
@@ -101,7 +101,7 @@ def test_checkpoint_period(self) -> None:
             os.mkdir(checkpoint_folder)
 
             # call the on start function
-            checkpoint_hook.on_start(task, local_variables)
+            checkpoint_hook.on_start(task)
 
             # shouldn't create any checkpoints until there are checkpoint_period
             # phases which are in phase_types
diff --git a/test/hooks_exponential_moving_average_model_hook_test.py b/test/hooks_exponential_moving_average_model_hook_test.py
@@ -47,7 +47,7 @@ def _test_exponential_moving_average_hook(self, model_device, hook_device):
             decay=decay, device=hook_device
         )
 
-        exponential_moving_average_hook.on_start(task, local_variables)
+        exponential_moving_average_hook.on_start(task)
         exponential_moving_average_hook.on_phase_start(task, local_variables)
         # set the weights to all ones and simulate 10 updates
         task.base_model.update_fc_weight()
diff --git a/test/hooks_profiler_hook_test.py b/test/hooks_profiler_hook_test.py
@@ -32,13 +32,12 @@ def test_profiler(
 
         for task in [get_test_classy_task(), get_test_classy_video_task()]:
             task.prepare()
-            local_variables = {}
 
             # create a model tensorboard hook
             profiler_hook = ProfilerHook()
 
             with self.assertLogs():
-                profiler_hook.on_start(task, local_variables)
+                profiler_hook.on_start(task)
 
             # a new profile should be created with use_cuda=True
             mock_profile_cls.assert_called_once_with(use_cuda=True)
diff --git a/test/manual/hooks_model_complexity_hook_test.py b/test/manual/hooks_model_complexity_hook_test.py
@@ -18,8 +18,6 @@ def test_model_complexity_hook(self) -> None:
         task = get_test_classy_task()
         task.prepare()
 
-        local_variables = {}
-
         # create a model complexity hook
         model_complexity_hook = ModelComplexityHook()
 
@@ -29,4 +27,4 @@ def test_model_complexity_hook(self) -> None:
             task.base_model = model
 
             with self.assertLogs():
-                model_complexity_hook.on_start(task, local_variables)
+                model_complexity_hook.on_start(task)
diff --git a/test/manual/hooks_model_tensorboard_hook_test.py b/test/manual/hooks_model_tensorboard_hook_test.py
@@ -28,7 +28,6 @@ def test_writer(self, mock_is_master_func: mock.MagicMock) -> None:
         for master in [False, True]:
             mock_is_master_func.return_value = master
             model_configs = get_test_model_configs()
-            local_variables = {}
 
             for model_config in model_configs:
                 model = build_model(model_config)
@@ -37,7 +36,7 @@ def test_writer(self, mock_is_master_func: mock.MagicMock) -> None:
                 # create a model tensorboard hook
                 model_tensorboard_hook = ModelTensorboardHook(mock_summary_writer)
 
-                model_tensorboard_hook.on_start(task, local_variables)
+                model_tensorboard_hook.on_start(task)
 
                 if master:
                     # SummaryWriter should have been init-ed with the correct
diff --git a/test/models_classy_model_test.py b/test/models_classy_model_test.py
@@ -59,7 +59,7 @@ def test_from_checkpoint(self):
 
         # Create checkpoint dir, save checkpoint
         os.mkdir(checkpoint_folder)
-        checkpoint_hook.on_start(task, local_variables)
+        checkpoint_hook.on_start(task)
 
         task.train = True
         checkpoint_hook.on_phase_end(task, local_variables)
