returnn 1.20230630.233220__tar.gz → 1.20230701.3857__tar.gz

{returnn-1.20230630.233220/returnn.egg-info → returnn-1.20230701.3857}/PKG-INFO

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

returnn-1.20230701.3857/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230701.003857'
+long_version = '1.20230701.003857+git.283731c'

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/datasets/basic.py

@@ -19,12 +19,13 @@ import os
 import numpy
 import functools
 import typing
-from typing import Optional, Union, Type, Dict, List
+from typing import Optional, Union, Type, Dict, Sequence, List
 
 from returnn.log import log
 from returnn.engine.batch import Batch, BatchSetGenerator
 from returnn.datasets.util.vocabulary import Vocabulary
 from returnn.util.basic import try_run, NumbersDict, OptionalNotImplementedError
+from returnn.tensor import TensorDict
 
 
 class Dataset(object):
@@ -1469,14 +1470,20 @@ def convert_data_dims(data_dims, leave_dict_as_is=False):
     return data_dims
 
 
-def shapes_for_batches(batches, *, data_keys, dataset=None, extern_data, enforce_min_len1=False):
+def shapes_for_batches(
+    batches: Sequence[Batch],
+    *,
+    data_keys: Sequence[str],
+    dataset: Optional[Dataset] = None,
+    extern_data: Optional[TensorDict],
+    enforce_min_len1: bool = False,
+) -> Optional[Dict[str, List[int]]]:
     """
-    :param list[EngineBatch.Batch] batches:
-    :param list[str] data_keys:
-    :param Dataset dataset:
-    :param returnn.tf.network.ExternData extern_data: detailed data description
-    :param bool enforce_min_len1:
-    :rtype: dict[str,list[int]] | None
+    :param batches:
+    :param data_keys:
+    :param dataset:
+    :param extern_data: detailed data description
+    :param enforce_min_len1:
     """
     assert dataset or extern_data
     all_data_keys = set(data_keys)

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/frontend/_backend.py

@@ -370,6 +370,15 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def squeeze(source: Tensor, axis: Dim) -> Tensor:
+        """
+        :param source:
+        :param axis:
+        :return: source with axis removed
+        """
+        raise NotImplementedError
+
     @staticmethod
     def concat(
         *sources: Tuple[Tensor, Dim],

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/frontend/array_.py

@@ -22,6 +22,7 @@ __all__ = [
     "reshape",
     "split",
     "expand_dim",
+    "squeeze",
     "concat",
     "concat_features",
     "pad",
@@ -239,6 +240,15 @@ def expand_dim(source: Tensor, dim: Dim) -> Tensor:
     return source._raw_backend.expand_dim(source, dim=dim)
 
 
+def squeeze(source: Tensor, axis: Dim) -> Tensor:
+    """
+    Removes the axis with dimension of extend 1 from the source.
+    """
+    assert axis.dimension == 1, f"squeeze {source}: axis {axis} is not of extend 1"
+    # noinspection PyProtectedMember
+    return source._raw_backend.squeeze(source, axis=axis)
+
+
 def concat(
     *sources: Tuple[Tensor, Dim],
     allow_broadcast: bool = False,

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/tf/data_pipeline.py

@@ -118,19 +118,16 @@ from __future__ import annotations
 import contextlib
 import sys
 import typing
+from typing import Optional, Dict
 
-try:
-    # noinspection PyCompatibility
-    from Queue import Queue
-except ImportError:
-    # noinspection PyCompatibility,PyUnresolvedReferences
-    from queue import Queue
+from queue import Queue
 from threading import Thread, Condition
 
 import numpy
 import tensorflow as tf
 
 from returnn.datasets.basic import Dataset, BatchSetGenerator
+from returnn.tensor import TensorDict
 from returnn.tf.network import ExternData
 import returnn.tf.compat as tf_compat
 import returnn.tf.horovod as tf_horovod
@@ -142,10 +139,10 @@ class DataProviderBase(object):
     Base class which wraps up the logic in this class. See derived classes.
     """
 
-    def __init__(self, extern_data, data_keys=None):
+    def __init__(self, *, extern_data: TensorDict, data_keys: Optional[str] = None):
         """
-        :param ExternData extern_data:
-        :param set(str)|None data_keys:
+        :param extern_data:
+        :param data_keys:
         """
         self.coord = tf.train.Coordinator()
         self.extern_data = extern_data
@@ -153,6 +150,16 @@ class DataProviderBase(object):
             data_keys = extern_data.data.keys()
         self.data_keys = sorted(data_keys)  # type: typing.List[str]
 
+    def _exclude_data_key(self, key: str) -> bool:
+        """
+        :param key:
+        :return: whether to exclude this data-key, even if in data_keys
+        """
+        if isinstance(self.extern_data, ExternData):
+            if key in self.extern_data.extra_added_keys:
+                return True
+        return False
+
     def start_threads(self, session):
         """
         Start threads.
@@ -218,27 +225,31 @@ class FeedDictDataProvider(DataProviderBase):
     """
     This class will fill all the placeholders used for training or forwarding or evaluation etc.
     of a :class:`returnn.tf.network.TFNetwork`.
-    It will run a background thread which reads the data from a dataset and puts it into a queue.
+
+    It will run a background thread which reads the data from a dataset and puts it into a queue
+    when you call start_threads().
+
+    In principle, this class can also be used independently of TensorFlow.
     """
 
     def __init__(
         self,
         *,
-        dataset,
-        batches,
-        enforce_min_len1=False,
-        capacity=10,
-        batch_slice=None,
+        dataset: Dataset,
+        batches: BatchSetGenerator,
+        enforce_min_len1: bool = False,
+        capacity: int = 10,
+        batch_slice: Optional[slice] = None,
         **kwargs,
     ):
         """
-        :param Dataset dataset:
-        :param BatchSetGenerator batches:
-        :param bool enforce_min_len1:
-        :param ExternData extern_data:
-        :param set(str)|None data_keys:
-        :param int capacity:
-        :param slice|None batch_slice: select a subset of the batches
+        :param dataset:
+        :param batches:
+        :param enforce_min_len1:
+        :param capacity:
+        :param batch_slice: select a subset of the batches
+        :param extern_data:
+        :param data_keys:
         """
         super(FeedDictDataProvider, self).__init__(**kwargs)
         self.dataset = dataset
@@ -275,13 +286,15 @@ class FeedDictDataProvider(DataProviderBase):
             self.thread = None
         self.dataset.finish_epoch()
 
-    def get_next_batch(self, consider_batch_slice):
+    def get_next_batch(self, *, consider_batch_slice: bool = False) -> Optional[Dict[str, numpy.ndarray]]:
         """
         This assumes that we have more data, i.e. self.batches.has_more().
 
-        :param bool consider_batch_slice:
+        The code is basically independent of TF now.
+        self.extern_data can be any TensorDict.
+
+        :param consider_batch_slice:
         :returns: batch-data-value-dict or None. if not consider_batch_slice, will never be None
-        :rtype: dict[str,numpy.ndarray]|None
         """
         # See EngineUtil.assign_dev_data() for reference.
         cur_batch_idx = self.cur_batch_idx
@@ -335,7 +348,7 @@ class FeedDictDataProvider(DataProviderBase):
                     # See also :func:`TFNetwork.get_extern_data`.
                     if k in ["seq_idx", "seq_tag"]:
                         continue  # handled below. will always be added
-                    if k in self.extern_data.extra_added_keys:
+                    if self._exclude_data_key(k):
                         continue
                     data_ = self.extern_data.data[k]
                     # Do not rely on time_dim_axis but check for any dynamic axes.
@@ -447,16 +460,12 @@ class FeedDictDataProvider(DataProviderBase):
             output = self.queue.get()
         assert isinstance(output, dict)
         # The data itself.
-        d = {
-            self.extern_data.get_data(k).placeholder: output[k]
-            for k in self.data_keys
-            if k not in self.extern_data.extra_added_keys
-        }
+        d = {self.extern_data.data[k].placeholder: output[k] for k in self.data_keys if not self._exclude_data_key(k)}
         # And seq lengths info.
         for k in self.data_keys:
-            if k in self.extern_data.extra_added_keys:
+            if self._exclude_data_key(k):
                 continue
-            data = self.extern_data.get_data(k)
+            data = self.extern_data.data[k]
             for dim, len_placeholder in data.size_placeholder.items():
                 if dim == 0:  # time-dim
                     d[len_placeholder] = output["%s_seq_lens" % k]
@@ -465,6 +474,7 @@ class FeedDictDataProvider(DataProviderBase):
                         "dataset currently does not support variable shape in other dimensions than the first. "
                         "dim=%i, placeholder=%r" % (dim, len_placeholder)
                     )
+        assert isinstance(self.extern_data, ExternData)
         batch_info = self.extern_data.get_batch_info()
         batch_dim = batch_info.dim
         if isinstance(batch_dim, int):

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/tf/frontend_layers/_backend.py

@@ -237,6 +237,11 @@ class ReturnnLayersBackend(Backend[Layer]):
             axis = "spatial"
         return rfl.make_layer({"class": "expand_dims", "from": source, "axis": axis, "dim": dim}, name="expand_dims")
 
+    @staticmethod
+    def squeeze(source: Tensor, axis: Dim) -> Tensor:
+        """squeeze"""
+        return rfl.make_layer({"class": "squeeze", "from": source, "axis": axis}, name="squeeze")
+
     @staticmethod
     def concat(
         *sources: Tuple[Tensor, Dim],

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/tf/layers/basic.py

@@ -7462,7 +7462,7 @@ class SqueezeLayer(_ConcatInputLayer):
     # noinspection PyUnusedLocal
     def __init__(self, axis, enforce_batch_dim_axis=None, allow_no_op=False, **kwargs):
         """
-        :param int|list[int]|str axis: one axis or multiple axis to squeeze.
+        :param Dim|int|list[int]|str axis: one axis or multiple axis to squeeze.
           this is counted with batch-dim, which by default is axis 0 (see enforce_batch_dim_axis).
           it also accepts the special tokens "B"|"batch", "spatial", "spatial_except_time", or "F"|"feature"
         :param int|None enforce_batch_dim_axis:
@@ -7486,7 +7486,7 @@ class SqueezeLayer(_ConcatInputLayer):
     @classmethod
     def _get_axes(cls, axis, input_data):
         """
-        :param int|list[int]|str axis: one axis or multiple axis to squeeze.
+        :param Dim|int|list[int]|str axis: one axis or multiple axis to squeeze.
         :param Data input_data:
         :rtype: list[int]
         """
@@ -7497,7 +7497,7 @@ class SqueezeLayer(_ConcatInputLayer):
     @classmethod
     def get_out_data_from_opts(cls, axis, enforce_batch_dim_axis=None, allow_no_op=False, sources=(), **kwargs):
         """
-        :param int|list[int]|str axis:
+        :param Dim|int|list[int]|str axis:
         :param int|None enforce_batch_dim_axis:
         :param bool allow_no_op:
         :param list[LayerBase] sources:

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/torch/distributed.py

@@ -27,7 +27,10 @@ class DistributedContext:
         """
         import torch.distributed as dist
 
-        dist.init_process_group("nccl")
+        # when no backend is specified, both gloo and nccl backends will be created
+        # the gloo backend will be used for collectives with CPU tensors and
+        # the nccl backend will be used for collectives with CUDA tensors
+        dist.init_process_group(backend=None)
 
         self._config = config
         self._local_rank = os.environ["LOCAL_RANK"]

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/torch/engine.py

@@ -248,11 +248,7 @@ class Engine(EngineBase):
 
         while True:
             extern_data_raw = next(data_iter, None)
-            # WARNING: torch.distributed works only for the registered device,
-            # as it uses only one mechanism for communication, like NCCL.
-            # This is suboptimal here as we have the roundtrip CPU -> GPU -> NCCL -> GPU -> CPU.
-            # TODO: Use more direct CPU -> Ethernet -> CPU communication.
-            _has_data = torch.tensor([extern_data_raw is not None], dtype=torch.int8).to(self._device)
+            _has_data = torch.tensor([extern_data_raw is not None], dtype=torch.int8)
 
             if self._use_torch_distributed:
                 # use all reduce to check if all workers have data, if at least one worker does not have data,

{returnn-1.20230630.233220 → returnn-1.20230701.3857}/returnn/torch/frontend/_backend.py

@@ -315,6 +315,14 @@ class TorchBackend(Backend[torch.Tensor]):
         out.raw_tensor = out_raw
         return out
 
+    @staticmethod
+    def squeeze(source: Tensor, axis: Dim) -> Tensor:
+        """squeeze"""
+        axis = source.get_axis_from_description(axis)
+        out = source.copy_template_excluding_axis(axis)
+        out.raw_tensor = torch.squeeze(source.raw_tensor, axis)
+        return out
+
     @staticmethod
     def concat(
         *sources: Tuple[Tensor, Dim],

{returnn-1.20230630.233220 → returnn-1.20230701.3857/returnn.egg-info}/PKG-INFO

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

returnn-1.20230630.233220/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230630.233220'
-long_version = '1.20230630.233220+git.1128360'