returnn 1.20230411.4301__tar.gz → 1.20230411.133415__tar.gz

{returnn-1.20230411.4301/returnn.egg-info → returnn-1.20230411.133415}/PKG-INFO

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

returnn-1.20230411.133415/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230411.133415'
+long_version = '1.20230411.133415+git.9593d64'

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/demos/demo-iter-dataset.py

@@ -36,10 +36,10 @@ def iterate_epochs():
     """
     Iterate through epochs.
     """
-    start_epoch, start_batch = EngineBase.get_train_start_epoch_batch(config)
+    start_epoch = EngineBase.get_train_start_epoch(config)
     final_epoch = EngineBase.config_get_final_epoch(config)
 
-    print("Starting with epoch %i, batch %i." % (start_epoch, start_batch), file=log.v3)
+    print("Starting with epoch %i." % (start_epoch,), file=log.v3)
     print("Final epoch is: %i" % final_epoch, file=log.v3)
 
     recurrent_net = "lstm" in config.value("hidden_type", "")  # good enough...

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/returnn/engine/base.py

@@ -4,9 +4,9 @@ Provides :class:`EngineBase`.
 
 from __future__ import annotations
 
+from typing import Optional
 import os
 import sys
-from typing import Optional
 
 from returnn.config import Config, get_global_config
 from returnn.learning_rate_control import load_learning_rate_control_from_config, LearningRateControl
@@ -164,35 +164,27 @@ class EngineBase(object):
         return epoch_model
 
     @classmethod
-    def get_train_start_epoch_batch(cls, config):
+    def get_train_start_epoch(cls, config: Config) -> int:
         """
         We will always automatically determine the best start (epoch,batch) tuple
         based on existing model files.
         This ensures that the files are present and enforces that there are
         no old outdated files which should be ignored.
         Note that epochs start at idx 1 and batches at idx 0.
-        :type config: returnn.config.Config
-        :returns (epoch,batch)
-        :rtype (int,int)
+
+        :param config:
+        :return: epoch
         """
         start_batch_mode = config.value("start_batch", "auto")
-        if start_batch_mode == "auto":
-            start_batch_config = None
-        else:
-            start_batch_config = int(start_batch_mode)
+        if start_batch_mode != "auto":
+            raise Exception(f"custom start_batch {start_batch_mode!r} not supported")
         last_epoch, _ = cls.get_epoch_model(config)
         if last_epoch is None:
             start_epoch = 1
-            start_batch = start_batch_config or 0
-        elif start_batch_config is not None:
-            # We specified a start batch. Stay in the same epoch, use that start batch.
-            start_epoch = last_epoch
-            start_batch = start_batch_config
         else:
             # Start with next epoch.
             start_epoch = last_epoch + 1
-            start_batch = 0
-        return start_epoch, start_batch
+        return start_epoch
 
     @classmethod
     def epoch_model_filename(cls, model_filename, epoch, is_pretrain):

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/returnn/sprint/interface.py

@@ -355,8 +355,8 @@ def getSegmentList(corpusName, segmentList, **kwargs):
     sprintDataset.use_multiple_epochs()
 
     finalEpoch = get_final_epoch()
-    startEpoch, startSegmentIdx = Engine.get_train_start_epoch_batch(config)
-    print("Sprint: Starting with epoch %i, segment-idx %s." % (startEpoch, startSegmentIdx))
+    startEpoch = Engine.get_train_start_epoch(config)
+    print("Sprint: Starting with epoch %i." % (startEpoch,))
     print("Final epoch is: %i" % finalEpoch)
 
     # Loop over multiple epochs. Epochs start at 1.

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/returnn/tensor/_tensor_extra.py

@@ -787,18 +787,12 @@ class _TensorMixin(_TensorMixinBase):
 
     def copy_add_batch_dim(self, batch_dim_axis, batch=None, dim_tag=None) -> _t.Tensor:
         """
-        Warning: Assumes TensorFlow.
-
         :param int batch_dim_axis:
         :param BatchInfo|None batch:
         :param Dim|None dim_tag:
         :return: copy of myself with added batch-dim
         """
         assert self.batch_dim_axis is None
-        if not batch:
-            from returnn.tf.layers.base import LayerBase
-
-            batch = LayerBase.get_recent_layer().get_batch_info()
         if batch_dim_axis < 0:
             assert batch_dim_axis + self.batch_ndim + 1 >= 0
             batch_dim_axis += self.batch_ndim + 1
@@ -808,19 +802,36 @@ class _TensorMixin(_TensorMixinBase):
         if placeholder is not None:
             backend = self._raw_backend
             placeholder = backend.expand_dims_raw(placeholder, batch_dim_axis)
-            if not isinstance(batch.dim, int) or batch.dim != 1:
-                placeholder = backend.expand_raw(placeholder, batch_dim_axis, batch.dim)
+            if batch:
+                batch_dim_ = batch.dim
+            elif dim_tag:
+                if dim_tag.dyn_size_ext:
+                    assert dim_tag.dyn_size_ext.dims == ()
+                    assert dim_tag.dyn_size_ext.raw_tensor is not None
+                    batch_dim_ = dim_tag.dyn_size_ext.raw_tensor
+                elif dim_tag.dimension:
+                    batch_dim_ = dim_tag.dimension
+                else:
+                    raise Exception(f"{self} copy_add_batch_dim: unknown batch dim for {dim_tag!r}")
+            else:
+                raise Exception(f"{self} copy_add_batch_dim: unknown batch dim ")
+            if not isinstance(batch_dim_, int) or batch_dim_ != 1:
+                placeholder = backend.expand_raw(placeholder, batch_dim_axis, batch_dim_)
         dim_tags = list(self.dim_tags)
         if dim_tag:
             assert dim_tag.is_batch_dim()
-            assert dim_tag.dimension == batch.static_dim or dim_tag.dimension is None
             assert dim_tag.batch == batch
+            if batch:
+                assert dim_tag.dimension == batch.static_dim or dim_tag.dimension is None
         else:
-            dim_tag = Dim(kind=Dim.Types.Batch, description="batch", dimension=batch.static_dim, batch=batch)
+            dim_tag = Dim(
+                kind=Dim.Types.Batch, description="batch", dimension=batch.static_dim if batch else None, batch=batch
+            )
         dim_tags.insert(batch_dim_axis, dim_tag)
         data_opts["dims"] = dim_tags
-        data_opts["batch"] = batch
-        data_opts["beam"] = batch.beam
+        if batch:
+            data_opts["batch"] = batch
+            data_opts["beam"] = batch.beam
         other_special_axes = self.get_special_axes_dict(counted_with_batch_dim=True, only_available=True)
         for k, a in other_special_axes.items():
             data_opts[k] = a if (a < batch_dim_axis) else (a + 1)
@@ -909,14 +920,17 @@ class _TensorMixin(_TensorMixinBase):
             if unbroadcast:
                 return self.copy_add_batch_dim(batch_dim_axis=axis, batch=dim_tag.batch, dim_tag=dim_tag)
             else:
-                from returnn.tf.util.data import BatchInfo
+                if dim_tag.batch or self.batch:
+                    from returnn.tf.util.data import BatchInfo
 
-                batch_info = BatchInfo.make_global_broadcast_batch_info()
-                return self.copy_add_batch_dim(
-                    batch_dim_axis=axis,
-                    batch=batch_info,
-                    dim_tag=dim_tag if (dim_tag.dimension == 1 and dim_tag.batch == batch_info) else None,
-                )
+                    batch_info = BatchInfo.make_global_broadcast_batch_info()
+                else:
+                    batch_info = None
+                if dim_tag and dim_tag.dimension == 1 and dim_tag.batch == batch_info:
+                    pass  # keep it
+                else:
+                    dim_tag = Dim(kind=Dim.Types.Batch, description="batch-broadcast", dimension=1, batch=batch_info)
+                return self.copy_add_batch_dim(batch_dim_axis=axis, batch=batch_info, dim_tag=dim_tag)
 
         data_opts = self.get_kwargs()
         # Note: if dim_tag is feature, but we are sparse, we just make it spatial
@@ -2812,9 +2826,11 @@ class _TensorMixin(_TensorMixinBase):
         :rtype: int|None
         """
         # Do not fallback to get_batch_dim or get_recent_layer or so. This should be safe.
-        if not self.batch:
-            return None
-        return self.batch.static_dim
+        if self.batch:
+            return self.batch.static_dim
+        if self.have_batch_axis():
+            return self.get_batch_dim_tag().dimension
+        return None
 
     def get_spatial_batch_axes(self):
         """

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/returnn/tf/engine.py

@@ -1105,7 +1105,7 @@ class Engine(EngineBase):
         if config.has("eval_datasets"):
             for dataset_name, dataset_opts in config.typed_value("eval_datasets", {}).items():
                 self.eval_datasets[dataset_name] = init_dataset(dataset_opts, default_kwargs={"name": dataset_name})
-        self.start_epoch, self.start_batch = self.get_train_start_epoch_batch(config)
+        self.start_epoch = self.get_train_start_epoch(config)
         self.batch_size = config.typed_value("batch_size", 1)
         self.shuffle_batches = config.bool("shuffle_batches", False)
         self.update_batch_size = config.int("update_batch_size", 0)
@@ -1577,7 +1577,7 @@ class Engine(EngineBase):
         """
         Does the whole training, i.e. the loop over all the epochs.
         """
-        print("start training at epoch %i and step %i" % (self.start_epoch, self.start_batch), file=log.v3)
+        print("start training at epoch %i" % (self.start_epoch,), file=log.v3)
         print("using batch size: %r, max seqs: %i" % (self.batch_size, self.max_seqs), file=log.v4)
         print("learning rate control:", self.learning_rate_control, file=log.v4)
         print("pretrain:", self.pretrain, file=log.v4)

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/returnn/torch/engine.py

@@ -90,7 +90,7 @@ class Engine(EngineBase):
         for dataset_name, dataset in self.eval_datasets.items():
             self._eval_dataloaders[dataset_name] = self._create_data_loader(dataset)
 
-        self._start_epoch, _ = self.get_train_start_epoch_batch(self.config)
+        self._start_epoch = self.get_train_start_epoch(self.config)
         self._final_epoch = self.config_get_final_epoch(self.config)
 
         self._load_model(epoch=self._start_epoch)

{returnn-1.20230411.4301 → returnn-1.20230411.133415/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/tests/_setup_test_env.py

@@ -64,13 +64,16 @@ def setup():
     log.initialize(verbosity=[5], propagate=False)
 
     # TF is optional.
-    # Note that importing TF still has a small side effect:
-    # BackendEngine._get_default_engine() will return TF by default, if TF is already loaded.
-    # For most tests, this does not matter.
-    try:
-        import tensorflow as tf
-    except ImportError:
+    if "RETURNN_DISABLE_TF" in os.environ and int(os.environ["RETURNN_DISABLE_TF"]) == 1:
         tf = None
+    else:
+        # Note that importing TF still has a small side effect:
+        # BackendEngine._get_default_engine() will return TF by default, if TF is already loaded.
+        # For most tests, this does not matter.
+        try:
+            import tensorflow as tf
+        except ImportError:
+            tf = None
 
     if tf:
         import returnn.tf.util.basic as tf_util

{returnn-1.20230411.4301 → returnn-1.20230411.133415}/tests/test_TFUtil.py

@@ -424,7 +424,7 @@ def test_Data_find_matching_dim_map_broadcast_matches():
     assert mapping[0] == 1 and mapping[1] == 2
 
     copied = d2.copy_compatible_to(d1)
-    assert copied.batch_ndim == d1.batch_ndim and copied.batch.static_dim == 1 and copied.batch_shape == (1, 5, 1)
+    assert copied.batch_ndim == d1.batch_ndim and copied.get_static_batch_dim() == 1 and copied.batch_shape == (1, 5, 1)
     print("copied compatible:", copied)
 
 
@@ -633,14 +633,14 @@ def test_Data_copy_compatible_to_src_no_batch():
     d2 = Data(name="d2", shape=(), batch_dim_axis=None, time_dim_axis=None)
     d2.placeholder = tf.zeros([d if (d is not None) else 1 for d in d2.batch_shape])
     d3 = d2.copy_compatible_to(d1)
-    assert d3.batch.static_dim == 1 and d3.batch_shape == (1, 1, 1)
+    assert d3.get_static_batch_dim() == 1 and d3.batch_shape == (1, 1, 1)
 
 
 def test_Data_copy_compatible_to_add_batch_dim():
     common_data = Data(name="accum_att_weights_output", shape=(None, 1))
     d1 = Data(name="att_weights_avg_output", shape=(1,), batch_dim_axis=None)
     d2 = d1.copy_compatible_to(common_data)
-    assert d2.have_batch_axis() and d2.batch.static_dim == 1 and d2.batch_shape == (1, 1, 1)
+    assert d2.have_batch_axis() and d2.get_static_batch_dim() == 1 and d2.batch_shape == (1, 1, 1)
 
 
 def test_Data_copy_compatible_to_add_feature_dim():
@@ -658,7 +658,7 @@ def test_Data_copy_compatible_to_add_batch_feature_dim():
     d2 = d1.copy_compatible_to(common_data)
     print(d2)
     assert d2.have_batch_axis() and d2.have_time_axis()
-    assert d2.batch.static_dim == 1  # batch-dim is broadcasted
+    assert d2.get_static_batch_dim() == 1  # batch-dim is broadcasted
     assert d2.batch_shape == (1, None, 1)
 
 

returnn-1.20230411.4301/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230411.004301'
-long_version = '1.20230411.004301+git.86c1436'