Merge pull request #3084 from plonerma/pluggable_trainer

Proposal: Pluggable `ModelTrainer` train function

Author: alanakbik

flair/trainers/plugins/__init__.py

@@ -2,7 +2,6 @@
     AmpPlugin,
     CheckpointPlugin,
     ModelCardPlugin,
-    RegularLoggingPlugin,
     SchedulerPlugin,
     SWAPlugin,
     WeightExtractorPlugin,
@@ -13,51 +12,35 @@
     MetricHistoryPlugin,
     TensorboardLogger,
 )
-from flair.trainers.plugins.metrics import (
-    BasicEvaluationPlugin,
-    MetricBasePlugin,
-    MetricName,
-    MetricRecord,
-    TrainingBehaviorPlugin,
-)
 
 from .base import BasePlugin, Pluggable, TrainerPlugin, TrainingInterrupt
+from .metric_records import MetricName, MetricRecord
 
 default_plugins = [
-    BasicEvaluationPlugin,
-    TrainingBehaviorPlugin,
-    RegularLoggingPlugin,
-    AmpPlugin,
     CheckpointPlugin,
     ModelCardPlugin,
     SchedulerPlugin,
-    SWAPlugin,
     WeightExtractorPlugin,
     LossFilePlugin,
     MetricHistoryPlugin,
-    TensorboardLogger,
     LogFilePlugin,
 ]
 
 __all__ = [
     "AmpPlugin",
     "CheckpointPlugin",
     "ModelCardPlugin",
-    "RegularLoggingPlugin",
     "SchedulerPlugin",
     "SWAPlugin",
     "WeightExtractorPlugin",
     "LogFilePlugin",
     "LossFilePlugin",
     "MetricHistoryPlugin",
     "TensorboardLogger",
-    "BasicEvaluationPlugin",
-    "TrainingBehaviorPlugin",
     "BasePlugin",
     "Pluggable",
     "TrainerPlugin",
     "TrainingInterrupt",
-    "MetricBasePlugin",
     "MetricName",
     "MetricRecord",
 ]

flair/trainers/plugins/base.py

@@ -6,7 +6,6 @@
 from typing import (
     Callable,
     Dict,
-    Iterable,
     Iterator,
     List,
     NewType,
@@ -36,7 +35,7 @@ class Pluggable:
 
     valid_events: Optional[Set[EventIdenifier]] = None
 
-    def __init__(self, *, plugins: Sequence[PluginArgument] = None):
+    def __init__(self, *, plugins: Sequence[PluginArgument] = []):
         """Initialize a `Pluggable`.
 
         :param plugins: Plugins which should be attached to this `Pluggable`.
@@ -50,14 +49,13 @@ def __init__(self, *, plugins: Sequence[PluginArgument] = None):
         self._event_queue: Queue = Queue()
         self._processing_events = False
 
-        if plugins is not None:
-            for plugin in plugins:
-                if isclass(plugin):
-                    # instantiate plugin
-                    plugin = plugin()
+        for plugin in plugins:
+            if isclass(plugin):
+                # instantiate plugin
+                plugin = plugin()
 
-                plugin = cast("BasePlugin", plugin)
-                plugin.attach_to(self)
+            plugin = cast("BasePlugin", plugin)
+            plugin.attach_to(self)
 
     @property
     def plugins(self):
@@ -72,7 +70,7 @@ def validate_event(self, *events: EventIdenifier):
 
             if self.valid_events is not None:
                 if event not in self.valid_events:
-                    raise RuntimeError(f"Event '{event}' not recognized (available {self.valid_events})")
+                    raise RuntimeError(f"Event '{event}' not recognized. Available: {', '.join(self.valid_events)}")
             return event
 
     def register_hook(self, func: Callable, *events: EventIdenifier):
@@ -92,31 +90,23 @@ def register_hook(self, func: Callable, *events: EventIdenifier):
             self._hook_handles[event][handle.id] = handle
         return handle
 
-    def dispatch(self, event: EventIdenifier, *args, **kwargs) -> dict:
+    def dispatch(self, event: EventIdenifier, *args, **kwargs) -> None:
         """Call all functions hooked to a certain event."""
         self.validate_event(event)
 
-        events_return_value: dict = {}
-        self._event_queue.put((event, args, kwargs, events_return_value))
+        self._event_queue.put((event, args, kwargs))
 
         if not self._processing_events:
             self._processing_events = True
 
             while not self._event_queue.empty():
-                event, args, kwargs, combined_return_values = self._event_queue.get()
+                event, args, kwargs = self._event_queue.get()
 
                 for hook in self._hook_handles[event].values():
-                    returned = hook(*args, **kwargs)
-
-                    if returned is not None:
-                        combined_return_values.update(returned)
+                    hook(*args, **kwargs)
 
             self._processing_events = False
 
-        # this dict may be empty and will be complete once all events have been
-        # processed
-        return events_return_value
-
     def remove_hook(self, handle: "HookHandle"):
         """Remove a hook handle from this instance."""
         for event in handle.events:
@@ -146,6 +136,10 @@ def id(self) -> HookHandleId:
         """Return the id of this `HookHandle`."""
         return self._id
 
+    @property
+    def func_name(self):
+        return self._func.__qualname__
+
     @property
     def events(self) -> Iterator[EventIdenifier]:
         """Return iterator of events whis `HookHandle` is registered for."""
@@ -165,7 +159,7 @@ def __call__(self, *args, **kw):
             for name in kw.keys():
                 if name not in sig.parameters:
                     raise TypeError(
-                        f"Hook callback {self._func.__qualname__}() does not accept keyword argument '{name}'"
+                        f"Hook callback {self.func_name}() does not accept keyword argument '{name}'"
                     ) from err
 
             raise err
@@ -174,10 +168,6 @@ def __call__(self, *args, **kw):
 class BasePlugin:
     """Base class for all plugins."""
 
-    provided_events: Optional[Set[EventIdenifier]] = None
-
-    dependencies: Iterable[Type["BasePlugin"]] = ()
-
     def __init__(self):
         """Initialize the base plugin."""
         self._hook_handles: List[HookHandle] = []
@@ -190,23 +180,6 @@ def attach_to(self, pluggable: Pluggable):
 
         self._pluggable = pluggable
 
-        for dep in self.dependencies:
-            dep_satisfied = False
-
-            for plugin in pluggable.plugins:
-                if isinstance(plugin, dep):
-                    # there is already a plugin which satisfies this dependency
-                    dep_satisfied = True
-                    break
-
-            if not dep_satisfied:
-                # create a plugin of this type and attach it to the trainer
-                dep_plugin = dep()
-                dep_plugin.attach_to(pluggable)
-
-        if self.provided_events is not None and pluggable.valid_events is not None:
-            pluggable.valid_events = pluggable.valid_events | self.provided_events
-
         pluggable.append_plugin(self)
 
         # go through all attributes

flair/trainers/plugins/functional/__init__.py

@@ -1,6 +1,5 @@
 from .amp import AmpPlugin
 from .checkpoints import CheckpointPlugin
-from .logging import RegularLoggingPlugin
 from .model_card import ModelCardPlugin
 from .scheduler import SchedulerPlugin
 from .swa import SWAPlugin
@@ -13,5 +12,4 @@
     "SchedulerPlugin",
     "SWAPlugin",
     "WeightExtractorPlugin",
-    "RegularLoggingPlugin",
 ]

flair/trainers/plugins/functional/amp.py

@@ -2,11 +2,6 @@
 
 from flair.trainers.plugins.base import TrainerPlugin
 
-try:
-    from apex import amp
-except ImportError:
-    amp = None
-
 
 class AmpPlugin(TrainerPlugin):
     """
@@ -15,52 +10,53 @@ class AmpPlugin(TrainerPlugin):
 
     def __init__(self):
         super().__init__()
-        self.use = None  # TODO: can be removed
-        self.opt_level = None # TODO: I think this also since only used in 1 place
+
         self.wrapped_backward = None
+        self.amp = None
 
     @TrainerPlugin.hook
-    def before_training_setup(self, use_amp, amp_opt_level, **kw):
-        self.use = use_amp # TODO: can be removed
-        self.opt_level = amp_opt_level
-
-        if self.use:
-            if sys.version_info < (3, 0):
-                raise RuntimeError("Apex currently only supports Python 3. Aborting.")
-            if amp is None:
-                raise RuntimeError(
-                    "Failed to import apex. Please install apex from "
-                    "https://www.github.com/nvidia/apex "
-                    "to enable mixed-precision training."
-                )
+    def before_training_setup(self, **kw):
+        if sys.version_info < (3, 0):
+            raise RuntimeError("Apex currently only supports Python 3. Aborting.")
+
+        try:
+            from apex import amp
+
+            self.amp = amp
+        except ImportError as exc:
+            raise RuntimeError(
+                "Failed to import apex. Please install apex from "
+                "https://www.github.com/nvidia/apex "
+                "to enable mixed-precision training."
+            ) from exc
 
     def detach(self, *args, **kwargs):
         # TODO: what does this do?
         super().detach(*args, **kwargs)
 
+        # unwrap trainer backward function
         self.trainer.backward = self.wrapped_backward
         self.wrapped_backward = None
 
     def backward(self, loss):
+        assert self.amp is not None
         optimizer = self.trainer.optimizer
 
-        if self.use:  # TODO: can be removed
-            with amp.scale_loss(loss, optimizer) as scaled_loss:
-                scaled_loss.backward()
+        with self.amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
 
     @TrainerPlugin.hook
-    def after_optimizer_setup(self, **kw):
+    def after_optimizer_setup(self, amp_opt_level, **kw):
         """
         Wraps with AMP
         :param kw:
         :return:
         """
         optimizer = self.trainer.optimizer
 
-        if self.use:  # TODO: can be removed
-            self.trainer.model, self.trainer.optimizer = amp.initialize(self.model, optimizer, opt_level=self.opt_level)
+        self.trainer.model, self.trainer.optimizer = self.amp.initialize(self.model, optimizer, opt_level=amp_opt_level)
 
-            # replace trainers backward function
-            self.wrapped_backward = self.trainer.backward
+        # replace trainers backward function
+        self.wrapped_backward = self.trainer.backward
 
-            self.trainer.backward = self.backward
+        self.trainer.backward = self.backward