returnn 1.20240829.174139__tar.gz → 1.20240830.140746__tar.gz

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/PKG-INFO

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

returnn-1.20240830.140746/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240830.140746'
+long_version = '1.20240830.140746+git.d8709ff'

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/returnn/datasets/postprocessing.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 
 from itertools import islice
 from numpy.random import RandomState
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar
 
 from returnn.datasets.basic import DatasetSeq
 from returnn.datasets.util.vocabulary import Vocabulary
@@ -15,7 +15,7 @@ from returnn.tensor.dim import Dim
 from .basic import init_dataset
 from .cached2 import CachedDataset2
 
-__all__ = ["PostprocessingDataset"]
+__all__ = ["PostprocessingDataset", "LaplaceOrdering", "Sequential"]
 
 
 class PostprocessingDataset(CachedDataset2):
@@ -54,6 +54,30 @@ class PostprocessingDataset(CachedDataset2):
                 "data": {"dims": [time_dim, new_data_dim]},
             },
         }
+
+    The dataset itself does not support its own seq ordering and relies on the wrapped
+    dataset for seq ordering instead. Specifying a ``seq_ordering`` other than ``default``
+    results in an error.
+
+    However, we provide an iterator that implements the common `laplace:.NUM_SEQS_PER_BIN`-variant
+    of seq ordering that any custom ``map_seq_stream``-style postprocessing iterator can be composed
+    with to implement the ordering via :class:`LaplaceOrdering`.
+
+    Like this::
+
+        from returnn.datasets.postprocessing import LaplaceOrdering, Sequential
+
+        def my_map_seq_stream(iterator):
+            ...
+
+        train = {
+            "class": "PostprocessingDataset",
+            # ...
+            "map_seq_stream": Sequential(
+                my_map_seq_stream,
+                LaplaceOrdering(num_seqs_per_bin=1000),
+            ),
+        }
     """
 
     def __init__(
@@ -233,6 +257,8 @@ class LaplaceOrdering(Callable[[Iterator[TensorDict]], Iterator[TensorDict]]):
     """
     Iterator compatible with :class:`PostprocessingDataset`'s ``map_seq_stream`` applying
     laplace sequence ordering based on the number of segments per bin.
+
+    To be composed with any custom data postprocessing logic via :class:`Sequential`.
     """
 
     def __init__(self, num_seqs_per_bin: int, length_key: str = "data"):
@@ -264,3 +290,29 @@ class LaplaceOrdering(Callable[[Iterator[TensorDict]], Iterator[TensorDict]]):
         :return: segment length of the segment in `tdict` as measured by `self.length_key` for comparison.
         """
         return tdict.data[self.length_key].raw_tensor.shape[0]
+
+
+T = TypeVar("T", TensorDict, Iterator[TensorDict])
+
+
+class Sequential:
+    """
+    Callable that composes multiple postprocessing functions into one by sequential application,
+    i.e. Sequential(f, g)(x) = (g ∘ f)(x) = g(f(x)).
+
+    Can either compose ``map_seq``-style single-segment processor functions or ``map_seq_stream``-style
+    iterators operating on multiple segments. Just make sure not to mix both styles.
+    """
+
+    def __init__(self, *postprocessing_funcs: Callable):
+        """
+        :param postprocessing_funcs: Postprocessing functions to compose.
+        """
+        self.funcs = postprocessing_funcs
+
+    def __call__(self, arg: T, **kwargs) -> T:
+        """:return: result of sequential application of the postprocessing functions"""
+
+        for func in self.funcs:
+            arg = func(arg, **kwargs)
+        return arg

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/returnn/frontend/attention.py

@@ -16,6 +16,7 @@ __all__ = [
     "SelfAttention",
     "CausalSelfAttention",
     "CausalSelfAttentionState",
+    "RotaryPosSelfAttention",
     "RotaryPosCausalSelfAttention",
     "RelPosSelfAttention",
     "RelPosCausalSelfAttention",
@@ -264,6 +265,36 @@ class CausalSelfAttentionState(rf.State):
             self.accum_axis = accum_axis
 
 
+class RotaryPosSelfAttention(SelfAttention):
+    """
+    Rotary positional encoding (RoPE)-based self attention
+    """
+
+    def __call__(self, source: Tensor, *, axis: Dim) -> Tensor:
+        """forward"""
+        q, k, v = self.forward_qkv(source)
+
+        # Apply RoPE using sinusoidal positional encoding.
+        # Note: base is a bit different in rf.sinusoidal_positional_encoding (like the original)
+        # vs how it's commonly used for RoPE.
+        # log(base) / (dim / 2 - 1) = log(10_000) * 2 / dim
+        # <=> log(base) = log(10_000) * (dim / 2 - 1) * 2 / dim = log(10_000) * (1 - 2 / dim)
+        # <=> base = 10_000 ** (1 - 2 / dim)
+        pos_enc = rf.sinusoidal_positional_encoding(
+            spatial_dim=axis,
+            feat_dim=self.key_dim_per_head,
+            base=10_000 ** (1 - 2 / self.key_dim_per_head.dimension),
+        )  # [T,D]
+        q = _apply_rope(q, pos_enc, self.key_dim_per_head)
+        k = _apply_rope(k, pos_enc, self.key_dim_per_head)
+
+        kv_axis = Dim(None, name=f"{axis.name}-kv")
+        k, _ = rf.replace_dim(k, in_dim=axis, out_dim=kv_axis)
+        v, _ = rf.replace_dim(v, in_dim=axis, out_dim=kv_axis)
+        output = self.attention(q, k, v, kv_axis=kv_axis)
+        return output
+
+
 class RotaryPosCausalSelfAttention(CausalSelfAttention):
     """
     Rotary positional encoding (RoPE)-based causal self attention

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/returnn/frontend/container.py

@@ -66,14 +66,17 @@ class ModuleList(rf.Module, Generic[__ModT]):
         """module items"""
         return self._get_modules().items()
 
-    def __getitem__(self, idx: Union[slice, int]) -> Union[ModuleList[__ModT], __ModT]:
+    def __getitem__(self, idx: Union[slice, int, str]) -> Union[ModuleList[__ModT], __ModT]:
         if isinstance(idx, slice):
             return self.__class__(dict(list(self._get_modules().items())[idx]))
-        else:
+        elif isinstance(idx, str):
+            return getattr(self, idx)
+        elif isinstance(idx, int):
             return list(self._get_modules().values())[idx]
+        else:
+            raise TypeError(f"{self} __getitem__ Invalid idx type {type(idx).__name__}")
 
     def __setitem__(self, idx: Union[slice, int], module: Union[__ModT, Iterable[__ModT]]) -> None:
-        key = list(self._get_modules().keys())[idx]
         if isinstance(idx, slice):
             assert not idx.step or idx.step == 1  # not supported
             mod_items = list(self._get_modules().items())
@@ -95,8 +98,13 @@ class ModuleList(rf.Module, Generic[__ModT]):
                 assert not hasattr(self, k)
                 setattr(self, k, v)
                 i += 1
-        else:
+        elif isinstance(idx, str):
+            setattr(self, idx, module)
+        elif isinstance(idx, int):
+            key = list(self._get_modules().keys())[idx]
             setattr(self, key, _convert_to_module(module))
+        else:
+            raise TypeError(f"{self} __setitem___ Invalid idx type {type(idx).__name__}")
 
     def __delitem__(self, key: Union[slice, int]):
         if isinstance(key, slice):

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/returnn/frontend/encoder/conformer.py

@@ -295,6 +295,7 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         *,
         num_layers: int,
         input_layer: Optional[Union[ConformerConvSubsample, ISeqDownsamplingEncoder, rf.Module, Any]],
+        input_embedding_scale: float = 1.0,
         input_dropout: float = 0.1,
         ff_dim: Dim = NotSpecified,
         ff_activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module] = NotSpecified,
@@ -312,6 +313,8 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         :param num_layers: the number of encoder layers
         :param input_layer: input/frontend/prenet with potential subsampling.
             (x, in_spatial_dim) -> (y, out_spatial_dim)
+        :param input_embedding_scale: applied after input_layer. 1.0 by default for historic reasons.
+            In std Transformer, also ESPnet E-Branchformer and Conformer, this is sqrt(out_dim).
         :param input_dropout: applied after input_projection(input_layer(x))
         :param ff_dim: the dimension of feed-forward layers. 2048 originally, or 4 times out_dim
         :param ff_activation: activation function for feed-forward network
@@ -338,6 +341,7 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
             if input_layer
             else None
         )
+        self.input_embedding_scale = input_embedding_scale
         self.input_dropout = input_dropout
 
         if not encoder_layer or isinstance(encoder_layer, (dict, type)):
@@ -389,6 +393,8 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         else:
             x_subsample, out_spatial_dim = source, in_spatial_dim
         x = self.input_projection(x_subsample) if self.input_projection else x_subsample
+        if self.input_embedding_scale != 1.0:
+            x = x * self.input_embedding_scale
         x = rf.dropout(x, self.input_dropout, axis=self.dropout_broadcast and self.out_dim)
         x = self.layers(x, spatial_dim=out_spatial_dim, collected_outputs=collected_outputs)
         return x, out_spatial_dim

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/returnn/frontend/encoder/e_branchformer.py

@@ -173,7 +173,13 @@ class FeedForwardConvGated(rf.Module):
     ):
         """
         :param out_dim: the encoder (e.g. E-Branchformer) model dim. (usually 256 or 512)
-        :param ff_dim: intermediate dimension. (usually 2048 or 3072, not necessarily factor 4 or 4*2/3)
+        :param ff_dim: intermediate dimension.
+            This is like cgmlp_linear_units/2 in ESPnet.
+            Note the 1/2 factor, which is because in ESPnet, you specify the total dimension,
+            before it is split for the gating,
+            while here, you specify the dimension for the gating part.
+            Common settings are 2048/2 or 3072/2.
+            In the paper, they mention a factor of 3 of the model dimension (factor 6 for ESPnet setting).
         :param kernel_size: for the depthwise convolution (usually 31)
         :param dropout:
         :param activation: activation function after the first linear layer, for both parts.
@@ -190,7 +196,7 @@ class FeedForwardConvGated(rf.Module):
         super().__init__()
 
         if ff_dim is NotSpecified:
-            ff_dim = out_dim * 6  # somewhat arbitrary. with 512, this is 3072.
+            ff_dim = out_dim * 3  # somewhat arbitrary. with 512, this is 3072/2.
         if isinstance(ff_dim, int):
             ff_dim = Dim(ff_dim, name="e-branchformer-ff-dim")
         if not isinstance(ff_dim, Dim):

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240829.174139 → returnn-1.20240830.140746}/tests/test_Dataset.py

@@ -1111,6 +1111,12 @@ def test_PostprocessingDataset():
             assert prev_len is None or classes.shape[0] <= prev_len or i == 3
             prev_len = classes.shape[0]
 
+    # test composition
+    from returnn.datasets.postprocessing import Sequential
+
+    func = Sequential(lambda x: x * 10, lambda y: y + 1)
+    assert func(2) == 21
+
 
 if __name__ == "__main__":
     better_exchook.install()

returnn-1.20240829.174139/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240829.174139'
-long_version = '1.20240829.174139+git.a19cf76'