From 351d727206441111a176972981b81f8a234b28d8 Mon Sep 17 00:00:00 2001 From: Aaron Adcock Date: Thu, 2 Apr 2020 07:18:19 -0700 Subject: [PATCH] Add hooks building to classification task (#402) Summary: Pull Request resolved: https://github.com/facebookresearch/ClassyVision/pull/402 Pull Request resolved: https://github.com/fairinternal/ClassyVision/pull/62 Add configurable hooks to classification task, had to remove the typehints from ClassyHook to avoid a circular dependency. As discussed, there are other options (for example importing in from_config), but since we keep running into this issue we are going to remove them to try and improve dev velocity. One more note, as requested I made the checkpoint hook a default option, but this led to some additional changes (need to allow input_args in checkpoint hook to be None as well as empty dict and had to add a make dir call in the task build). Differential Revision: D19770583 fbshipit-source-id: 15a3e6de204a7d89da7312c8a1dc58b085e4768e --- classy_vision/generic/util.py | 2 +- classy_vision/hooks/checkpoint_hook.py | 5 ++--- classy_vision/hooks/classy_hook.py | 12 +++++------- .../hooks/exponential_moving_average_model_hook.py | 11 +++++------ classy_vision/hooks/loss_lr_meter_logging_hook.py | 9 ++++----- classy_vision/hooks/model_complexity_hook.py | 3 +-- classy_vision/hooks/model_tensorboard_hook.py | 3 +-- classy_vision/hooks/profiler_hook.py | 3 +-- classy_vision/hooks/progress_bar_hook.py | 7 +++---- classy_vision/hooks/tensorboard_plot_hook.py | 7 +++---- classy_vision/hooks/visdom_hook.py | 3 +-- classy_vision/tasks/classification_task.py | 10 ++++++++++ test/tasks_classification_task_test.py | 7 +++++++ 13 files changed, 44 insertions(+), 38 deletions(-) diff --git a/classy_vision/generic/util.py b/classy_vision/generic/util.py index 691dbb986a..cf413e4420 100644 --- a/classy_vision/generic/util.py +++ b/classy_vision/generic/util.py @@ -432,7 +432,7 @@ def compute_pr_curves(class_hist, total_hist): def get_checkpoint_dict(task, input_args, deep_copy=False): - assert isinstance( + assert input_args is None or isinstance( input_args, dict ), f"Unexpected input_args of type: {type(input_args)}" return { diff --git a/classy_vision/hooks/checkpoint_hook.py b/classy_vision/hooks/checkpoint_hook.py index 29a84a75da..5b100b3923 100644 --- a/classy_vision/hooks/checkpoint_hook.py +++ b/classy_vision/hooks/checkpoint_hook.py @@ -7,7 +7,6 @@ import logging from typing import Any, Collection, Dict, Optional -from classy_vision import tasks from classy_vision.generic.distributed_util import is_master from classy_vision.generic.util import get_checkpoint_dict, save_checkpoint from classy_vision.hooks import register_hook @@ -85,7 +84,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") -> None: + def on_start(self, task) -> None: if not is_master() or getattr(task, "test_only", False): return if not PathManager.exists(self.checkpoint_folder): @@ -94,7 +93,7 @@ def on_start(self, task: "tasks.ClassyTask") -> None: ) raise FileNotFoundError(err_msg) - def on_phase_end(self, task: "tasks.ClassyTask") -> None: + def on_phase_end(self, task) -> None: """Checkpoint the task every checkpoint_period phases. We do not necessarily checkpoint the task at the end of every phase. diff --git a/classy_vision/hooks/classy_hook.py b/classy_vision/hooks/classy_hook.py index 0bb661f16d..250d2566dc 100644 --- a/classy_vision/hooks/classy_hook.py +++ b/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 diff --git a/classy_vision/hooks/exponential_moving_average_model_hook.py b/classy_vision/hooks/exponential_moving_average_model_hook.py index cd72fb901d..c33b614766 100644 --- a/classy_vision/hooks/exponential_moving_average_model_hook.py +++ b/classy_vision/hooks/exponential_moving_average_model_hook.py @@ -12,7 +12,6 @@ import torch.nn as nn from classy_vision.hooks import register_hook from classy_vision.hooks.classy_hook import ClassyHook -from classy_vision.tasks import ClassyTask @register_hook("ema_model_weights") @@ -78,7 +77,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) -> None: + def on_start(self, task) -> None: if self.state.model_state: # loaded state from checkpoint, do not re-initialize, only move the state # to the right device @@ -93,17 +92,17 @@ 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) -> None: + def on_phase_start(self, task) -> 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) -> None: + def on_phase_end(self, task) -> None: if task.train: # save the current model state since this will be overwritten by the ema # state in the test phase self._save_current_model_state(task.base_model, self.state.model_state) - def on_step(self, task: ClassyTask) -> None: + def on_step(self, task) -> None: if not task.train: return @@ -117,7 +116,7 @@ def on_step(self, task: ClassyTask) -> None: device=self.device ) - def set_model_state(self, task: ClassyTask, use_ema: bool) -> None: + def set_model_state(self, task, use_ema: bool) -> None: """ Depending on use_ema, set the appropriate state for the model. """ diff --git a/classy_vision/hooks/loss_lr_meter_logging_hook.py b/classy_vision/hooks/loss_lr_meter_logging_hook.py index 9d1bd89b38..b6c6150536 100644 --- a/classy_vision/hooks/loss_lr_meter_logging_hook.py +++ b/classy_vision/hooks/loss_lr_meter_logging_hook.py @@ -7,7 +7,6 @@ import logging from typing import Any, Dict, Optional -from classy_vision import tasks from classy_vision.generic.distributed_util import get_rank from classy_vision.hooks import register_hook from classy_vision.hooks.classy_hook import ClassyHook @@ -36,7 +35,7 @@ def __init__(self, log_freq: Optional[int] = None) -> None: ), "log_freq must be an int or None" self.log_freq: Optional[int] = log_freq - def on_phase_end(self, task: "tasks.ClassyTask") -> None: + def on_phase_end(self, task) -> None: """ Log the loss, optimizer LR, and meters for the phase. """ @@ -51,7 +50,7 @@ def on_phase_end(self, task: "tasks.ClassyTask") -> None: if task.train: self._log_lr(task) - def on_step(self, task: "tasks.ClassyTask") -> None: + def on_step(self, task) -> None: """ Log the LR every log_freq batches, if log_freq is not None. """ @@ -63,14 +62,14 @@ def on_step(self, task: "tasks.ClassyTask") -> None: logging.info("Local unsynced metric values:") self._log_loss_meters(task) - def _log_lr(self, task: "tasks.ClassyTask") -> None: + def _log_lr(self, task) -> None: """ Compute and log the optimizer LR. """ optimizer_lr = task.optimizer.parameters.lr logging.info("Learning Rate: {}\n".format(optimizer_lr)) - def _log_loss_meters(self, task: "tasks.ClassyTask") -> None: + def _log_loss_meters(self, task) -> None: """ Compute and log the loss and meters. """ diff --git a/classy_vision/hooks/model_complexity_hook.py b/classy_vision/hooks/model_complexity_hook.py index e8dc5666f7..d97ee3876a 100644 --- a/classy_vision/hooks/model_complexity_hook.py +++ b/classy_vision/hooks/model_complexity_hook.py @@ -7,7 +7,6 @@ import logging from typing import Any, Dict -from classy_vision import tasks from classy_vision.generic.profiler import ( compute_activations, compute_flops, @@ -28,7 +27,7 @@ class ModelComplexityHook(ClassyHook): on_phase_end = ClassyHook._noop on_end = ClassyHook._noop - def on_start(self, task: "tasks.ClassyTask") -> None: + def on_start(self, task) -> 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 index b4510ce736..735678ef71 100644 --- a/classy_vision/hooks/model_tensorboard_hook.py +++ b/classy_vision/hooks/model_tensorboard_hook.py @@ -7,7 +7,6 @@ import logging from typing import Any, Dict -from classy_vision import tasks from classy_vision.generic.distributed_util import is_master from classy_vision.generic.visualize import plot_model from classy_vision.hooks import register_hook @@ -62,7 +61,7 @@ def from_config(cls, config: [Dict[str, Any]]) -> "ModelTensorboardHook": tb_writer = SummaryWriter(**config["summary_writer"]) return cls(tb_writer=tb_writer) - def on_start(self, task: "tasks.ClassyTask") -> None: + def on_start(self, task) -> None: """ Plot the model on Tensorboard. """ diff --git a/classy_vision/hooks/profiler_hook.py b/classy_vision/hooks/profiler_hook.py index 31945dcd5a..ab4910cc5a 100644 --- a/classy_vision/hooks/profiler_hook.py +++ b/classy_vision/hooks/profiler_hook.py @@ -7,7 +7,6 @@ import logging from typing import Any, Dict -from classy_vision import tasks from classy_vision.generic.profiler import profile, summarize_profiler_info from classy_vision.hooks import register_hook from classy_vision.hooks.classy_hook import ClassyHook @@ -25,7 +24,7 @@ class ProfilerHook(ClassyHook): on_phase_end = ClassyHook._noop on_end = ClassyHook._noop - def on_start(self, task: "tasks.ClassyTask") -> None: + def on_start(self, task) -> None: """Profile the forward pass.""" logging.info("Profiling forward pass...") batchsize_per_replica = getattr( diff --git a/classy_vision/hooks/progress_bar_hook.py b/classy_vision/hooks/progress_bar_hook.py index 670a036057..4d561da740 100644 --- a/classy_vision/hooks/progress_bar_hook.py +++ b/classy_vision/hooks/progress_bar_hook.py @@ -6,7 +6,6 @@ from typing import Any, Dict, Optional -from classy_vision import tasks from classy_vision.generic.distributed_util import is_master from classy_vision.hooks import register_hook from classy_vision.hooks.classy_hook import ClassyHook @@ -36,7 +35,7 @@ def __init__(self) -> None: self.bar_size: int = 0 self.batches: int = 0 - def on_phase_start(self, task: "tasks.ClassyTask") -> None: + def on_phase_start(self, task) -> None: """Create and display a progress bar with 0 progress.""" if not progressbar_available: raise RuntimeError( @@ -49,13 +48,13 @@ def on_phase_start(self, task: "tasks.ClassyTask") -> None: self.progress_bar = progressbar.ProgressBar(self.bar_size) self.progress_bar.start() - def on_step(self, task: "tasks.ClassyTask") -> None: + def on_step(self, task) -> None: """Update the progress bar with the batch size.""" if task.train and is_master() and self.progress_bar is not None: self.batches += 1 self.progress_bar.update(min(self.batches, self.bar_size)) - def on_phase_end(self, task: "tasks.ClassyTask") -> None: + def on_phase_end(self, task) -> 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() diff --git a/classy_vision/hooks/tensorboard_plot_hook.py b/classy_vision/hooks/tensorboard_plot_hook.py index b8d5d7414a..18dde65f97 100644 --- a/classy_vision/hooks/tensorboard_plot_hook.py +++ b/classy_vision/hooks/tensorboard_plot_hook.py @@ -8,7 +8,6 @@ import time from typing import Any, Dict, List, Optional -from classy_vision import tasks from classy_vision.generic.distributed_util import is_master from classy_vision.hooks import register_hook from classy_vision.hooks.classy_hook import ClassyHook @@ -70,7 +69,7 @@ def from_config(cls, config: Dict[str, Any]) -> "TensorboardPlotHook": log_period = config.get("log_period", 10) return cls(tb_writer=tb_writer, log_period=log_period) - def on_phase_start(self, task: "tasks.ClassyTask") -> None: + def on_phase_start(self, task) -> None: """Initialize losses and learning_rates.""" self.learning_rates = [] self.wall_times = [] @@ -87,7 +86,7 @@ def on_phase_start(self, task: "tasks.ClassyTask") -> None: f"Parameters/{name}", parameter, global_step=-1 ) - def on_step(self, task: "tasks.ClassyTask") -> None: + def on_step(self, task) -> None: """Store the observed learning rates.""" if self.learning_rates is None: logging.warning("learning_rates is not initialized") @@ -106,7 +105,7 @@ def on_step(self, task: "tasks.ClassyTask") -> None: self.step_idx += 1 - def on_phase_end(self, task: "tasks.ClassyTask") -> None: + def on_phase_end(self, task) -> None: """Add the losses and learning rates to tensorboard.""" if self.learning_rates is None: logging.warning("learning_rates is not initialized") diff --git a/classy_vision/hooks/visdom_hook.py b/classy_vision/hooks/visdom_hook.py index be323483ee..69ba8757bd 100644 --- a/classy_vision/hooks/visdom_hook.py +++ b/classy_vision/hooks/visdom_hook.py @@ -8,7 +8,6 @@ import logging from typing import Any, Dict -from classy_vision import tasks from classy_vision.generic.distributed_util import is_master from classy_vision.generic.util import flatten_dict from classy_vision.generic.visualize import plot_learning_curves @@ -60,7 +59,7 @@ def __init__( self.metrics: Dict = {} self.visdom: Visdom = Visdom(self.server, self.port) - def on_phase_end(self, task: "tasks.ClassyTask") -> None: + def on_phase_end(self, task) -> None: """ Plot the metrics on visdom. """ diff --git a/classy_vision/tasks/classification_task.py b/classy_vision/tasks/classification_task.py index decf0efebf..8c8c089aa7 100644 --- a/classy_vision/tasks/classification_task.py +++ b/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,13 @@ 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 = [] + if hooks_config is not None: + hooks = build_hooks(hooks_config) + optimizer = build_optimizer(optimizer_config) task = ( @@ -348,7 +356,9 @@ 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) diff --git a/test/tasks_classification_task_test.py b/test/tasks_classification_task_test.py index 44617e3e46..7580189e5f 100644 --- a/test/tasks_classification_task_test.py +++ b/test/tasks_classification_task_test.py @@ -44,6 +44,13 @@ def test_build_task(self): task = build_task(config) self.assertTrue(isinstance(task, ClassificationTask)) + def test_hooks_config_builds_correctly(self): + config = get_test_task_config() + config["hooks"] = [{"name": "loss_lr_meter_logging"}] + task = build_task(config) + self.assertTrue(len(task.hooks) == 1) + self.assertTrue(isinstance(task.hooks[0], LossLrMeterLoggingHook)) + def test_get_state(self): config = get_test_task_config() loss = build_loss(config["loss"])