returnn 1.20240607.133858__tar.gz → 1.20240610.111814__tar.gz

{returnn-1.20240607.133858/returnn.egg-info → returnn-1.20240610.111814}/PKG-INFO

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

returnn-1.20240610.111814/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240610.111814'
+long_version = '1.20240610.111814+git.b8dc968'

{returnn-1.20240607.133858 → returnn-1.20240610.111814}/returnn/datasets/basic.py

@@ -28,6 +28,9 @@ from returnn.util.basic import try_run, NumbersDict, OptionalNotImplementedError
 from returnn.tensor import TensorDict
 
 
+RANDOM_SEED_OFFSET_ENV_VAR = "RETURNN_RANDOM_SEED_OFFSET"
+
+
 class Dataset(object):
     """
     Base class for any dataset. This defines the dataset API.
@@ -245,7 +248,11 @@ class Dataset(object):
         config = get_global_config(raise_exception=False)
         if not config:
             return 0
-        if config.typed_value("torch_distributed") is not None:
+
+        env_val = os.environ.get(RANDOM_SEED_OFFSET_ENV_VAR)
+        if env_val is not None:
+            return int(env_val)
+        elif config.typed_value("torch_distributed") is not None:
             import returnn.torch.distributed
 
             return returnn.torch.distributed.get_ctx(config=config).rank() * 16127

{returnn-1.20240607.133858 → returnn-1.20240610.111814}/returnn/datasets/concat_files.py

@@ -11,10 +11,10 @@ import os
 import sys
 import numpy
 from returnn.log import log
-from returnn.util.basic import try_run
+from returnn.util.basic import override_env_var, try_run
 from returnn.util.multi_proc_non_daemonic_spawn import NonDaemonicSpawnContext
 from returnn.config import SubProcCopyGlobalConfigPreInitFunc
-from .basic import init_dataset, DatasetSeq
+from .basic import init_dataset, DatasetSeq, RANDOM_SEED_OFFSET_ENV_VAR
 from .cached2 import CachedDataset2
 
 # noinspection PyProtectedMember
@@ -423,6 +423,9 @@ class _WorkerProcParent:
         buffer_size: int,
         exit_hook: Optional[Callable[[], None]] = None,
     ):
+        # the dataset makes sure this is set
+        assert "random_seed_offset" in dataset_dict
+
         self.epoch = epoch
         self.full_epoch_0idx = full_epoch_0idx
         self.dataset_dict = dataset_dict
@@ -431,13 +434,16 @@ class _WorkerProcParent:
         parent_conn, child_conn = _mp.Pipe()
         self.parent_conn: mpConnection = parent_conn
 
-        self.worker_proc = _mp.Process(
-            name=f"{name} worker ep {epoch}",
-            target=_worker_proc_loop,
-            args=(epoch, buffer_size, dataset_dict, child_conn),
-            daemon=True,
-        )
-        self.worker_proc.start()
+        # the env will be forwarded to the child process
+        with override_env_var(RANDOM_SEED_OFFSET_ENV_VAR, str(dataset_dict["random_seed_offset"])):
+            self.worker_proc = _mp.Process(
+                name=f"{name} worker ep {epoch}",
+                target=_worker_proc_loop,
+                args=(epoch, buffer_size, dataset_dict, child_conn),
+                daemon=True,
+            )
+            self.worker_proc.start()
+
         # Make sure the child connection is closed here.
         # It stays open in the child, until the child dies.
         # When that happens, now any consecutive read on the pipe

{returnn-1.20240607.133858 → returnn-1.20240610.111814}/returnn/datasets/hdf.py

@@ -3,6 +3,7 @@ Provides :class:`HDFDataset`.
 """
 
 from __future__ import annotations
+import bisect
 import typing
 import collections
 import gc
@@ -457,10 +458,10 @@ class HDFDataset(CachedDataset):
         return ", ".join(["HDF dataset", "sequences: %i" % self.num_seqs, "frames: %i" % self.get_num_timesteps()])
 
     def _get_file_index(self, real_seq_idx):
-        file_index = 0
-        while file_index < len(self.file_start) - 1 and real_seq_idx >= self.file_start[file_index + 1]:
-            file_index += 1
-        return file_index
+        # bisect() returns the position for which all elements to the left of the returned index are <= real_seq_idx,
+        # so it actually returns the next file index in which the sequence can be found.
+        # Therefore, we subtract 1 to the index provided.
+        return bisect.bisect(self.file_start, real_seq_idx) - 1
 
 
 # ------------------------------------------------------------------------------

{returnn-1.20240607.133858 → returnn-1.20240610.111814}/returnn/datasets/meta.py

@@ -23,10 +23,10 @@ The dataset classes MetaDataset and CombinedDataset which perform these tasks ar
 
 from __future__ import annotations
 
-from typing import List
+from typing import Optional, Any, Sequence, List, Dict
 from returnn.datasets.basic import Dataset, DatasetSeq, init_dataset, convert_data_dims
 from .cached2 import CachedDataset2
-from returnn.util.basic import NumbersDict, load_json
+from returnn.util.basic import NumbersDict, load_json, OptionalNotImplementedError
 from returnn.log import log
 from random import Random
 import numpy
@@ -949,7 +949,15 @@ class CombinedDataset(CachedDataset2):
         self.sampling_sizes = sampling_sizes
 
         # This will only initialize datasets needed for features occurring in data_map
-        self.datasets = {key: init_dataset(datasets[key]) for key in self.dataset_keys}
+        self.datasets = {
+            key: init_dataset(
+                {
+                    **datasets[key],
+                    "random_seed_offset": datasets[key].get("random_seed_offset", self.random_seed_offset),
+                }
+            )
+            for key in self.dataset_keys
+        }
 
         self._estimated_num_seqs = sum([self.datasets[k].estimated_num_seqs for k in sorted(self.datasets.keys())])
         self.estimated_num_seq_per_subset = [self.datasets[k].estimated_num_seqs for k in sorted(self.datasets.keys())]
@@ -1772,13 +1780,245 @@ class ChunkShuffleDataset(CachedDataset2):
         return self.dataset.get_target_list()
 
 
-def _simple_to_bool(v):
-    if v == 0:
-        v = False
-    if v == 1:
-        v = True
-    assert isinstance(v, bool)
-    return v
+class VariableDataset(Dataset):
+    """
+    For every (sub)epoch, it would generate a new subdataset,
+    based on a user-provided function.
+    """
+
+    def __init__(self, *, get_dataset, **kwargs):
+        """
+        :param get_dataset: function (*, epoch: int, **_) -> Dict[str,Any], will be called for every sub-epoch.
+            It will cache the dict from the prev call, and if the dict is the same, it will not recreate the dataset.
+        """
+        super().__init__(**kwargs)
+        self._get_dataset = get_dataset
+        self._dataset_dict: Optional[Dict[str, Any]] = None
+        self._dataset: Optional[Dataset] = None
+        self._load_dataset(epoch=1)
+        self.num_inputs = self._dataset.num_inputs
+        self.num_outputs = self._dataset.num_outputs
+        self.labels = self._dataset.labels
+
+    def _load_dataset(self, epoch: int):
+        dataset_dict = self._get_dataset(epoch=epoch)
+        if dataset_dict != self._dataset_dict:
+            self._dataset_dict = dataset_dict
+            self._dataset = init_dataset(dataset_dict)
+
+    def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
+        """init seq order"""
+        super().init_seq_order()
+        if epoch is None:
+            if seq_list is not None or seq_order is not None:
+                raise ValueError(f"{self}: epoch is None, but given seq_list or seq_order, not supported")
+            return
+        self._load_dataset(epoch)
+        self._dataset.init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
+
+    def finish_epoch(self, *, free_resources: bool = False):
+        """finish epoch"""
+        super().finish_epoch(free_resources=free_resources)
+        self._dataset.finish_epoch(free_resources=free_resources)
+
+    def supports_seq_order_sorting(self) -> bool:
+        """supports sorting"""
+        return self._dataset.supports_seq_order_sorting()
+
+    def get_current_seq_order(self) -> Sequence[int]:
+        """current seq order"""
+        return self._dataset.get_current_seq_order()
+
+    def get_all_tags(self) -> List[str]:
+        """all tags"""
+        return self._dataset.get_all_tags()
+
+    def get_total_num_seqs(self) -> int:
+        """total num seqs"""
+        return self._dataset.get_total_num_seqs()
+
+    def get_seq_length(self, sorted_seq_idx: int) -> NumbersDict:
+        """seq len"""
+        return self._dataset.get_seq_length(sorted_seq_idx)
+
+    def get_tag(self, sorted_seq_idx: int) -> str:
+        """tag"""
+        return self._dataset.get_tag(sorted_seq_idx)
+
+    def get_data_keys(self) -> List[str]:
+        """data keys"""
+        return self._dataset.get_data_keys()
+
+    def get_target_list(self) -> List[str]:
+        """target list"""
+        return self._dataset.get_target_list()
+
+    def is_cached(self, start: int, end: int) -> bool:
+        """is cached"""
+        return self._dataset.is_cached(start, end)
+
+    @property
+    def num_seqs(self) -> int:
+        """num seqs"""
+        return self._dataset.num_seqs
+
+    def is_less_than_num_seqs(self, n: int) -> bool:
+        """n < num_seqs"""
+        return self._dataset.is_less_than_num_seqs(n)
+
+    def get_num_timesteps(self) -> int:
+        """num timesteps"""
+        return self._dataset.get_num_timesteps()
+
+    def load_seqs(self, start: int, end: int):
+        """load seqs"""
+        self._dataset.load_seqs(start, end)
+
+    def _load_seqs(self, start: int, end: int):
+        raise Exception(f"{self}: _load_seqs: should not get here")
+
+    def get_data(self, seq_idx: int, key: str) -> numpy.ndarray:
+        """data"""
+        return self._dataset.get_data(seq_idx, key)
+
+    def get_input_data(self, seq_idx: int) -> numpy.ndarray:
+        """input data"""
+        return self._dataset.get_input_data(seq_idx)
+
+    def get_targets(self, target: str, seq_idx: int) -> numpy.ndarray:
+        """target data"""
+        return self._dataset.get_targets(target, seq_idx)
+
+    def get_data_dim(self, key: str) -> int:
+        """data dim"""
+        return self._dataset.get_data_dim(key)
+
+    def get_data_shape(self, data_key: str) -> List[int]:
+        """data shape"""
+        return self._dataset.get_data_shape(data_key)
+
+    def get_data_dtype(self, key: str) -> str:
+        """data dtype"""
+        return self._dataset.get_data_dtype(key)
+
+    def is_data_sparse(self, key: str) -> bool:
+        """is data sparse"""
+        return self._dataset.is_data_sparse(key)
+
+
+class AnythingDataset(Dataset):
+    """
+    An infinite dataset, creating dummy (zero) data on the fly,
+    given the data-keys and their shapes.
+
+    When this is used inside a :class:`MetaDataset`,
+    controlled by the seqs list from another dataset,
+    it will just take over whatever seq list is given.
+    """
+
+    def __init__(self, *, data_keys: Dict[str, Dict[str, Any]], **kwargs):
+        """
+        :param data_keys: similar like extern_data. defines shape, dtype, sparse, dim, etc
+        """
+        super().__init__(**kwargs)
+        self._data_keys = data_keys
+        self._input_key = next(iter(data_keys))
+        self.num_inputs = data_keys[self._input_key]["dim"]
+        self.num_outputs = {k: (v["dim"], len(v["shape"])) for (k, v) in data_keys.items()}
+        self.labels = None
+        self._seq_list = None
+        self._seq_order = None
+
+    def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
+        """init seq order"""
+        super().init_seq_order()
+        self._seq_list = seq_list
+        self._seq_order = seq_order
+
+    def supports_seq_order_sorting(self) -> bool:
+        """supports sorting"""
+        return True
+
+    def get_current_seq_order(self) -> Sequence[int]:
+        """current seq order"""
+        if self._seq_order is not None:
+            return self._seq_order
+        raise OptionalNotImplementedError
+
+    def have_seqs(self) -> bool:
+        """total num seqs"""
+        return True
+
+    def get_seq_length(self, sorted_seq_idx: int) -> NumbersDict:
+        """seq len"""
+        return NumbersDict(1)
+
+    def get_tag(self, sorted_seq_idx: int) -> str:
+        """tag"""
+        if self._seq_list:
+            return self._seq_list[sorted_seq_idx]
+        return f"seq-{sorted_seq_idx}"
+
+    def get_data_keys(self) -> List[str]:
+        """data keys"""
+        return list(self._data_keys)
+
+    def get_target_list(self) -> List[str]:
+        """target list"""
+        return [key for key in self._data_keys if key != self._input_key]
+
+    def is_cached(self, start: int, end: int) -> bool:
+        """is cached"""
+        return True
+
+    @property
+    def num_seqs(self) -> int:
+        """num seqs"""
+        if self._seq_order:
+            return len(self._seq_order)
+        if self._seq_list:
+            return len(self._seq_list)
+        raise NotImplementedError
+
+    def is_less_than_num_seqs(self, n: int) -> bool:
+        """n < num_seqs"""
+        if self._seq_order:
+            return n < len(self._seq_order)
+        if self._seq_list:
+            return n < len(self._seq_list)
+        return True  # we can always generate more
+
+    def _load_seqs(self, start: int, end: int):
+        pass  # noop
+
+    def get_data(self, seq_idx: int, key: str) -> numpy.ndarray:
+        """data"""
+        data = self._data_keys[key]
+        return numpy.zeros([d or 1 for d in data["shape"]], dtype=data["dtype"])
+
+    def get_input_data(self, seq_idx: int) -> numpy.ndarray:
+        """input data"""
+        return self.get_data(seq_idx, self._input_key)
+
+    def get_targets(self, target: str, seq_idx: int) -> numpy.ndarray:
+        """target data"""
+        return self.get_data(seq_idx, target)
+
+    def get_data_dim(self, key: str) -> int:
+        """data dim"""
+        return self._data_keys[key]["dim"]
+
+    def get_data_shape(self, data_key: str) -> List[int]:
+        """data shape"""
+        return list(self._data_keys[data_key]["shape"])
+
+    def get_data_dtype(self, key: str) -> str:
+        """data dtype"""
+        return self._data_keys[key]["dtype"]
+
+    def is_data_sparse(self, key: str) -> bool:
+        """is data sparse"""
+        return self._data_keys[key].get("sparse", False)
 
 
 def _select_dtype(key, data_dims, data_dtypes):

{returnn-1.20240607.133858 → returnn-1.20240610.111814}/returnn/util/basic.py

@@ -4519,3 +4519,23 @@ def find_libcudart_from_runtime():
             return fn
     _find_libcudart_from_runtime_cached = [None]
     return None
+
+
+@contextlib.contextmanager
+def override_env_var(var_name: str, value: str):
+    """
+    context manager for temporarily overriding the value of an env var
+
+    :param var_name: the name of the environment variable to override
+    :param value: the value to set while the context mgr is active
+    """
+
+    cur_val = os.environ.get(var_name)
+    os.environ[var_name] = value
+    try:
+        yield
+    finally:
+        if cur_val is not None:
+            os.environ[var_name] = cur_val
+        else:
+            os.environ.pop(var_name)

{returnn-1.20240607.133858 → returnn-1.20240610.111814}/returnn/util/multi_proc_non_daemonic_spawn.py

@@ -94,7 +94,7 @@ class NonDaemonicSpawnProcess(SpawnProcess):
             signals = ["SIGINT"] * 3 + ["SIGTERM", "SIGKILL"]
             for sig in signals:
                 try:
-                    print(f"Send signal {sig} to pid {self.ident}")
+                    print(f"Send signal {sig} to pid {self.ident}/'{self.name or 'N/A'}'")
                     os.kill(self.ident, getattr(signal, sig))
                 except ProcessLookupError:
                     break

{returnn-1.20240607.133858 → returnn-1.20240610.111814/returnn.egg-info}/PKG-INFO

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

returnn-1.20240607.133858/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240607.133858'
-long_version = '1.20240607.133858+git.ffd35e5'