From 0b2368efce6597dd71f9b748e359dbfd1d05b745 Mon Sep 17 00:00:00 2001 From: Vinicius Reis Date: Thu, 5 Mar 2020 18:41:33 -0800 Subject: [PATCH] Remove local_variables from on_phase_start (#416) Summary: Pull Request resolved: https://github.com/facebookresearch/ClassyVision/pull/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 --- classy_vision/hooks/classy_hook.py | 4 +--- .../hooks/exponential_moving_average_model_hook.py | 2 +- classy_vision/hooks/progress_bar_hook.py | 4 +--- classy_vision/hooks/tensorboard_plot_hook.py | 4 +--- classy_vision/hooks/time_metrics_hook.py | 4 +--- classy_vision/tasks/classification_task.py | 5 +++-- classy_vision/tasks/classy_task.py | 2 +- classy_vision/trainer/classy_trainer.py | 2 +- classy_vision/trainer/elastic_trainer.py | 2 +- test/hooks_exponential_moving_average_model_hook_test.py | 6 +++--- test/hooks_time_metrics_hook_test.py | 2 +- test/manual/hooks_progress_bar_hook_test.py | 4 ++-- test/manual/hooks_tensorboard_plot_hook_test.py | 2 +- 13 files changed, 18 insertions(+), 25 deletions(-) diff --git a/classy_vision/hooks/classy_hook.py b/classy_vision/hooks/classy_hook.py index ad5c0a900f..1f98006c47 100644 --- a/classy_vision/hooks/classy_hook.py +++ b/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 diff --git a/classy_vision/hooks/exponential_moving_average_model_hook.py b/classy_vision/hooks/exponential_moving_average_model_hook.py index 284d9415d9..0b19aa5f06 100644 --- a/classy_vision/hooks/exponential_moving_average_model_hook.py +++ b/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) diff --git a/classy_vision/hooks/progress_bar_hook.py b/classy_vision/hooks/progress_bar_hook.py index fc5dfde405..3ce5da3d24 100644 --- a/classy_vision/hooks/progress_bar_hook.py +++ b/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( diff --git a/classy_vision/hooks/tensorboard_plot_hook.py b/classy_vision/hooks/tensorboard_plot_hook.py index 6abeec04ba..ed76e18b42 100644 --- a/classy_vision/hooks/tensorboard_plot_hook.py +++ b/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 = [] diff --git a/classy_vision/hooks/time_metrics_hook.py b/classy_vision/hooks/time_metrics_hook.py index 475282a699..7d789ee04d 100644 --- a/classy_vision/hooks/time_metrics_hook.py +++ b/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 """ diff --git a/classy_vision/tasks/classification_task.py b/classy_vision/tasks/classification_task.py index 899fd99dc0..16922983fa 100644 --- a/classy_vision/tasks/classification_task.py +++ b/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() diff --git a/classy_vision/tasks/classy_task.py b/classy_vision/tasks/classy_task.py index b2c7c46842..3e976c5d4d 100644 --- a/classy_vision/tasks/classy_task.py +++ b/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. diff --git a/classy_vision/trainer/classy_trainer.py b/classy_vision/trainer/classy_trainer.py index cb6d10e191..e0491dbd8b 100644 --- a/classy_vision/trainer/classy_trainer.py +++ b/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) diff --git a/classy_vision/trainer/elastic_trainer.py b/classy_vision/trainer/elastic_trainer.py index c5adde617a..b85ea78609 100644 --- a/classy_vision/trainer/elastic_trainer.py +++ b/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 diff --git a/test/hooks_exponential_moving_average_model_hook_test.py b/test/hooks_exponential_moving_average_model_hook_test.py index 6a1b89a33d..c4b17608ce 100644 --- a/test/hooks_exponential_moving_average_model_hook_test.py +++ b/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)) diff --git a/test/hooks_time_metrics_hook_test.py b/test/hooks_time_metrics_hook_test.py index ab9c0358bd..05412ef88d 100644 --- a/test/hooks_time_metrics_hook_test.py +++ b/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)) diff --git a/test/manual/hooks_progress_bar_hook_test.py b/test/manual/hooks_progress_bar_hook_test.py index bfbd3a9747..0df3fca5ed 100644 --- a/test/manual/hooks_progress_bar_hook_test.py +++ b/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: diff --git a/test/manual/hooks_tensorboard_plot_hook_test.py b/test/manual/hooks_tensorboard_plot_hook_test.py index 033fa46d70..2585cde842 100644 --- a/test/manual/hooks_tensorboard_plot_hook_test.py +++ b/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)