returnn 1.20240120.124314__tar.gz → 1.20240129.140703__tar.gz

{returnn-1.20240120.124314/returnn.egg-info → returnn-1.20240129.140703}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240120.124314
+Version: 1.20240129.140703
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240129.140703/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240129.140703'
+long_version = '1.20240129.140703+git.e6f3f5d'

{returnn-1.20240120.124314 → returnn-1.20240129.140703}/returnn/datasets/hdf.py

@@ -438,7 +438,7 @@ class HDFDataset(CachedDataset):
         :param str key:
         :rtype: bool
         """
-        if self.get_data_dtype(key).startswith("int"):
+        if "int" in self.get_data_dtype(key):
             if key in self.num_outputs:
                 return self.num_outputs[key][1] <= 1
         return False

{returnn-1.20240120.124314 → returnn-1.20240129.140703}/returnn/frontend/__init__.py

@@ -49,4 +49,5 @@ from .types import *
 
 # Modules not in the main namespace but in sub namespaces.
 from . import audio
+from . import hooks
 from . import init

returnn-1.20240129.140703/returnn/frontend/hooks.py

@@ -0,0 +1,165 @@
+"""
+Hooks logic
+"""
+
+from __future__ import annotations
+from typing import Optional, Any, Dict, Callable, TypeVar, Tuple
+from types import MethodType
+from collections import OrderedDict
+import weakref
+
+
+__all__ = ["setup_post_hook_on_method", "RemovableHandle"]
+
+
+T = TypeVar("T")
+
+
+def setup_post_hook_on_method(
+    obj: Any,
+    attr: str,
+    hook: Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
+    *,
+    prepend: bool = False,
+) -> RemovableHandle:
+    """
+    :param obj: e.g. module
+    :param attr: attribute name to get the method, via ``getattr(obj, attr)``
+    :param hook:
+    :param prepend: if there are multiple hooks, this will be registered in front of all, otherwise at the end
+    """
+    method = MethodWithHooks.get(obj, attr)
+    return _register_new_hook(method, method.post_hooks, hook, prepend=prepend)
+
+
+class MethodWithHooks:
+    """
+    Method with hooks
+
+    (Not sure yet whether this should be exposed in the API...)
+    """
+
+    @classmethod
+    def get(cls, obj: Any, attr: str) -> MethodWithHooks:
+        """get existing or init new :class:`MethodWithHooks`"""
+        method = getattr(obj, attr)
+        if not isinstance(method, MethodWithHooks):
+            assert (
+                isinstance(method, MethodType)
+                and method.__self__ is obj
+                and method.__func__ is getattr(obj.__class__, attr)
+            ), (
+                f"{obj}.{attr} is {method}, {method.__self__} is not obj"
+                f" or {method.__func__} is not {getattr(obj.__class__, attr)}"
+            )  # not necessary, just what we expect here...
+            method = MethodWithHooks(obj, attr)
+            method.setup()
+        return method
+
+    def __init__(self, obj: Any, attr: str):
+        """
+        :param obj:
+        :param attr:
+        """
+        self.obj = obj
+        self.attr = attr
+        self.post_hooks: Dict[int, Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]]] = OrderedDict()
+
+    def setup(self):
+        """setup on object"""
+        setattr(self.obj, self.attr, self)
+        if self.attr == "__call__":
+            self.obj.__class__.__call__ = _CallWrapperClass(self.obj.__class__)
+
+    def uninstall(self):
+        """undo the setup"""
+        assert not self.post_hooks
+        delattr(self.obj, self.attr)
+        if self.attr == "__call__":
+            call_wrapper = self.obj.__class__.__call__
+            assert isinstance(call_wrapper, _CallWrapperClass)
+            self.obj.__class__.__call__ = call_wrapper.orig_call
+
+    def has_any_hook(self) -> bool:
+        """whether there is any hook"""
+        return bool(self.post_hooks)
+
+    def __call__(self, *args, **kwargs):
+        function = getattr(self.obj.__class__, self.attr)
+        if isinstance(function, _CallWrapperClass):
+            function = function.orig_call
+        result = function(self.obj, *args, **kwargs)
+        for hook in self.post_hooks.values():
+            result = hook(self.obj, args, kwargs, result)
+        return result
+
+
+class _CallWrapperClass:
+    """
+    We need special handling for ``obj(*args, **kwargs)`` - simply wrapping ``__call__`` is not enough.
+    An instance of this class here would overwrite the obj.__class__.__call__.
+    """
+
+    def __init__(self, cls):
+        self.cls = cls
+        self.orig_call = getattr(cls, "__call__")
+
+    def __call__(self, obj, *args, **kwargs):  # called on the class
+        call_func = getattr(obj, "__call__")
+        if isinstance(call_func, MethodWithHooks):  # obj has hooks installed
+            return call_func(*args, **kwargs)
+        # This is what we would expect currently, but the user might have dome some weird things,
+        # e.g. overwritten obj.__call__ or obj.__class__.__call__ in some other way,
+        # so maybe extend this then, if we encounter this.
+        # (Only do it when we really now the case properly, and add a test case for it.)
+        assert isinstance(call_func, MethodType)
+        assert call_func.__func__ is self
+        assert call_func.__self__ is obj
+        return self.orig_call(obj, *args, **kwargs)
+
+    def __get__(self, obj, objtype=None):
+        if obj is not None:
+            return MethodType(self, obj)
+        return self
+
+
+def _register_new_hook(
+    method: MethodWithHooks,
+    hooks_dict: Dict[int, Any],
+    hook: Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
+    *,
+    prepend: bool = False,
+) -> RemovableHandle:
+    """register"""
+    handle = RemovableHandle(method, hooks_dict)
+    hooks_dict[handle.id] = hook
+    if prepend:
+        assert isinstance(hooks_dict, OrderedDict)
+        hooks_dict.move_to_end(handle.id, last=False)
+    return handle
+
+
+class RemovableHandle:
+    """
+    A handle which provides the capability to remove a hook.
+
+    Code partly copied from PyTorch.
+    """
+
+    id: int
+    next_id: int = 0
+
+    def __init__(self, method: MethodWithHooks, hooks_dict: Dict[int, Any]):
+        self.method_ref = weakref.ref(method)
+        self.hooks_dict_ref = weakref.ref(hooks_dict)
+        self.id = RemovableHandle.next_id
+        RemovableHandle.next_id += 1
+
+    def remove(self) -> None:
+        """remove hook"""
+        method: MethodWithHooks = self.method_ref()
+        hooks_dict = self.hooks_dict_ref()
+        if hooks_dict is not None and self.id in hooks_dict:
+            del hooks_dict[self.id]
+            if method and not method.has_any_hook():  # we removed the last hook
+                method.uninstall()

{returnn-1.20240120.124314 → returnn-1.20240129.140703}/returnn/frontend/module.py

@@ -3,7 +3,7 @@ Base module class, :class:`Module`.
 """
 
 from __future__ import annotations
-from typing import Any, Optional, Sequence, List, Tuple, Union, Set, Iterator, Callable, TypeVar
+from typing import Any, Optional, Sequence, List, Tuple, Union, Set, Iterator, Callable, TypeVar, Dict
 from returnn.util.basic import OptionalNotImplementedError, RefIdEq
 from returnn.tensor import Tensor, Dim
 from .. import frontend as rf
@@ -247,6 +247,24 @@ class Module:
             param: rf.Parameter
             param.to(device=device, dtype=dtype)
 
+    def register_forward_hook(
+        self,
+        hook: Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
+        *,
+        prepend: bool = False,
+    ) -> rf.hooks.RemovableHandle:
+        """
+        Register forward hook. Very similar to PyTorch. (Code even partly copied from PyTorch.)
+        This will be called after __call__, following this logic (where module=self):
+
+            result = module(*args, **kwargs)
+            result = hook(args, kwargs, result)
+
+        :param hook:
+        :param prepend: if there are multiple hooks, this will be registered in front of all, otherwise at the end
+        """
+        return rf.hooks.setup_post_hook_on_method(self, "__call__", hook=hook, prepend=prepend)
+
 
 class Functional(Module):
     """

{returnn-1.20240120.124314 → returnn-1.20240129.140703/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240120.124314
+Version: 1.20240129.140703
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20240120.124314 → returnn-1.20240129.140703}/returnn.egg-info/SOURCES.txt

@@ -172,6 +172,7 @@ returnn/frontend/dropout.py
 returnn/frontend/dtype.py
 returnn/frontend/gradient.py
 returnn/frontend/graph.py
+returnn/frontend/hooks.py
 returnn/frontend/init.py
 returnn/frontend/label_smoothing.py
 returnn/frontend/linear.py

{returnn-1.20240120.124314 → returnn-1.20240129.140703}/tests/test_torch_frontend.py

@@ -533,6 +533,54 @@ def test_Data_copy_tranpose_match_priority():
     numpy.testing.assert_equal(x_np, raw_np)
 
 
+def test_forward_hook():
+    from returnn.frontend import hooks
+    from types import MethodType, FunctionType
+
+    call_count = 0
+    hook_call_count = 0
+
+    class _Module(rf.Module):
+        def __call__(self, x: Tensor) -> Tensor:
+            nonlocal call_count
+            call_count += 1
+            return x * 0.5
+
+    def _hook(module, args, kwargs, result):
+        nonlocal hook_call_count
+        assert isinstance(module, _Module) and module is mod
+        assert isinstance(args, tuple) and len(args) == 1 and isinstance(args[0], Tensor)
+        assert isinstance(kwargs, dict) and len(kwargs) == 0
+        assert isinstance(result, Tensor)
+        hook_call_count += 1
+        return result * 2.0
+
+    mod = _Module()
+    time_dim = Dim(None, name="time")
+    x = Tensor("x", [time_dim], "float32", raw_tensor=torch.range(0.0, 5.0))
+    mod(x)
+    assert call_count == 1 and hook_call_count == 0
+
+    handle = mod.register_forward_hook(_hook)
+    assert isinstance(handle, hooks.RemovableHandle)
+    assert isinstance(mod.__call__, hooks.MethodWithHooks)
+    assert isinstance(mod.__class__.__call__, hooks._CallWrapperClass)
+    mod(x)
+    assert call_count == 2 and hook_call_count == 1
+
+    mod.__call__(x)
+    assert call_count == 3 and hook_call_count == 2
+
+    handle.remove()
+    assert isinstance(mod.__call__, MethodType)
+    assert isinstance(mod.__class__.__call__, FunctionType)
+    mod(x)
+    assert call_count == 4 and hook_call_count == 2
+
+    mod.__call__(x)
+    assert call_count == 5 and hook_call_count == 2
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

{returnn-1.20240120.124314 → returnn-1.20240129.140703}/tools/torch_avg_checkpoints.py

@@ -67,7 +67,11 @@ def merge_checkpoints(in_ckpts: Sequence[str], out_ckpt: str, extra_state: Optio
     out_state: Dict[str, Any] = {"model": out_model_state, "merged_epochs": [], "merged_steps": []}
     for in_ckpt in in_ckpts:
         print("read ckpt:", in_ckpt)
-        in_state = torch.load(in_ckpt, map_location=torch.device("cpu"), mmap=True)
+        torch_version = tuple(int(s) for s in str(torch.__version__).split(".")[:2])
+        load_kwargs = dict(map_location=torch.device("cpu"), mmap=True)
+        if torch_version < (2, 1):
+            load_kwargs.pop("mmap")  # mmap flag only introduced from 2.1.0 onwards
+        in_state = torch.load(in_ckpt, **load_kwargs)
         assert isinstance(in_state, dict)
 
         assert "model" in in_state and isinstance(in_state["model"], dict)

returnn-1.20240120.124314/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240120.124314'
-long_version = '1.20240120.124314+git.91ba673'