returnn 1.20230815.191535__tar.gz → 1.20230816.2116__tar.gz

{returnn-1.20230815.191535/returnn.egg-info → returnn-1.20230816.2116}/PKG-INFO

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

returnn-1.20230816.2116/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230816.002116'
+long_version = '1.20230816.002116+git.e1ca657'

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/returnn/frontend/__init__.py

@@ -45,6 +45,7 @@ from .signal import *
 from .types import *
 
 # Modules not in the main namespace but in sub namespaces.
+from . import audio
 from . import init
 
 # And some functions from the internal backend API.

returnn-1.20230816.2116/returnn/frontend/audio/__init__.py

@@ -0,0 +1,5 @@
+"""
+Audio utils
+"""
+
+from .mel import *

returnn-1.20230815.191535/returnn/frontend/signal.py → returnn-1.20230816.2116/returnn/frontend/audio/mel.py

@@ -1,103 +1,19 @@
 """
-stft etc
+Mel filterbank.
 """
 
 
 from __future__ import annotations
 from typing import Optional, Union, Tuple
+import functools
 import math
 import numpy
-import functools
 from returnn.util import math as util_math
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 
 
-__all__ = ["stft", "mel_filterbank", "log_mel_filterbank_from_raw"]
-
-
-def stft(
-    x: Tensor,
-    *,
-    in_spatial_dim: Dim,
-    frame_step: int,
-    frame_length: int,
-    fft_length: Optional[int] = None,
-    window_use_frame_length: bool = True,
-    align_window_left: bool = True,
-    window_enforce_even: bool = True,
-    out_spatial_dim: Optional[Dim] = None,
-    out_dim: Optional[Dim] = None,
-) -> Tuple[Tensor, Dim, Dim]:
-    """
-    Calculate the short-time Fourier transform (STFT) of a signal.
-
-    We currently always use the Hann window.
-
-    Note that there are inconsistencies between frameworks (e.g. PyTorch vs TensorFlow).
-    We have options where you can explicitly specify the behavior,
-    to match either PyTorch or TensorFlow.
-    See here for some discussion and demonstration:
-    https://github.com/pytorch/pytorch/issues/100177
-    https://github.com/librosa/librosa/issues/1701
-    https://github.com/albertz/playground/blob/master/tf_pt_stft.py
-
-    :param x: (..., in_spatial_dim, ...). any other dims are treated as batch dims.
-    :param in_spatial_dim:
-    :param frame_length: win_length in PT/librosa, frame_size in StftLayer, frame_length in TF
-    :param frame_step: hop_length in PT/librosa, frame_shift in StftLayer, frame_step in TF
-    :param fft_length: n_fft in PT/librosa, fft_size in StftLayer, fft_length in TF.
-        If not given, will use frame_length.
-    :param window_use_frame_length: Whether to use window size = frame_length.
-        If False, uses window size = fft_length.
-        This has only an effect if frame_length != fft_length.
-        This has an effect on the output seq length.
-        If True: out_seq_len = ⌈(T - frame_length + 1) / frame_step⌉ (T = in_seq_len).
-        If False: out_seq_len = ⌈(T - fft_length + 1) / frame_step⌉.
-        Note that in TF/SciPy, the behavior matches to window_use_frame_length=True,
-        but in PT/librosa, the behavior matches to window_use_frame_length=False.
-    :param align_window_left: whether to align the window to the left inside the potential wider fft_length window.
-        This only has an effect if frame_length != fft_length.
-        If True, the window will be aligned to the left and the right side will be padded with zeros.
-        Then effectively the remaining phases are not used.
-        If False, the window will be aligned to the center and the left and right side will be padded with zeros.
-        Then effectively the beginning and end phases are not used.
-        In TF/SciPy, the behavior matches to align_window_left=True,
-        but in PT/librosa, the behavior matches to align_window_left=False.
-    :param window_enforce_even: enforce that the window size is even, potentially lowering the window size by 1.
-        This only has an effect if the window size is uneven.
-        In TF, the behavior matches to window_enforce_even=True,
-        but in most other frameworks, the behavior matches to window_enforce_even=False.
-    :param out_spatial_dim:
-    :param out_dim:
-    :return: (stft, out_spatial_dim, out_dim)
-    """
-    fft_length = fft_length or frame_length
-    if out_dim is None:
-        out_dim = Dim(fft_length // 2 + 1, name="stft-freq")
-    if out_spatial_dim is None:
-        from .conv import make_conv_out_spatial_dims
-
-        (out_spatial_dim,) = make_conv_out_spatial_dims(
-            [in_spatial_dim], filter_size=frame_length, strides=frame_step, padding="valid"
-        )
-    # noinspection PyProtectedMember
-    return (
-        x._raw_backend.stft(
-            x,
-            in_spatial_dim=in_spatial_dim,
-            frame_step=frame_step,
-            frame_length=frame_length,
-            fft_length=fft_length,
-            window_use_frame_length=window_use_frame_length,
-            align_window_left=align_window_left,
-            window_enforce_even=window_enforce_even,
-            out_spatial_dim=out_spatial_dim,
-            out_dim=out_dim,
-        ),
-        out_spatial_dim,
-        out_dim,
-    )
+__all__ = ["mel_filterbank", "log_mel_filterbank_from_raw"]
 
 
 def mel_filterbank(
@@ -275,7 +191,7 @@ def log_mel_filterbank_from_raw(
         fft_length=n_fft,
     )
     power_spectrogram = rf.abs(spectrogram) ** 2.0
-    mel_fbank = rf.mel_filterbank(power_spectrogram, in_dim=in_dim_, out_dim=out_dim, sampling_rate=sampling_rate)
+    mel_fbank = mel_filterbank(power_spectrogram, in_dim=in_dim_, out_dim=out_dim, sampling_rate=sampling_rate)
     log_mel_fbank = rf.safe_log(mel_fbank, eps=1e-10)
     if log_base != math.e:
         log_mel_fbank = log_mel_fbank * (1.0 / math.log(log_base))

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/returnn/frontend/run_ctx.py

@@ -29,20 +29,20 @@ def reset_run_ctx():
     _run_ctx = None
 
 
-def init_train_step_run_ctx(*, train_flag: Union[bool, Tensor]):
+def init_train_step_run_ctx(*, train_flag: Union[bool, Tensor], step: Union[int, Tensor]):
     """
     Call this at the beginning of a new train step.
     """
     global _run_ctx
-    _run_ctx = RunCtx(stage="train_step", train_flag=train_flag)
+    _run_ctx = RunCtx(stage="train_step", train_flag=train_flag, step=step)
 
 
-def init_forward_step_run_ctx(expected_outputs: Optional[TensorDict] = None):
+def init_forward_step_run_ctx(*, expected_outputs: Optional[TensorDict] = None, step: Union[int, Tensor]):
     """
     Call this at the beginning of a new forward step.
     """
     global _run_ctx
-    _run_ctx = RunCtx(stage="forward_step", expected_outputs=expected_outputs)
+    _run_ctx = RunCtx(stage="forward_step", expected_outputs=expected_outputs, step=step)
 
 
 def get_run_ctx() -> RunCtx:
@@ -68,7 +68,12 @@ class RunCtx:
     """
 
     def __init__(
-        self, *, stage: str, train_flag: Union[bool, Tensor] = False, expected_outputs: Optional[TensorDict] = None
+        self,
+        *,
+        stage: str,
+        train_flag: Union[bool, Tensor] = False,
+        step: Union[int, Tensor] = 0,
+        expected_outputs: Optional[TensorDict] = None,
     ):
         """
         :param stage:
@@ -78,6 +83,7 @@ class RunCtx:
         """
         self._stage = stage
         self._train_flag = train_flag
+        self._step = step
         self.losses = {}  # type: Dict[str, Loss]
         self.outputs = TensorDict()
         self.expected_outputs = expected_outputs
@@ -98,6 +104,14 @@ class RunCtx:
         """
         return self._train_flag
 
+    @property
+    def step(self) -> Union[int, Tensor]:
+        """
+        :return: global train step, starting with 0, not reset after an epoch, i.e. ignoring the epochs.
+            In a graph-based backend, this can be dynamic.
+        """
+        return self._step
+
     def mark_as_loss(
         self,
         loss: Union[Tensor, Any],

returnn-1.20230816.2116/returnn/frontend/signal.py

@@ -0,0 +1,95 @@
+"""
+stft etc
+"""
+
+
+from __future__ import annotations
+from typing import Optional, Tuple
+from returnn.tensor import Tensor, Dim
+
+
+__all__ = ["stft"]
+
+
+def stft(
+    x: Tensor,
+    *,
+    in_spatial_dim: Dim,
+    frame_step: int,
+    frame_length: int,
+    fft_length: Optional[int] = None,
+    window_use_frame_length: bool = True,
+    align_window_left: bool = True,
+    window_enforce_even: bool = True,
+    out_spatial_dim: Optional[Dim] = None,
+    out_dim: Optional[Dim] = None,
+) -> Tuple[Tensor, Dim, Dim]:
+    """
+    Calculate the short-time Fourier transform (STFT) of a signal.
+
+    We currently always use the Hann window.
+
+    Note that there are inconsistencies between frameworks (e.g. PyTorch vs TensorFlow).
+    We have options where you can explicitly specify the behavior,
+    to match either PyTorch or TensorFlow.
+    See here for some discussion and demonstration:
+    https://github.com/pytorch/pytorch/issues/100177
+    https://github.com/librosa/librosa/issues/1701
+    https://github.com/albertz/playground/blob/master/tf_pt_stft.py
+
+    :param x: (..., in_spatial_dim, ...). any other dims are treated as batch dims.
+    :param in_spatial_dim:
+    :param frame_length: win_length in PT/librosa, frame_size in StftLayer, frame_length in TF
+    :param frame_step: hop_length in PT/librosa, frame_shift in StftLayer, frame_step in TF
+    :param fft_length: n_fft in PT/librosa, fft_size in StftLayer, fft_length in TF.
+        If not given, will use frame_length.
+    :param window_use_frame_length: Whether to use window size = frame_length.
+        If False, uses window size = fft_length.
+        This has only an effect if frame_length != fft_length.
+        This has an effect on the output seq length.
+        If True: out_seq_len = ⌈(T - frame_length + 1) / frame_step⌉ (T = in_seq_len).
+        If False: out_seq_len = ⌈(T - fft_length + 1) / frame_step⌉.
+        Note that in TF/SciPy, the behavior matches to window_use_frame_length=True,
+        but in PT/librosa, the behavior matches to window_use_frame_length=False.
+    :param align_window_left: whether to align the window to the left inside the potential wider fft_length window.
+        This only has an effect if frame_length != fft_length.
+        If True, the window will be aligned to the left and the right side will be padded with zeros.
+        Then effectively the remaining phases are not used.
+        If False, the window will be aligned to the center and the left and right side will be padded with zeros.
+        Then effectively the beginning and end phases are not used.
+        In TF/SciPy, the behavior matches to align_window_left=True,
+        but in PT/librosa, the behavior matches to align_window_left=False.
+    :param window_enforce_even: enforce that the window size is even, potentially lowering the window size by 1.
+        This only has an effect if the window size is uneven.
+        In TF, the behavior matches to window_enforce_even=True,
+        but in most other frameworks, the behavior matches to window_enforce_even=False.
+    :param out_spatial_dim:
+    :param out_dim:
+    :return: (stft, out_spatial_dim, out_dim)
+    """
+    fft_length = fft_length or frame_length
+    if out_dim is None:
+        out_dim = Dim(fft_length // 2 + 1, name="stft-freq")
+    if out_spatial_dim is None:
+        from .conv import make_conv_out_spatial_dims
+
+        (out_spatial_dim,) = make_conv_out_spatial_dims(
+            [in_spatial_dim], filter_size=frame_length, strides=frame_step, padding="valid"
+        )
+    # noinspection PyProtectedMember
+    return (
+        x._raw_backend.stft(
+            x,
+            in_spatial_dim=in_spatial_dim,
+            frame_step=frame_step,
+            frame_length=frame_length,
+            fft_length=fft_length,
+            window_use_frame_length=window_use_frame_length,
+            align_window_left=align_window_left,
+            window_enforce_even=window_enforce_even,
+            out_spatial_dim=out_spatial_dim,
+            out_dim=out_dim,
+        ),
+        out_spatial_dim,
+        out_dim,
+    )

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/returnn/tf/frontend_layers/config_entry_points.py

@@ -47,15 +47,19 @@ def get_net_dict(
     model = get_model_func(epoch=epoch, step=step)
 
     task = config.value("task", None)
+    step_tensor = rfl.make_layer({"class": "global_train_step"}, name="global_train_step")
     if task in {"train", "eval"}:
-        rf.init_train_step_run_ctx(train_flag=rfl.make_layer({"class": "train_flag"}, name="train_flag"))
+        rf.init_train_step_run_ctx(
+            train_flag=rfl.make_layer({"class": "train_flag"}, name="train_flag"),
+            step=step_tensor,
+        )
         train_step_func = get_global_config().typed_value("train_step")
         train_step_func(
             model=model,
             extern_data=extern_data,
         )
     elif task in {"forward", "search"}:
-        rf.init_forward_step_run_ctx()
+        rf.init_forward_step_run_ctx(step=step_tensor)
         forward_step_func = get_global_config().typed_value("forward_step")
         forward_step_func(
             model=model,

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/returnn/tf/frontend_layers/layer.py

@@ -429,9 +429,13 @@ class Layer:
         # However, this can always be done via the `name_scope` option,
         # if the name would not match otherwise.
         queue = [self.root]  # type: List[Layer]
+        visited: Set[Layer] = set()
         mod_in_layer = {}  # type: Dict[Tuple[Layer, RefIdEq[rf.Module]], Layer]
         while queue:
             ctx = queue.pop(0)
+            if ctx in visited:  # can happen when parents are reassigned
+                continue
+            visited.add(ctx)
 
             assert not ctx.parent or ctx.parent.children[ctx.name] is ctx
             for child in ctx.children.values():

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/returnn/torch/engine.py

@@ -459,10 +459,12 @@ class Engine(EngineBase):
         """
         if train_func:
             assert self._train_step_func is not None
-            rf.init_train_step_run_ctx(train_flag=train_flag)
+            rf.init_train_step_run_ctx(train_flag=train_flag, step=self.global_train_step)
         else:
             assert self._forward_step_func is not None, "define forward_step in the config"
-            rf.init_forward_step_run_ctx(expected_outputs=self._forward_step_expected_outputs)
+            rf.init_forward_step_run_ctx(
+                expected_outputs=self._forward_step_expected_outputs, step=self.global_train_step
+            )
 
         from returnn.torch.distributed import ddp_train_forward_ctx
 

{returnn-1.20230815.191535 → returnn-1.20230816.2116/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/returnn.egg-info/SOURCES.txt

@@ -186,6 +186,8 @@ returnn/frontend/signal.py
 returnn/frontend/state.py
 returnn/frontend/tensor_array.py
 returnn/frontend/types.py
+returnn/frontend/audio/__init__.py
+returnn/frontend/audio/mel.py
 returnn/frontend/encoder/__init__.py
 returnn/frontend/encoder/base.py
 returnn/frontend/encoder/conformer.py

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/tests/rf_utils.py

@@ -99,7 +99,7 @@ def run_model_torch(extern_data: TensorDict, get_model: rf.GetModelFunc, forward
     tensor_dict_numpy_to_torch_(extern_data)
 
     model = get_model(epoch=1, step=0)
-    rf.init_forward_step_run_ctx()
+    rf.init_forward_step_run_ctx(step=0)
     forward_step(model=model, extern_data=extern_data)
     outputs = rf.get_run_ctx().outputs
     assert outputs.data

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/tests/test_rf_rec.py

@@ -285,7 +285,7 @@ def test_zoneout_lstm_tf_layers_vs_rf_pt():
         new_model.lstm.bias.raw_tensor.copy_(torch.from_numpy(tf_lstm_bias))
 
     print("*** Forward")
-    rf.init_train_step_run_ctx(train_flag=False)
+    rf.init_train_step_run_ctx(train_flag=False, step=0)
     extern_data.reset_content()
     extern_data.assign_from_raw_tensor_dict_(extern_data_numpy_raw_dict)
     tensor_dict_numpy_to_torch_(extern_data)

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/tests/test_rf_signal.py

@@ -46,7 +46,7 @@ def test_mel_filterbank():
                 source, in_spatial_dim=in_spatial_dim, frame_step=3, frame_length=5, fft_length=6
             )
             source = rf.abs(source) ** 2.0
-            source = rf.mel_filterbank(source, in_dim=in_dim_, out_dim=feat_dim, sampling_rate=16)
+            source = rf.audio.mel_filterbank(source, in_dim=in_dim_, out_dim=feat_dim, sampling_rate=16)
             return source, in_spatial_dim
 
     # noinspection PyShadowingNames

{returnn-1.20230815.191535 → returnn-1.20230816.2116}/tools/torch_export_to_onnx.py

@@ -192,13 +192,14 @@ def main():
     ), "The specified config needs to have explicit model outputs. Please define `model_outputs` in your config."
     model_outputs = TensorDict()
     model_outputs.update(model_outputs_dict, auto_convert=True)
-    rf.init_forward_step_run_ctx(expected_outputs=model_outputs)
-    rf.set_random_seed(42)
 
     loaded_checkpoint = torch.load(args.checkpoint, map_location=torch.device(args.device))
     epoch = loaded_checkpoint["epoch"]
     step = loaded_checkpoint["step"]
 
+    rf.init_forward_step_run_ctx(expected_outputs=model_outputs, step=step)
+    rf.set_random_seed(42)
+
     get_model_func = config.typed_value("get_model")
     assert get_model_func, "get_model() isn't specified in the config passed as a parameter."
     sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}

returnn-1.20230815.191535/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230815.191535'
-long_version = '1.20230815.191535+git.2a78f6e'