Remove local_variables from on_phase_end (#421)

Summary:
Pull Request resolved: #421

This is part of a series of diffs to eliminate local_variables (see D20171981).
Final step! Remove local_variables from on_phase_end.

Reviewed By: mannatsingh

Differential Revision: D20178293

fbshipit-source-id: 635001ba5331a1f43581c653c8f54932227c28f8

classy_vision/hooks/checkpoint_hook.py

@@ -87,9 +87,7 @@ def on_start(self, task: "tasks.ClassyTask") -> None:
             )
             raise FileNotFoundError(err_msg)
 
-    def on_phase_end(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """Checkpoint the task every checkpoint_period phases.
 
         We do not necessarily checkpoint the task at the end of every phase.

classy_vision/hooks/classy_hook.py

@@ -80,9 +80,7 @@ def on_step(self, task: "tasks.ClassyTask") -> None:
         pass
 
     @abstractmethod
-    def on_phase_end(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """Called at the end of each phase (epoch)."""
         pass
 

classy_vision/hooks/exponential_moving_average_model_hook.py

@@ -97,7 +97,7 @@ 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)
 
-    def on_phase_end(self, task: ClassyTask, local_variables: Dict[str, Any]) -> None:
+    def on_phase_end(self, task: ClassyTask) -> None:
         if task.train:
             # save the current model state since this will be overwritten by the ema
             # state in the test phase

classy_vision/hooks/loss_lr_meter_logging_hook.py

@@ -32,9 +32,7 @@ def __init__(self, log_freq: Optional[int] = None) -> None:
         super().__init__()
         self.log_freq: Optional[int] = log_freq
 
-    def on_phase_end(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """
         Log the loss, optimizer LR, and meters for the phase.
         """

classy_vision/hooks/progress_bar_hook.py

@@ -55,9 +55,7 @@ def on_step(self, task: "tasks.ClassyTask") -> None:
             self.batches += 1
             self.progress_bar.update(min(self.batches, self.bar_size))
 
-    def on_phase_end(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """Clear the progress bar at the end of the phase."""
         if is_master() and self.progress_bar is not None:
             self.progress_bar.finish()

classy_vision/hooks/tensorboard_plot_hook.py

@@ -78,9 +78,7 @@ def on_step(self, task: "tasks.ClassyTask") -> None:
         self.wall_times.append(time.time())
         self.num_steps_global.append(task.num_updates)
 
-    def on_phase_end(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """Add the losses and learning rates to tensorboard."""
         if self.learning_rates is None:
             logging.warning("learning_rates is not initialized")

classy_vision/hooks/time_metrics_hook.py

@@ -50,9 +50,7 @@ def on_step(self, task: "tasks.ClassyTask") -> None:
         if batches and batches % self.log_freq == 0:
             self._log_performance_metrics(task)
 
-    def on_phase_end(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """
         Log metrics at the end of a phase if log_freq is None.
         """

classy_vision/hooks/visdom_hook.py

@@ -60,9 +60,7 @@ def __init__(
         self.metrics: Dict = {}
         self.visdom: Visdom = Visdom(self.server, self.port)
 
-    def on_phase_end(
-        self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
-    ) -> None:
+    def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """
         Plot the metrics on visdom.
         """

classy_vision/tasks/classification_task.py

@@ -868,7 +868,7 @@ def on_phase_start(self):
 
         self.phase_start_time_train = time.perf_counter()
 
-    def on_phase_end(self, local_variables):
+    def on_phase_end(self):
         self.log_phase_end("train")
 
         logging.info("Syncing meters on phase end...")
@@ -877,7 +877,8 @@ def on_phase_end(self, local_variables):
         logging.info("...meters synced")
         barrier()
 
-        self.run_hooks(local_variables, ClassyHookFunctions.on_phase_end.name)
+        for hook in self.hooks:
+            hook.on_phase_end(self)
         self.perf_log = []
 
         self.log_phase_end("total")

classy_vision/tasks/classy_task.py

@@ -137,7 +137,7 @@ def on_phase_start(self):
         pass
 
     @abstractmethod
-    def on_phase_end(self, local_variables):
+    def on_phase_end(self):
         """
         Epoch end.
 

classy_vision/trainer/classy_trainer.py

@@ -70,8 +70,6 @@ def train(self, task: ClassyTask):
         # this helps catch hangs which would have happened elsewhere
         barrier()
 
-        local_variables = {}
-
         task.on_start()
         while not task.done_training():
             task.on_phase_start()
@@ -80,5 +78,5 @@ def train(self, task: ClassyTask):
                     task.step(self.use_gpu)
                 except StopIteration:
                     break
-            task.on_phase_end(local_variables)
+            task.on_phase_end()
         task.on_end()

classy_vision/trainer/elastic_trainer.py

@@ -70,8 +70,6 @@ def train(self, task):
         )
         state = self._ClassyElasticState(task, self.input_args)
 
-        local_variables = {}
-
         state.advance_to_next_phase = True
 
         def elastic_train_step(orig_state):
@@ -82,13 +80,13 @@ def elastic_train_step(orig_state):
                 state.run_start_hooks = False
                 return state, self._ClassyWorkerStats(None)
 
-            return self._run_step(orig_state, local_variables, self.use_gpu)
+            return self._run_step(orig_state, self.use_gpu)
 
         torchelastic.train(self.elastic_coordinator, elastic_train_step, state)
 
         task.on_end()
 
-    def _run_step(self, state, local_variables, use_gpu):
+    def _run_step(self, state, use_gpu):
         # Check for training complete but only terminate when the last phase is done
         if state.task.done_training() and state.advance_to_next_phase:
             raise StopIteration
@@ -112,10 +110,10 @@ def _run_step(self, state, local_variables, use_gpu):
 
         if state.advance_to_next_phase:
             self.elastic_coordinator.barrier()
-            state.task.on_phase_end(local_variables)
+            state.task.on_phase_end()
 
         progress_rate = None  # using None to signal 'unknown'
-        perf_stats = local_variables.get("perf_stats", None)
+        perf_stats = getattr(state.task, "perf_stats", None)
         if perf_stats is not None:
             batch_time = perf_stats._cuda_stats["train_step_total"].smoothed_value
             if batch_time is not None and batch_time > 0.0:

test/hooks_checkpoint_hook_test.py

@@ -33,7 +33,6 @@ def test_state_checkpointing(self) -> None:
         task = build_task(config)
         task.prepare()
 
-        local_variables = {}
         checkpoint_folder = self.base_dir + "/checkpoint_end_test/"
         input_args = {"foo": "bar"}
 
@@ -48,7 +47,7 @@ def test_state_checkpointing(self) -> None:
             checkpoint_hook.on_start(task)
         # call the on end phase function
         with self.assertRaises(AssertionError):
-            checkpoint_hook.on_phase_end(task, local_variables)
+            checkpoint_hook.on_phase_end(task)
         # try loading a non-existent checkpoint
         checkpoint = load_checkpoint(checkpoint_folder)
         self.assertIsNone(checkpoint)
@@ -60,13 +59,13 @@ def test_state_checkpointing(self) -> None:
         # Phase_type is test, expect no checkpoint
         task.train = False
         # call the on end phase function
-        checkpoint_hook.on_phase_end(task, local_variables)
+        checkpoint_hook.on_phase_end(task)
         checkpoint = load_checkpoint(checkpoint_folder)
         self.assertIsNone(checkpoint)
 
         task.train = True
         # call the on end phase function
-        checkpoint_hook.on_phase_end(task, local_variables)
+        checkpoint_hook.on_phase_end(task)
         # model should be checkpointed. load and compare
         checkpoint = load_checkpoint(checkpoint_folder)
         self.assertIsNotNone(checkpoint)
@@ -84,7 +83,6 @@ def test_checkpoint_period(self) -> None:
         task = build_task(config)
         task.prepare()
 
-        local_variables = {}
         checkpoint_folder = self.base_dir + "/checkpoint_end_test/"
         checkpoint_period = 10
 
@@ -110,7 +108,7 @@ def test_checkpoint_period(self) -> None:
             while valid_phase_count < checkpoint_period - 1:
                 task.train = count % 2 == 0
                 # call the on end phase function
-                checkpoint_hook.on_phase_end(task, local_variables)
+                checkpoint_hook.on_phase_end(task)
                 checkpoint = load_checkpoint(checkpoint_folder)
                 self.assertIsNone(checkpoint)
                 valid_phase_count += 1 if task.phase_type in phase_types else 0
@@ -119,7 +117,7 @@ def test_checkpoint_period(self) -> None:
             # create a phase which is in phase_types
             task.train = True
             # call the on end phase function
-            checkpoint_hook.on_phase_end(task, local_variables)
+            checkpoint_hook.on_phase_end(task)
             # model should be checkpointed. load and compare
             checkpoint = load_checkpoint(checkpoint_folder)
             self.assertIsNotNone(checkpoint)
@@ -137,13 +135,11 @@ def test_checkpointing(self):
 
         task.prepare(use_gpu=cuda_available)
 
-        local_variables = {}
-
         # create a checkpoint hook
         checkpoint_hook = CheckpointHook(checkpoint_folder, {}, phase_types=["train"])
 
         # call the on end phase function
-        checkpoint_hook.on_phase_end(task, local_variables)
+        checkpoint_hook.on_phase_end(task)
 
         # we should be able to train a task using the checkpoint on all available
         # devices

test/hooks_exponential_moving_average_model_hook_test.py

@@ -37,7 +37,6 @@ def _map_device_string(self, device):
     def _test_exponential_moving_average_hook(self, model_device, hook_device):
         task = mock.MagicMock()
         model = TestModel().to(device=self._map_device_string(model_device))
-        local_variables = {}
         task.base_model = model
         task.train = True
         decay = 0.5
@@ -54,14 +53,14 @@ def _test_exponential_moving_average_hook(self, model_device, hook_device):
         fc_weight = model.fc.weight.clone()
         for _ in range(num_updates):
             exponential_moving_average_hook.on_step(task)
-        exponential_moving_average_hook.on_phase_end(task, local_variables)
+        exponential_moving_average_hook.on_phase_end(task)
         # the model weights shouldn't have changed
         self.assertTrue(torch.allclose(model.fc.weight, fc_weight))
 
         # simulate a test phase now
         task.train = False
         exponential_moving_average_hook.on_phase_start(task)
-        exponential_moving_average_hook.on_phase_end(task, local_variables)
+        exponential_moving_average_hook.on_phase_end(task)
 
         # the model weights should be updated to the ema weights
         self.assertTrue(

test/hooks_loss_lr_meter_logging_hook_test.py

@@ -34,7 +34,6 @@ def test_logging(self, mock_get_rank: mock.MagicMock) -> None:
 
         losses = [1.2, 2.3, 3.4, 4.5]
 
-        local_variables = {}
         task.phase_idx = 0
 
         for log_freq in [5, None]:
@@ -61,7 +60,7 @@ def test_logging(self, mock_get_rank: mock.MagicMock) -> None:
                         mock_fn.assert_not_called()
                         mock_lr_fn.assert_not_called()
 
-                    loss_lr_meter_hook.on_phase_end(task, local_variables)
+                    loss_lr_meter_hook.on_phase_end(task)
                     mock_fn.assert_called_with(task)
                     if task.train:
                         mock_lr_fn.assert_called_with(task)

test/hooks_time_metrics_hook_test.py

@@ -32,7 +32,6 @@ def test_time_metrics(
         mock_get_rank.return_value = rank
 
         mock_report_str.return_value = ""
-        local_variables = {}
 
         for log_freq, train in product([5, None], [True, False]):
             # create a time metrics hook
@@ -53,7 +52,7 @@ def test_time_metrics(
 
             # test that the code doesn't raise an exception if losses is empty
             try:
-                time_metrics_hook.on_phase_end(task, local_variables)
+                time_metrics_hook.on_phase_end(task)
             except Exception as e:
                 self.fail("Received Exception when losses is []: {}".format(e))
 
@@ -73,7 +72,7 @@ def test_time_metrics(
                         continue
                     mock_fn.assert_not_called()
 
-                time_metrics_hook.on_phase_end(task, local_variables)
+                time_metrics_hook.on_phase_end(task)
                 mock_fn.assert_called_with(task)
 
             task.losses = [0.23, 0.45, 0.34, 0.67]
@@ -116,7 +115,7 @@ def test_time_metrics(
             time_metrics_hook_new = TimeMetricsHook()
 
             with self.assertLogs() as log_watcher:
-                time_metrics_hook_new.on_phase_end(task, local_variables)
+                time_metrics_hook_new.on_phase_end(task)
 
             self.assertEqual(len(log_watcher.output), 2)
             self.assertTrue(

test/manual/hooks_progress_bar_hook_test.py

@@ -26,8 +26,6 @@ def test_progress_bar(
 
         mock_is_master.return_value = True
 
-        local_variables = {}
-
         task = get_test_classy_task()
         task.prepare()
         task.advance_phase()
@@ -60,7 +58,7 @@ def test_progress_bar(
             mock_progress_bar.update.reset_mock()
 
         # finish should be called on the progress bar
-        progress_bar_hook.on_phase_end(task, local_variables)
+        progress_bar_hook.on_phase_end(task)
         mock_progress_bar.finish.assert_called_once_with()
         mock_progress_bar.finish.reset_mock()
 
@@ -69,7 +67,7 @@ def test_progress_bar(
         progress_bar_hook = ProgressBarHook()
         try:
             progress_bar_hook.on_step(task)
-            progress_bar_hook.on_phase_end(task, local_variables)
+            progress_bar_hook.on_phase_end(task)
         except Exception as e:
             self.fail(
                 "Received Exception when on_phase_start() isn't called: {}".format(e)
@@ -82,7 +80,7 @@ def test_progress_bar(
         try:
             progress_bar_hook.on_phase_start(task)
             progress_bar_hook.on_step(task)
-            progress_bar_hook.on_phase_end(task, local_variables)
+            progress_bar_hook.on_phase_end(task)
         except Exception as e:
             self.fail("Received Exception when is_master() is False: {}".format(e))
         self.assertIsNone(progress_bar_hook.progress_bar)

test/manual/hooks_tensorboard_plot_hook_test.py

@@ -49,8 +49,6 @@ def test_writer(self, mock_is_master_func: mock.MagicMock) -> None:
 
             losses = [1.23, 4.45, 12.3, 3.4]
 
-            local_variables = {}
-
             summary_writer = SummaryWriter(self.base_dir)
             # create a spy on top of summary_writer
             summary_writer = mock.MagicMock(wraps=summary_writer)
@@ -74,7 +72,7 @@ def test_writer(self, mock_is_master_func: mock.MagicMock) -> None:
             # the writer if on_phase_start() is not called for initialization
             # if on_phase_end() is called.
             with self.assertLogs() as log_watcher:
-                tensorboard_plot_hook.on_phase_end(task, local_variables)
+                tensorboard_plot_hook.on_phase_end(task)
 
             self.assertTrue(
                 len(log_watcher.records) == 1
@@ -90,7 +88,7 @@ def test_writer(self, mock_is_master_func: mock.MagicMock) -> None:
                 task.losses.append(loss)
                 tensorboard_plot_hook.on_step(task)
 
-            tensorboard_plot_hook.on_phase_end(task, local_variables)
+            tensorboard_plot_hook.on_phase_end(task)
 
             if master:
                 # add_scalar() should have been called with the right scalars