returnn 1.20231215.23621__tar.gz → 1.20231219.120347__tar.gz

{returnn-1.20231215.23621/returnn.egg-info → returnn-1.20231219.120347}/PKG-INFO

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

returnn-1.20231219.120347/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20231219.120347'
+long_version = '1.20231219.120347+git.d0d59e6'

{returnn-1.20231215.23621 → returnn-1.20231219.120347}/returnn/frontend/attention.py

@@ -7,40 +7,23 @@ from __future__ import annotations
 from typing import Tuple, Union, Optional, Sequence
 import weakref
 import logging
-from returnn.util.py_compat import Protocol
 from returnn.tensor import Tensor, Dim, single_step_dim
 import returnn.frontend as rf
 
 
 __all__ = [
-    "AttentionFunc",
     "dot_attention",
     "SelfAttentionBase",
     "SelfAttention",
     "CausalSelfAttention",
     "CausalSelfAttentionState",
     "RelPosSelfAttention",
+    "CrossAttention",
     "LearnedRelativePositionalEncoding",
     "relative_positional_encoding",
 ]
 
 
-class AttentionFunc(Protocol):
-    """Protocol defining a generic attention function"""
-
-    def __call__(
-        self,
-        query: Tensor,
-        keys: Tensor,
-        values: Tensor,
-        *,
-        key_dim: Dim,
-        axis: Dim,
-        att_dropout: float = 0.1,
-    ):
-        ...
-
-
 def dot_attention(
     query: Tensor,
     keys: Tensor,
@@ -106,13 +89,13 @@ class SelfAttentionBase(rf.Module):
         """
         :param in_dim: input dim
         :param proj_dim: if given, will add a final linear projection to this dim.
-          otherwise no projection after the attention
+            otherwise no projection after the attention
         :param key_dim_total: total key dim. should be a multiple of num_heads
         :param value_dim_total: total value dim. should be a multiple of num_heads
         :param num_heads: number of heads
         :param with_bias: whether to add bias to qkv and proj linear projections.
-          Was False in original Transformer, but many recent implementations use True by default.
-          Also see: https://github.com/rwth-i6/returnn_common/issues/234.
+            Was False in original Transformer, but many recent implementations use True by default.
+            Also see: https://github.com/rwth-i6/returnn_common/issues/234.
         :param att_dropout: dropout for attention weights
         :param att_dropout_broadcast: whether to broadcast over all but ``axis``.
             normally not wanted. disabled by default since behavior version 19.
@@ -414,6 +397,133 @@ class RelPosSelfAttention(SelfAttentionBase):
         return x_padded
 
 
+class CrossAttention(rf.Module):
+    """
+    Cross attention
+
+    It uses :func:`dot_attention` for multi-headed dot-attention.
+    """
+
+    def __init__(
+        self,
+        encoder_dim: Dim,
+        query_in_dim: Dim,
+        proj_dim: Optional[Dim],
+        *,
+        key_dim_total: Dim,
+        value_dim_total: Dim,
+        num_heads: Union[int, Dim],
+        with_bias: bool = True,
+        att_dropout: float = 0.1,
+        att_dropout_broadcast: Optional[bool] = None,
+    ):
+        """
+        :param encoder_dim: encoder output dim = input dim for key-value
+        :param query_in_dim: input dim for query
+        :param proj_dim: if given, will add a final linear projection to this dim.
+            otherwise no projection after the attention
+        :param key_dim_total: total key dim. should be a multiple of num_heads
+        :param value_dim_total: total value dim. should be a multiple of num_heads
+        :param num_heads: number of heads
+        :param with_bias: whether to add bias to qkv and proj linear projections.
+            Was False in original Transformer, but many recent implementations use True by default.
+            Also see: https://github.com/rwth-i6/returnn_common/issues/234.
+        :param att_dropout: dropout for attention weights
+        :param att_dropout_broadcast: whether to broadcast over all but ``axis``.
+            normally not wanted. disabled by default since behavior version 19.
+        """
+        super().__init__()
+        self.encoder_dim = encoder_dim
+        self.query_in_dim = query_in_dim
+        self.out_dim = proj_dim if proj_dim else value_dim_total
+        if isinstance(num_heads, int):
+            num_heads = Dim(num_heads, name="num_heads")
+        self.key_dim_total = key_dim_total
+        self.key_dim_per_head = key_dim_total.div_left(num_heads)
+        self.value_dim_total = value_dim_total
+        self.value_dim_per_head = value_dim_total.div_left(num_heads)
+        self.num_heads = num_heads
+        self.kv_dim_total = key_dim_total + value_dim_total
+        self.kv_dim_per_head = self.key_dim_per_head + self.value_dim_per_head
+        self.kv = rf.Linear(encoder_dim, self.kv_dim_total, with_bias=with_bias)
+        self.q = rf.Linear(query_in_dim, self.key_dim_total, with_bias=with_bias)
+        # In Fairseq MultiheadAttention, they use:
+        #   nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))  (same for q_proj, v_proj),
+        # where (Xavier Glorot) xavier_uniform_ means:
+        #   std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
+        #   a = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
+        #   _no_grad_uniform_(tensor, -a, a)
+        # Out nn.init.VarianceScaling with mode="fan_avg", distribution="uniform":
+        #   scale = scale * 2.0 / float(fan_in + fan_out)
+        #   limit = math.sqrt(3.0 * scale)
+        #   nn.random(distribution="uniform", minval=-limit, maxval=limit, ...)
+        # Xavier Glorot: VarianceScaling with mode="fan_avg", distribution="uniform", scale=1.0.
+        self.kv.weight.initial = rf.init.Glorot(scale=3 / 4)
+        self.q.weight.initial = rf.init.Glorot(scale=1 / 2)
+        # The bias init is different, but not sure how important this is.
+        if proj_dim:
+            self.proj = rf.Linear(value_dim_total, proj_dim, with_bias=with_bias)
+        else:
+            self.proj = None
+        self.att_dropout = att_dropout
+        if att_dropout_broadcast is None:
+            att_dropout_broadcast = _att_dropout_broadcast_default()
+        self.att_dropout_broadcast = att_dropout_broadcast
+
+    def transform_encoder(self, encoder: Tensor, *, axis: Dim) -> rf.State:
+        """
+        Transformer encoder output. This is intended as an initial API suggestion.
+        """
+        k, v = self.forward_kv(encoder)
+        return rf.State(k=k, v=v, kv_axis=axis)
+
+    def forward_kv(self, source: Tensor) -> Tuple[Tensor, Tensor]:
+        """
+        This would be calculated once for the whole sequence (batch)
+        and then always reused for :func:`attention`.
+
+        :return: k,v
+        """
+        qkv = self.kv(source)
+        qkv = rf.split_dims(qkv, axis=self.kv_dim_total, dims=(self.num_heads, self.kv_dim_per_head))
+        k, v = rf.split(
+            qkv,
+            axis=self.kv_dim_per_head,
+            out_dims=(self.key_dim_per_head, self.value_dim_per_head),
+        )
+        return k, v
+
+    def forward_query(self, source: Tensor) -> Tensor:
+        """
+        This is calculated for every different query.
+
+        :return: q
+        """
+        q = self.q(source)
+        q = rf.split_dims(q, axis=self.key_dim_total, dims=(self.num_heads, self.key_dim_per_head))
+        return q
+
+    def __call__(self, q: Tensor, encoder: rf.State) -> Tensor:
+        q = self.forward_query(q)
+        return self.attention(q=q, **encoder)
+
+    def attention(self, q: Tensor, k: Tensor, v: Tensor, *, kv_axis: Dim) -> Tensor:
+        """apply attention"""
+        att = dot_attention(
+            q,
+            k,
+            v,
+            key_dim=self.key_dim_per_head,
+            axis=kv_axis,
+            att_dropout=self.att_dropout,
+            att_dropout_broadcast=self.att_dropout_broadcast,
+        )
+        output, _ = rf.merge_dims(att, dims=(self.num_heads, self.value_dim_per_head), out_dim=self.value_dim_total)
+        if self.proj:
+            output = self.proj(output)
+        return output
+
+
 class LearnedRelativePositionalEncoding(rf.Module):
     """
     Learnable relative positional encoding.

{returnn-1.20231215.23621 → returnn-1.20231219.120347}/returnn/torch/engine.py

@@ -116,6 +116,7 @@ class Engine(EngineBase):
             torch.cuda.set_device(self._device)
 
         self._log_memory_usage = config.bool("torch_log_memory_usage", False)
+        self._reset_dev_memory_caches = config.bool("reset_dev_memory_caches", False)
 
         amp_options = self.config.opt_typed_value("torch_amp")
         grad_scaler_opts = self.config.typed_value("grad_scaler", NotSpecified)
@@ -248,6 +249,12 @@ class Engine(EngineBase):
             }
         )
 
+    def _maybe_reset_dev_memory_caches(self):
+        if not self._reset_dev_memory_caches:
+            return
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def _reset_dev_memory_stats(self):
         dev = torch.device(self._device)
         if dev.type == "cuda":
@@ -291,6 +298,7 @@ class Engine(EngineBase):
         elapsed_computation_time = 0
 
         self._pt_model.train()
+        self._maybe_reset_dev_memory_caches()
         self._reset_dev_memory_stats()
 
         if self.config.bool("debug_shell_before_train_loop", False):
@@ -413,6 +421,7 @@ class Engine(EngineBase):
         Runs model on all eval datasets and calculates the loss.
         """
         self._pt_model.eval()
+        self._maybe_reset_dev_memory_caches()
         self._reset_dev_memory_stats()
 
         eval_dump_str = []
@@ -807,6 +816,7 @@ class Engine(EngineBase):
         elapsed_computation_time = 0.0
 
         self._pt_model.eval()
+        self._maybe_reset_dev_memory_caches()
         self._reset_dev_memory_stats()
 
         if dataset.supports_seq_order_sorting():

{returnn-1.20231215.23621 → returnn-1.20231219.120347}/returnn/util/debug.py

@@ -2,6 +2,8 @@
 Some generic debugging utilities.
 """
 
+from __future__ import annotations
+from typing import Optional, Any, Dict
 import os
 import sys
 import signal
@@ -497,15 +499,19 @@ def init_cuda_not_in_main_proc_check():
     cuda.use.device_number = None
 
 
-def debug_shell(user_ns=None, user_global_ns=None, exit_afterwards=True):
+def debug_shell(
+    user_ns: Optional[Dict[str, Any]] = None,
+    user_global_ns: Optional[Dict[str, Any]] = None,
+    exit_afterwards: bool = True,
+):
     """
     Provides some interactive Python shell.
     Uses IPython if possible.
     Wraps to ``better_exchook.debug_shell``.
 
-    :param dict[str]|None user_ns:
-    :param dict[str]|None user_global_ns:
-    :param bool exit_afterwards: will do sys.exit(1) at the end
+    :param user_ns:
+    :param user_global_ns:
+    :param exit_afterwards: will do sys.exit(1) at the end
     """
     print("Debug shell:")
     from returnn.util.basic import ObjAsDict

{returnn-1.20231215.23621 → returnn-1.20231219.120347/returnn.egg-info}/PKG-INFO

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

returnn-1.20231215.23621/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20231215.023621'
-long_version = '1.20231215.023621+git.190c791'