Repository: pfnet/pytorch-pfn-extras

Issue:
Trainer: publicize ExtensionEntry
Allow `manager.extend(extension_entry)` so that users can specify custom trigger etc.


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/training/__init__.py b/pytorch_pfn_extras/training/__init__.py
--- a/pytorch_pfn_extras/training/__init__.py
+++ b/pytorch_pfn_extras/training/__init__.py
@@ -1,4 +1,5 @@
 from pytorch_pfn_extras.training.extension import Extension  # NOQA
+from pytorch_pfn_extras.training.extension import ExtensionEntry  # NOQA
 from pytorch_pfn_extras.training.extension import make_extension  # NOQA
 from pytorch_pfn_extras.training.extension import PRIORITY_EDITOR  # NOQA
 from pytorch_pfn_extras.training.extension import PRIORITY_READER  # NOQA
diff --git a/pytorch_pfn_extras/training/extension.py b/pytorch_pfn_extras/training/extension.py
--- a/pytorch_pfn_extras/training/extension.py
+++ b/pytorch_pfn_extras/training/extension.py
@@ -1,6 +1,8 @@
 from typing import Any, Callable, Dict, Optional, TYPE_CHECKING
 import types
 
+from pytorch_pfn_extras.training import _trigger_util
+
 if TYPE_CHECKING:
     from pytorch_pfn_extras.training.manager import _BaseExtensionsManager
     from pytorch_pfn_extras.training._trigger_util import TriggerLike
@@ -222,3 +224,58 @@ def decorator(ext: 'ExtensionLike') -> 'ExtensionLike':
 
 def _as_extension(ext: 'ExtensionLike') -> Extension:
     return ext if isinstance(ext, Extension) else _WrappedExtension(ext)
+
+
+class ExtensionEntry:
+    """Extension and options.
+    When name, priority, or trigger is not specified, it is copied from the
+    attributes of the given ``extension``.
+
+    Args:
+        extension: An extension.
+        name: Name of extension.
+        priority: Invocation priority of the extension.
+        trigger: Trigger object that determines when to invoke the extension.
+        call_before_training: Flag to call extension before training.
+
+    .. seealso::
+       :meth:`pytorch_pfn_extras.training.ExtensionsManager.extend`
+    """
+
+    def __init__(
+            self,
+            extension: 'ExtensionLike',
+            *,
+            name: Optional[str] = None,
+            priority: Optional[int] = None,
+            trigger: Optional['TriggerLike'] = None,
+            call_before_training: bool = False,
+    ) -> None:
+        self.extension = _as_extension(extension)
+        self.priority = priority or self.extension.priority
+        self.call_before_training = call_before_training
+
+        self._update_trigger(trigger or self.extension.trigger)
+        self._update_name(name or self.extension.name or self.extension.default_name)
+
+    def _update_trigger(self, trigger: 'TriggerLike') -> None:
+        self.trigger = _trigger_util.get_trigger(trigger)
+
+    def _update_name(self, name: str) -> None:
+        if name == 'training':
+            raise ValueError(
+                'the name "training" is prohibited as an extension name')
+        self.name = name
+        self.extension.name = name
+
+    def state_dict(self) -> Dict[str, Any]:
+        state = {}
+        state['extension'] = self.extension.state_dict()
+        state['trigger'] = self.trigger.state_dict()
+        return state
+
+    def load_state_dict(self, to_load: Dict[str, Any]) -> None:
+        if 'extension' in to_load:
+            self.extension.load_state_dict(to_load['extension'])
+        if 'trigger' in to_load:
+            self.trigger.load_state_dict(to_load['trigger'])
diff --git a/pytorch_pfn_extras/training/manager.py b/pytorch_pfn_extras/training/manager.py
--- a/pytorch_pfn_extras/training/manager.py
+++ b/pytorch_pfn_extras/training/manager.py
@@ -27,33 +27,6 @@ def defer(self):
         self._is_completed = False
 
 
-class _ExtensionEntry:
-
-    def __init__(
-            self,
-            extension: extension_module.Extension,
-            priority: int,
-            trigger: trigger_module.Trigger,
-            call_before_training: bool
-    ) -> None:
-        self.extension = extension
-        self.trigger = trigger
-        self.priority = priority
-        self.call_before_training = call_before_training
-
-    def state_dict(self) -> Dict[str, Any]:
-        state = {}
-        state['extension'] = self.extension.state_dict()
-        state['trigger'] = self.trigger.state_dict()
-        return state
-
-    def load_state_dict(self, to_load: Dict[str, Any]) -> None:
-        if 'extension' in to_load:
-            self.extension.load_state_dict(to_load['extension'])
-        if 'trigger' in to_load:
-            self.trigger.load_state_dict(to_load['trigger'])
-
-
 class _ManagerExecutionProxy:
     """
     Object that is passed to the triggers of extensions depending if they
@@ -149,7 +122,8 @@ def __init__(
         # Defer!
         self._start_time: Optional[float] = None
         self._iters_per_epoch: Optional[int] = None
-        self._extensions: Dict[str, _ExtensionEntry] = collections.OrderedDict()
+        self._extensions: Dict[
+            str, extension_module.ExtensionEntry] = collections.OrderedDict()
         for ext in extensions:
             self.extend(ext)
 
@@ -291,12 +265,15 @@ def start_extensions(self) -> None:
 
     def extend(
             self,
-            extension: 'extension_module.ExtensionLike',
+            extension: Union[
+                'extension_module.ExtensionLike',
+                'extension_module.ExtensionEntry',
+            ],
             name: Optional[str] = None,
             trigger: 'trigger_module.TriggerLike' = None,
             priority: Optional[int] = None,
             *,
-            call_before_training: bool = False,
+            call_before_training: Optional[bool] = None,
             **kwargs: Dict[str, Any],
     ) -> None:
         """Registers an extension to the manager.
@@ -340,29 +317,29 @@ def extend(
         if self._start_extensions_called:
             raise RuntimeError(
                 'extend called after the extensions were initialized')
-        ext = extension_module._as_extension(extension)
-        if name is None:
-            name = ext.name or ext.default_name
-        if name == 'training':
-            raise ValueError(
-                'the name "training" is prohibited as an extension name')
 
-        if trigger is None:
-            trigger = ext.trigger
-        trigger = trigger_module.get_trigger(trigger)
+        if isinstance(extension, extension_module.ExtensionEntry):
+            entry = extension
+        else:
+            entry = extension_module.ExtensionEntry(extension)
+
+        if trigger is not None:
+            entry._update_trigger(trigger)
+
+        if priority is not None:
+            entry.priority = priority
 
-        if priority is None:
-            priority = ext.priority
+        if call_before_training is not None:
+            entry.call_before_training = call_before_training
 
-        modified_name = name
+        modified_name = name or entry.name
         ordinal = 0
         while modified_name in self._extensions:
             ordinal += 1
             modified_name = '%s_%d' % (name, ordinal)
 
-        ext.name = modified_name
-        self._extensions[modified_name] = _ExtensionEntry(
-            ext, priority, trigger, call_before_training)
+        entry._update_name(modified_name)
+        self._extensions[modified_name] = entry
 
     def get_extension(self, name: str) -> extension_module.Extension:
         """Returns the extension of a given name.


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_extension_entry.py b/tests/pytorch_pfn_extras_tests/training_tests/test_extension_entry.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_extension_entry.py
@@ -0,0 +1,63 @@
+import torch
+
+import pytorch_pfn_extras as ppe
+
+
+def _get_dummy_manager():
+    model = torch.nn.Module()
+    return ppe.training.ExtensionsManager(
+        {'main': model},
+        [],  # optimizers
+        10,  # max_epochs
+        iters_per_epoch=1,
+    )
+
+
+def test_default_name():
+    class MyExtension(ppe.training.Extension):
+        name = None
+        default_name = 'defalut_name'
+
+    ext = MyExtension()
+    entry = ppe.training.ExtensionEntry(ext)
+    assert entry.name == MyExtension.default_name
+    entry = ppe.training.ExtensionEntry(ext, name='updated')
+    assert entry.name == 'updated'
+
+
+def test_name():
+    class MyExtension(ppe.training.Extension):
+        name = 'name'
+        default_name = 'defalut_name'
+
+    ext = MyExtension()
+    entry = ppe.training.ExtensionEntry(ext)
+    assert entry.name == MyExtension.name
+    entry = ppe.training.ExtensionEntry(ext, name='updated')
+    assert entry.name == 'updated'
+
+
+def test_priority():
+    class MyExtension(ppe.training.Extension):
+        priority = 100
+
+    ext = MyExtension()
+    entry = ppe.training.ExtensionEntry(ext)
+    assert entry.priority == MyExtension.priority
+    entry = ppe.training.ExtensionEntry(ext, priority=10)
+    assert entry.priority == 10
+
+
+def test_trigger():
+    class MyExtension(ppe.training.Extension):
+        trigger = (1, 'iteration')
+
+    ext = MyExtension()
+    entry = ppe.training.ExtensionEntry(ext)
+    assert isinstance(entry.trigger, ppe.training.triggers.IntervalTrigger)
+    assert entry.trigger.period == 1
+    assert entry.trigger.unit == 'iteration'
+    entry = ppe.training.ExtensionEntry(ext, trigger=(3, 'epoch'))
+    assert isinstance(entry.trigger, ppe.training.triggers.IntervalTrigger)
+    assert entry.trigger.period == 3
+    assert entry.trigger.unit == 'epoch'
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py b/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
@@ -69,6 +69,14 @@ def test_extensions_manager_extensions():
         _DummyExtensionInitialize(3, call_record, init_record),
         _DummyExtensionInitialize(4, call_record, init_record),
         lambda manager: dummy5(manager),
+        training.ExtensionEntry(
+            _DummyExtension(6, call_record, init_record),
+            name='ext6', priority=-3, call_before_training=True,
+        ),
+        training.ExtensionEntry(
+            _DummyExtension(7, call_record, init_record),
+            name='ext7_', priority=-2, call_before_training=True,
+        ),
     ]
 
     manager.extend(exts[0], 'ext0', priority=2, call_before_training=True)
@@ -77,12 +85,16 @@ def test_extensions_manager_extensions():
     manager.extend(exts[3], 'ext3', priority=0, call_before_training=True)
     manager.extend(exts[4], 'ext4', priority=4, call_before_training=True)
     manager.extend(exts[5], 'ext5', priority=-1, call_before_training=True)
+    manager.extend(exts[6])
+    manager.extend(exts[7], 'ext7', priority=-4, call_before_training=False)
 
     assert manager.get_extension('ext0') is exts[0]
     assert manager.get_extension('ext1') is exts[1]
     assert manager.get_extension('ext2') is exts[2]
     assert manager.get_extension('ext3') is exts[3]
     assert manager.get_extension('ext4') is exts[4]
+    assert manager.get_extension('ext6') is exts[6].extension
+    assert manager.get_extension('ext7') is exts[7].extension
 
     with pytest.raises(ValueError):
         manager.get_extension('ext10')
@@ -95,7 +107,7 @@ def test_extensions_manager_extensions():
             assert manager.iteration == it
 
             if it == 0:
-                assert call_record == [4, 0, 3, 5]
+                assert call_record == [4, 0, 3, 5, 6]
                 assert init_record == [4, 1, 3]
             else:
                 assert call_record == []
@@ -104,7 +116,7 @@ def test_extensions_manager_extensions():
             call_record.clear()
             init_record.clear()
 
-        assert call_record == [4, 2, 0, 1, 3, 5]
+        assert call_record == [4, 2, 0, 1, 3, 5, 6, 7]
         assert init_record == []