returnn 1.20240919.112211__tar.gz → 1.20240919.155206__tar.gz

{returnn-1.20240919.112211 → returnn-1.20240919.155206}/PKG-INFO

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

returnn-1.20240919.155206/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240919.155206'
+long_version = '1.20240919.155206+git.0947255'

{returnn-1.20240919.112211 → returnn-1.20240919.155206}/returnn/frontend/decoder/transformer.py

@@ -1,6 +1,8 @@
 """
 (Label-sync) Transformer decoder, optionally including cross attention to encoder
 
+Also see :mod:`returnn.frontend.encoder.transformer`.
+
 References:
 
     (Original paper of course)
@@ -25,7 +27,7 @@ from returnn.tensor import Tensor, Dim, single_step_dim
 
 class TransformerDecoder(rf.Module):
     """
-    Represents Transformer decoder architecture
+    Represents the Transformer decoder architecture
     """
 
     def __init__(
@@ -259,7 +261,7 @@ class TransformerDecoderLayer(rf.Module):
         :param ff_activation: activation function for feed-forward network
         :param dropout: the dropout value for the FF block
         :param num_heads: the number of attention heads
-        :param self_att: the self-attention layer. RelPosSelfAttention originally and default
+        :param self_att: the self-attention layer. CausalSelfAttention originally and default
         :param self_att_opts: options for the self-attention layer, for :class:`nn.RelPosSelfAttention`
         :param att_dropout: attention dropout value
         :param norm: pre-normalization for FF and attention blocks

returnn-1.20240919.155206/returnn/frontend/encoder/transformer.py

@@ -0,0 +1,253 @@
+"""
+Transformer encoder
+
+Also see :mod:`returnn.frontend.decoder.transformer`.
+"""
+
+from __future__ import annotations
+from typing import Optional, Union, Any, Callable, Sequence, Dict
+import functools
+import copy as _copy
+from types import FunctionType
+from returnn.util.basic import NotSpecified
+from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
+from returnn.frontend.decoder.transformer import make_norm, FeedForward
+
+
+class TransformerEncoder(rf.Module):
+    """
+    Represents the Transformer encoder architecture
+    """
+
+    def __init__(
+        self,
+        vocab_dim: Dim,
+        model_dim: Union[Dim, int] = Dim(512, name="transformer-enc-default-model-dim"),
+        *,
+        num_layers: int,
+        ff: Union[type, Dict[str, Any], rf.Module] = NotSpecified,
+        pos_enc: Union[None, Callable, Dict[str, Any], rf.Module] = rf.sinusoidal_positional_encoding,
+        dropout: float = 0.1,
+        num_heads: int = 8,
+        att_dropout: float = 0.1,
+        norm: Union[type, Dict[str, Any], rf.Module, Callable] = rf.LayerNorm,
+        decoder_layer: Optional[Union[TransformerEncoderLayer, rf.Module, type, Any]] = None,
+        embed_dim: Optional[Dim] = None,
+        input_embedding_scale: float = None,
+        input_dropout: float = None,
+        sequential=rf.Sequential,
+    ):
+        """
+        :param vocab_dim:
+        :param model_dim: the output feature dimension
+        :param num_layers: the number of encoder layers
+        :param ff: feed-forward / MLP block. Default is :class:`FeedForward`
+        :param pos_enc: positional encoding. Default is sinusoidal positional encoding.
+        :param dropout: the dropout value for the FF block
+        :param num_heads: the number of attention heads
+        :param att_dropout: attention dropout value
+        :param norm: pre-normalization for FF and attention blocks
+        :param decoder_layer: an instance of :class:`TransformerDecoderLayer` or similar
+        :param embed_dim: if given, will first have an embedding [vocab,embed] and then a linear [embed,model].
+        :param input_embedding_scale:
+        :param input_dropout:
+        :param sequential:
+        """
+        super().__init__()
+
+        if not isinstance(vocab_dim, Dim):
+            raise TypeError(f"TransformerDecoder: unexpected vocab_dim {vocab_dim!r} type {type(vocab_dim)}")
+        if isinstance(model_dim, int):
+            model_dim = Dim(model_dim, name="transformer-dec-model-dim")
+        if not isinstance(model_dim, Dim):
+            raise TypeError(f"TransformerDecoder: unexpected model_dim {model_dim!r} type {type(model_dim)}")
+
+        self.vocab_dim = vocab_dim
+        self.model_dim = model_dim
+        self.embed_dim = embed_dim
+
+        # We could make this optional or configurable if we ever need to.
+        # Or maybe you would just have another separate implementation of this module then...
+        self.input_embedding = rf.Embedding(vocab_dim, embed_dim or model_dim)
+
+        self.input_embedding_proj = None
+        if embed_dim:
+            self.input_embedding_proj = rf.Linear(embed_dim, model_dim, with_bias=False)
+
+        if pos_enc is None:
+            pass
+        elif isinstance(pos_enc, dict):
+            pos_enc = rf.build_from_dict(
+                pos_enc, feat_dim=embed_dim or model_dim, dtype=self.input_embedding.weight.dtype
+            )
+        elif isinstance(pos_enc, rf.Module):
+            pass
+        elif isinstance(pos_enc, FunctionType):
+            pos_enc = functools.partial(
+                pos_enc, feat_dim=embed_dim or model_dim, dtype=self.input_embedding.weight.dtype
+            )
+        else:
+            raise TypeError(f"unexpected pos_enc type {pos_enc!r}")
+        self.pos_enc = pos_enc
+        if input_embedding_scale is None:
+            input_embedding_scale = model_dim.dimension**0.5
+        self.input_embedding_scale = input_embedding_scale
+        if input_dropout is None:
+            input_dropout = dropout
+        self.input_dropout = input_dropout
+
+        if not decoder_layer or isinstance(decoder_layer, type):
+            decoder_layer_opts_ = dict(
+                out_dim=model_dim,
+                ff=ff,
+                dropout=dropout,
+                num_heads=num_heads,
+                att_dropout=att_dropout,
+                norm=norm,
+            )
+            if not decoder_layer:
+                decoder_layer = TransformerEncoderLayer(**decoder_layer_opts_)
+            elif isinstance(decoder_layer, type):
+                decoder_layer = decoder_layer(**decoder_layer_opts_)
+            else:
+                raise TypeError(f"unexpected decoder_layer {decoder_layer!r}")
+
+        self.layers = sequential(_copy.deepcopy(decoder_layer) for _ in range(num_layers))
+
+        self.final_layer_norm = make_norm(norm, model_dim)
+
+    def default_initial_state(self, *, batch_dims: Sequence[Dim]) -> rf.State:
+        """default initial state"""
+        state = rf.State({k: v.default_initial_state(batch_dims=batch_dims) for k, v in self.layers.items()})
+        state.pos = rf.zeros((), dtype="int32", device="cpu")
+        return state
+
+    def transform_encoder(self, encoder: Tensor, *, axis: Dim) -> rf.State:
+        """
+        Transform encoder output.
+        Note that the Transformer decoder usually expects that layer-norm was applied already on the encoder output.
+        """
+        return rf.State({k: v.transform_encoder(encoder, axis=axis) for k, v in self.layers.items()})
+
+    def __call__(
+        self, source: Tensor, *, spatial_dim: Dim, collected_outputs: Optional[Dict[str, Tensor]] = None
+    ) -> Tensor:
+        """
+        forward, single step or whole sequence.
+
+        :param source: labels
+        :param spatial_dim: single_step_dim or spatial dim of source
+        :param collected_outputs:
+        :return: final encoder output, after final layer norm
+        """
+        decoded = self.input_embedding(source) * self.input_embedding_scale
+        if self.pos_enc is not None:
+            decoded = decoded + self.pos_enc(spatial_dim=spatial_dim)
+        decoded = rf.dropout(decoded, self.input_dropout)
+        if self.input_embedding_proj is not None:
+            decoded = self.input_embedding_proj(decoded)
+
+        for layer_name, layer in self.layers.items():
+            layer: TransformerEncoderLayer  # or similar
+            decoded = layer(decoded, spatial_dim=spatial_dim)
+            if collected_outputs is not None:
+                collected_outputs[layer_name] = decoded
+
+        decoded = self.final_layer_norm(decoded)
+        return decoded
+
+
+class TransformerEncoderLayer(rf.Module):
+    """
+    Represents a Transformer encoder block
+    """
+
+    def __init__(
+        self,
+        out_dim: Dim = Dim(512, name="transformer-enc-default-out-dim"),
+        *,
+        ff: Union[type, Dict[str, Any], rf.Module] = NotSpecified,
+        dropout: float = 0.1,
+        num_heads: int = 8,
+        self_att: Optional[
+            Union[rf.CausalSelfAttention, rf.RelPosCausalSelfAttention, rf.Module, type, Dict[str, Any]]
+        ] = None,
+        att_dropout: float = 0.1,
+        norm: Union[type, Dict[str, Any], rf.Module, Callable] = rf.LayerNorm,
+    ):
+        """
+        :param out_dim: the output feature dimension
+        :param ff: feed-forward / MLP block. Default is :class:`FeedForward`
+        :param dropout: the dropout value for the FF block
+        :param num_heads: the number of attention heads
+        :param self_att: the self-attention layer. SelfAttention originally and default
+        :param att_dropout: attention dropout value
+        :param norm: pre-normalization for FF and attention blocks
+        """
+        super().__init__()
+
+        self.dropout = dropout
+        self.dropout_broadcast = rf.dropout_broadcast_default()
+        self.out_dim = out_dim
+
+        if ff is NotSpecified:
+            ff = FeedForward
+        if isinstance(ff, rf.Module):
+            ff = _copy.deepcopy(ff)
+        else:
+            ff_kwargs = dict(out_dim=out_dim, dropout=dropout)
+            ff_kwargs = {k: v for (k, v) in ff_kwargs.items() if v is not NotSpecified}
+            if isinstance(ff, type):
+                ff = ff(**ff_kwargs)
+            elif isinstance(ff, dict):
+                ff = rf.build_from_dict(ff, **ff_kwargs)
+            else:
+                raise TypeError(f"unexpected ff type {ff!r}")
+        assert isinstance(ff, rf.Module)
+
+        self.ff = ff
+        self.ff_layer_norm = make_norm(norm, out_dim)
+
+        if self_att is None or isinstance(self_att, type) or isinstance(self_att, dict):
+            self_att_opts_ = dict(
+                in_dim=out_dim,
+                proj_dim=out_dim,
+                key_dim_total=out_dim,
+                value_dim_total=out_dim,
+                num_heads=num_heads,
+                att_dropout=att_dropout,
+            )
+            if self_att is None:
+                self.self_att = rf.SelfAttention(**self_att_opts_)
+            elif isinstance(self_att, type):
+                self.self_att = self_att(**self_att_opts_)
+            elif isinstance(self_att, dict):
+                self.self_att = rf.build_from_dict(self_att, **self_att_opts_)
+            else:
+                raise TypeError(f"unexpected self_att type {self_att!r}")
+        elif isinstance(self_att, rf.Module):
+            self.self_att = _copy.deepcopy(self_att)
+        else:
+            raise TypeError(f"unexpected self_att type {self_att!r}")
+        self.self_att_layer_norm = make_norm(norm, out_dim)
+
+    def default_initial_state(self, *, batch_dims: Sequence[Dim]) -> rf.State:
+        """default initial state"""
+        return rf.State(self_att=self.self_att.default_initial_state(batch_dims=batch_dims))
+
+    def __call__(self, x: Tensor, *, spatial_dim: Dim) -> Tensor:
+        """forward"""
+        # (multi-head) self-attention (MHSA or simply SA)
+        x_sa_ln = self.self_att_layer_norm(x)
+        x_sa = self.self_att(x_sa_ln, axis=spatial_dim)
+        x_sa = rf.dropout(x_sa, self.dropout, axis=self.dropout_broadcast and self.out_dim)
+        x = x_sa + x
+
+        # feed-forward (FF)
+        x_ff_ln = self.ff_layer_norm(x)
+        x_ff = self.ff(x_ff_ln)
+        x_ff = rf.dropout(x_ff, self.dropout, axis=self.dropout_broadcast and self.out_dim)
+        x = x_ff + x
+
+        return x

{returnn-1.20240919.112211 → returnn-1.20240919.155206}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240919.112211 → returnn-1.20240919.155206}/returnn.egg-info/SOURCES.txt

@@ -219,6 +219,7 @@ returnn/frontend/encoder/__init__.py
 returnn/frontend/encoder/base.py
 returnn/frontend/encoder/conformer.py
 returnn/frontend/encoder/e_branchformer.py
+returnn/frontend/encoder/transformer.py
 returnn/import_/__init__.py
 returnn/import_/common.py
 returnn/import_/git.py

returnn-1.20240919.112211/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240919.112211'
-long_version = '1.20240919.112211+git.b4c8c79'