Remove local_variables from on_phase_start (#416)

Summary:
Pull Request resolved: #416

This is part of a series of diffs to eliminate local_variables (see D20171981).
Proceed removing local_variables from on_phase_start

Reviewed By: mannatsingh

Differential Revision: D20178268

fbshipit-source-id: 09f78810228b2fec9faa2205d92b108aea30aff9

Author: vreis

classy_vision/hooks/classy_hook.py

@@ -70,9 +70,7 @@ def on_start(self, task: "tasks.ClassyTask") -> None:
         pass
 
     @abstractmethod
-    def on_phase_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_start(self, task: "tasks.ClassyTask") -> None:
         """Called at the start of each phase."""
         pass
 

classy_vision/hooks/exponential_moving_average_model_hook.py

@@ -93,7 +93,7 @@ def on_start(self, task: ClassyTask) -> None:
         self._save_current_model_state(task.base_model, self.state.model_state)
         self._save_current_model_state(task.base_model, self.state.ema_model_state)
 
-    def on_phase_start(self, task: ClassyTask, local_variables: Dict[str, Any]) -> None:
+    def on_phase_start(self, task: ClassyTask) -> None:
         # restore the right state depending on the phase type
         self.set_model_state(task, use_ema=not task.train)
 

classy_vision/hooks/progress_bar_hook.py

@@ -36,9 +36,7 @@ def __init__(self) -> None:
         self.bar_size: int = 0
         self.batches: int = 0
 
-    def on_phase_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_start(self, task: "tasks.ClassyTask") -> None:
         """Create and display a progress bar with 0 progress."""
         if not progressbar_available:
             raise RuntimeError(

classy_vision/hooks/tensorboard_plot_hook.py

@@ -56,9 +56,7 @@ def __init__(self, tb_writer) -> None:
         self.wall_times: Optional[List[float]] = None
         self.num_steps_global: Optional[List[int]] = None
 
-    def on_phase_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_start(self, task: "tasks.ClassyTask") -> None:
         """Initialize losses and learning_rates."""
         self.learning_rates = []
         self.wall_times = []

classy_vision/hooks/time_metrics_hook.py

@@ -33,9 +33,7 @@ def __init__(self, log_freq: Optional[int] = None) -> None:
         self.log_freq: Optional[int] = log_freq
         self.start_time: Optional[float] = None
 
-    def on_phase_start(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_start(self, task: "tasks.ClassyTask") -> None:
         """
         Initialize start time and reset perf stats
         """

classy_vision/tasks/classification_task.py

@@ -858,12 +858,13 @@ def on_start(self):
         for hook in self.hooks:
             hook.on_start(self)
 
-    def on_phase_start(self, local_variables):
+    def on_phase_start(self):
         self.phase_start_time_total = time.perf_counter()
 
         self.advance_phase()
 
-        self.run_hooks(local_variables, ClassyHookFunctions.on_phase_start.name)
+        for hook in self.hooks:
+            hook.on_phase_start(self)
 
         self.phase_start_time_train = time.perf_counter()
 

classy_vision/tasks/classy_task.py

@@ -128,7 +128,7 @@ def on_start(self):
         pass
 
     @abstractmethod
-    def on_phase_start(self, local_variables):
+    def on_phase_start(self):
         """
         Epoch start.
 

classy_vision/trainer/classy_trainer.py

@@ -74,7 +74,7 @@ def train(self, task: ClassyTask):
 
         task.on_start()
         while not task.done_training():
-            task.on_phase_start(local_variables)
+            task.on_phase_start()
             while True:
                 try:
                     task.step(self.use_gpu)

classy_vision/trainer/elastic_trainer.py

@@ -96,7 +96,7 @@ def _run_step(self, state, local_variables, use_gpu):
         if state.advance_to_next_phase:
             self.elastic_coordinator.barrier()
             self.elastic_coordinator._log_event("on_phase_start")
-            state.task.on_phase_start(local_variables)
+            state.task.on_phase_start()
 
             state.advance_to_next_phase = False
 

test/hooks_exponential_moving_average_model_hook_test.py

@@ -48,7 +48,7 @@ def _test_exponential_moving_average_hook(self, model_device, hook_device):
         )
 
         exponential_moving_average_hook.on_start(task)
-        exponential_moving_average_hook.on_phase_start(task, local_variables)
+        exponential_moving_average_hook.on_phase_start(task)
         # set the weights to all ones and simulate 10 updates
         task.base_model.update_fc_weight()
         fc_weight = model.fc.weight.clone()
@@ -60,7 +60,7 @@ def _test_exponential_moving_average_hook(self, model_device, hook_device):
 
         # simulate a test phase now
         task.train = False
-        exponential_moving_average_hook.on_phase_start(task, local_variables)
+        exponential_moving_average_hook.on_phase_start(task)
         exponential_moving_average_hook.on_phase_end(task, local_variables)
 
         # the model weights should be updated to the ema weights
@@ -72,7 +72,7 @@ def _test_exponential_moving_average_hook(self, model_device, hook_device):
 
         # simulate a train phase again
         task.train = True
-        exponential_moving_average_hook.on_phase_start(task, local_variables)
+        exponential_moving_average_hook.on_phase_start(task)
 
         # the model weights should be back to the old value
         self.assertTrue(torch.allclose(model.fc.weight, fc_weight))

test/hooks_time_metrics_hook_test.py

@@ -47,7 +47,7 @@ def test_time_metrics(
             # on_phase_start() should set the start time and perf_stats
             start_time = 1.2
             mock_time.return_value = start_time
-            time_metrics_hook.on_phase_start(task, local_variables)
+            time_metrics_hook.on_phase_start(task)
             self.assertEqual(time_metrics_hook.start_time, start_time)
             self.assertTrue(isinstance(task.perf_stats, PerfStats))
 

test/manual/hooks_progress_bar_hook_test.py

@@ -40,7 +40,7 @@ def test_progress_bar(
         progress_bar_hook = ProgressBarHook()
 
         # progressbar.ProgressBar should be init-ed with num_batches
-        progress_bar_hook.on_phase_start(task, local_variables)
+        progress_bar_hook.on_phase_start(task)
         mock_progressbar_pkg.ProgressBar.assert_called_once_with(num_batches)
         mock_progress_bar.start.assert_called_once_with()
         mock_progress_bar.start.reset_mock()
@@ -80,7 +80,7 @@ def test_progress_bar(
         mock_is_master.return_value = False
         progress_bar_hook = ProgressBarHook()
         try:
-            progress_bar_hook.on_phase_start(task, local_variables)
+            progress_bar_hook.on_phase_start(task)
             progress_bar_hook.on_step(task)
             progress_bar_hook.on_phase_end(task, local_variables)
         except Exception as e:

test/manual/hooks_tensorboard_plot_hook_test.py

@@ -84,7 +84,7 @@ def test_writer(self, mock_is_master_func: mock.MagicMock) -> None:
             summary_writer.add_scalar.reset_mock()
 
             # run the hook in the correct order
-            tensorboard_plot_hook.on_phase_start(task, local_variables)
+            tensorboard_plot_hook.on_phase_start(task)
 
             for loss in losses:
                 task.losses.append(loss)