returnn 1.20240917.180949__tar.gz → 1.20240918.234323__tar.gz

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/PKG-INFO

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

returnn-1.20240918.234323/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240918.234323'
+long_version = '1.20240918.234323+git.9c0f6a9'

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/returnn/datasets/audio.py

@@ -146,7 +146,7 @@ class OggZipDataset(CachedDataset2):
         # However, some other code might expect that the labels are all strings, not bytes,
         # and the API requires the labels to be strings.
         # The code in Dataset.serialize_data tries to decode this case as utf8 (if possible).
-        self.labels["orth"] = [chr(i) for i in range(255)]
+        self.labels["orth"] = [chr(i) for i in range(256)]
         if self.targets:
             self.num_outputs["classes"] = [self.targets.num_labels, 1]
         if self.feature_extractor:

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/returnn/datasets/cached.py

@@ -222,7 +222,7 @@ class CachedDataset(Dataset):
         assert start >= 0
         assert start <= end
 
-        if self.is_cached(start, end, blocking=True):
+        if self.is_cached(start, end, blocking=False):
             return
 
         if self.cache_byte_size_limit_at_start > 0:  # If the cache is enabled.

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/returnn/datasets/hdf.py

@@ -94,23 +94,12 @@ class HDFDataset(CachedDataset):
             from returnn.util.basic import cf
 
             filename = cf(filename)
+        print("parsing file", filename, file=log.v5)
         fin = h5py.File(filename, "r")
-        if "targets" in fin:
-            self.labels = {
-                k: [self._decode(item) for item in fin["targets/labels"][k][...].tolist()]
-                for k in fin["targets/labels"]
-            }
-        if not self.labels and "labels" in fin:
-            labels = [item.split("\0")[0] for item in fin["labels"][...].tolist()]  # type: typing.List[str]
-            self.labels = {"classes": labels}
-            assert len(self.labels["classes"]) == len(labels), (
-                "expected " + str(len(self.labels["classes"])) + " got " + str(len(labels))
-            )
         self.files.append(filename)
         self.h5_files.append(fin)
         self.cached_h5_datasets.append({})
-        print("parsing file", filename, file=log.v5)
-        if "times" in fin:
+        if attr_times in fin:
             if self.timestamps is None:
                 self.timestamps = fin[attr_times][...]
             else:
@@ -119,10 +108,36 @@ class HDFDataset(CachedDataset):
         if len(self.files) >= 2:
             prev_target_keys = self.target_keys
         if "targets" in fin:
+            # Note: Earlier RETURNN versions used "targets/labels" to determine target_keys.
+            # https://github.com/rwth-i6/returnn/blob/c2d8fed877022d1ac1bf68b801604733db51223e/HDFDataset.py#L60
             self.target_keys = sorted(set(fin["targets/data"].keys()) | set(fin["targets/size"].attrs.keys()))
         else:
+            # Actually this "classes" target key is never used.
+            # Only if there are "targets" in the HDF file, we use the keys from there.
+            # HDFDataset.get_target_list() returns an empty list,
+            # and HDFDataset.get_data_keys() uses get_target_list() + optional ["data"].
+            # HDFDataset.get_target_list() returns self.targets.keys().
+            # self.targets is set below but only if "targets" in fin,
+            # which is not the case here.
+            # However, for historical reasons, the shape of seq_lengths (and seq_start)
+            # will count with this dummy target key,
+            # although the seq_lengths/seq_start values are never used.
+            # Thus, we cannot change this now, because then we couldn't handle old HDF files anymore.
             self.target_keys = ["classes"]
 
+        if "targets" in fin:
+            for k in fin["targets/labels"]:
+                if k not in self.labels:
+                    self.labels[k] = [self._decode(item) for item in fin["targets/labels"][k][...].tolist()]
+        # Note: "labels" in fin was not usd since quite a while in most HDFs
+        # (which I can tell because the code was broken and would have resulted in an exception,
+        # specifically self._decode was missing, and it used item.split("\0")[0] instead,
+        # which does not work because we get bytes since Python 3, thus the split would not work).
+        # However, SimpleHDFWriter might have written it if the provided data has it,
+        # thus those are the labels for "data", not for "classes", as we had it earlier.
+        if "labels" in fin and "data" not in self.labels and "inputs" in fin:
+            self.labels["data"] = [self._decode(item) for item in fin["labels"][...].tolist()]
+
         seq_lengths = fin[attr_seqLengths][...]  # shape (num_seqs,num_target_keys + 1)
         num_input_keys = 1 if "inputs" in fin else 0
         if len(seq_lengths.shape) == 1:
@@ -1110,7 +1125,8 @@ class SimpleHDFWriter:
         # seq_length idx represents (seq_idx,data_key_idx),
         # where data_key_idx == 0 is for the main input data,
         # and otherwise data_key_idx == 1 + sorted(self._prepared_extra).index(data_key).
-        # data_key_idx must allow for 2 entries by default, as HDFDataset assumes 'classes' by default.
+        # data_key_idx must allow for 2 entries by default,
+        # as HDFDataset assumes 'classes' by default, as long as there is no targets/data or targets/labels.
         if extend_existing_file:
             self._seq_lengths = self._file["seqLengths"]
         else:

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/returnn/frontend/array_.py

@@ -549,14 +549,16 @@ def masked_select(
         return tensor._raw_backend.masked_select(tensor, mask=mask, dims=dims, out_dim=out_dim)
     # Separate implementation for the case where we have a subset of the mask dims, specifically one single dim.
     # See https://github.com/rwth-i6/returnn/issues/1605 for discussion.
-    if len(dims) != 1:
-        # Please check https://github.com/rwth-i6/returnn/issues/1605 when you need this.
-        raise NotImplementedError(f"masked_select: dims {dims} len {len(dims)} > 1 not supported")
-    (in_dim,) = dims
-    in_dim: Dim
     mask = mask.copy_masked(mask_value=False, dims=dims)
+    if len(dims) > 1:
+        # Flatten it, in the specified order.
+        tensor, in_dim = rf.merge_dims(tensor, dims=dims)
+        mask, _ = rf.merge_dims(mask, dims=dims, out_dim=in_dim)
+    else:
+        (in_dim,) = dims
+    in_dim: Dim
     idxs = rf.cumsum(rf.cast(mask, "int32"), spatial_dim=in_dim)  # [T,B] -> idx in T' + 1
-    new_size = rf.gather(idxs, indices=in_dim.get_size_tensor() - 1, axis=in_dim)  # [B]
+    new_size = rf.gather(idxs, indices=in_dim.get_dim_value_tensor() - 1, axis=in_dim)  # [B]
     if out_dim is None:
         out_dim = Dim(new_size, name="masked_select")
     elif out_dim.dyn_size_ext is None:
@@ -620,7 +622,6 @@ def pack_padded(
     """
     assert not enforce_sorted  # not implemented yet...
     mask = rf.sequence_mask(dims, device=source.device)
-    assert mask.dims_set == set(dims)
     # Note: We could already calculate out_dim here, as follows:
     #   out_dim = Dim(rf.num_elements_of_shape(dims), name="packed")
     # This could trigger a more efficient calculation path in masked_select,

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/returnn/tf/layers/basic.py

@@ -4589,7 +4589,6 @@ class MergeDimsLayer(_ConcatInputLayer):
             or input_data.feature_dim_axis_or_unspecified is not NotSpecified
         ):
             data.feature_dim_axis = new_feature_dim_axis  # explicitly set
-            data.dim = data.batch_shape[data.feature_dim_axis] if data.feature_dim_axis is not None else None
 
         if input_data.batch_dim_axis in axes and data.batch:
             for axis in axes:

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/returnn/torch/engine.py

@@ -1086,6 +1086,9 @@ class Engine(EngineBase):
             step_idx = 0
             for extern_data_raw in data_loader:
                 step_begin_time = time.time()
+                if self._forward_step_expected_outputs:
+                    # Also resets any dyn dims, which might have been set in the prev step.
+                    self._forward_step_expected_outputs.reset_content()
                 extern_data = extern_data_util.raw_dict_to_extern_data(
                     extern_data_raw, extern_data_template=self.extern_data, device=self._device
                 )

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/tests/test_HDFDataset.py

@@ -4,7 +4,7 @@ tests for HDF dataset
 
 from __future__ import annotations
 
-from typing import Dict, Any
+from typing import Dict, Any, List, Tuple, Optional
 import os
 import sys
 import _setup_test_env  # noqa
@@ -254,6 +254,8 @@ def test_SimpleHDFWriter():
     seq_lens = seq_lens1 + seq_lens2
 
     dataset = HDFDataset(files=[fn])
+    assert dataset.get_data_keys() == ["data"]
+    assert dataset.get_target_list() == []
     reader = DatasetTestReader(dataset=dataset)
     reader.read_all()
     assert "data" in reader.data_keys  # "classes" might be in there as well, although not really correct/existing
@@ -480,6 +482,248 @@ def test_SimpleHDFWriter_swmr():
     writer.close()  # Should not matter.
 
 
+def test_SimpleHDFWriter_labels():
+    fn = get_test_tmp_file(suffix=".hdf")
+    os.remove(fn)  # SimpleHDFWriter expects that the file does not exist
+    n_dim = 3
+    writer = SimpleHDFWriter(filename=fn, dim=n_dim, labels=[" ", "a", "b"])
+    seq_lens = [2, 3]
+    writer.insert_batch(
+        inputs=numpy.random.normal(size=(len(seq_lens), max(seq_lens), n_dim)).astype("float32"),
+        seq_len=seq_lens,
+        seq_tag=["seq-%i" % i for i in range(len(seq_lens))],
+    )
+    writer.close()
+
+    old_dataset = Old2018HDFDataset()
+    old_dataset.add_file(fn)
+    old_dataset.initialize()
+    print("Old dataset:", old_dataset)
+    print("Old dataset outputs:", old_dataset.num_outputs)
+    print("Old dataset target keys:", old_dataset.target_keys)
+    old_dataset.init_seq_order(epoch=1)
+    old_dataset.load_seqs(0, 1)
+
+    dataset = HDFDataset(files=[fn])
+    print("Dataset:")
+    print("  input:", dataset.num_inputs, "x", dataset.window)
+    print("  output:", dataset.num_outputs)
+    print(" ", dataset.len_info(fast=True) or "no info")
+    print("Dataset keys:", dataset.get_data_keys())
+    print("Dataset target keys:", dataset.get_target_list())
+    print("Dataset labels:", ", ".join(f"{k!r}: {v[:3]}... len {len(v)}" for k, v in dataset.labels.items()) or "None")
+    assert dataset.get_data_keys() == ["data"]
+    assert dataset.get_target_list() == []
+
+    assert dataset.labels["data"] == [" ", "a", "b"]
+    reader = DatasetTestReader(dataset=dataset)
+    reader.read_all()
+    assert reader.data_keys == ["data"]
+    assert reader.data_sparse["data"] is False
+    assert list(reader.data_shape["data"]) == [n_dim]
+    assert reader.data_dtype["data"] == "float32"
+    assert len(seq_lens) == reader.num_seqs
+    for i, seq_len in enumerate(seq_lens):
+        assert reader.seq_lens[i]["data"] == seq_len
+
+
+class Old2018HDFDataset(CachedDataset):
+    """
+    Copied and adapted from an early RETURNN version:
+    2018-03-09: https://github.com/rwth-i6/returnn/blob/c2d8fed877022d1ac1bf68b801604733db51223e/HDFDataset.py
+
+    (Not really fully functional here though... we just use it to test whether loading works.)
+
+    Some version history:
+    2015-01-05: https://github.com/rwth-i6/returnn/blob/3be0fe0906212d1ffdd93807c0a3854a38842eb8/Dataset.py
+    2015-08-04: https://github.com/rwth-i6/returnn/blob/995e87184abc6e07256417cb533163ac0a7d7dd8/HDFDataset.py
+    2018-03-09: https://github.com/rwth-i6/returnn/blob/c2d8fed877022d1ac1bf68b801604733db51223e/HDFDataset.py
+    """
+
+    def __init__(self, **kwargs):
+        super(Old2018HDFDataset, self).__init__(cache_byte_size=100, **kwargs)
+        self.files: List[str] = []
+        self.file_start = [0]
+        self.file_seq_start: List[List[int]] = []
+        self.file_index: List[int] = []
+        self.data_dtype: Dict[str, str] = {}
+        self.data_sparse: Dict[str, bool] = {}
+
+        # Copied from old base CachedDataset
+        self._seq_lengths: List[Tuple[int, int]] = []  # uses real seq idx
+        self.tags: List[str] = []  # uses real seq idx
+        self.tag_idx: Dict[str, int] = {}  # map of tag -> real-seq-idx
+        self.targets = {}
+        self.target_keys = []
+
+        # Copied from old base Dataset
+        self._num_codesteps: Optional[int] = None  # Num output frames, could be different from input, seq2seq, ctc.
+
+    def add_file(self, filename):
+        """
+        Setups data:
+          self.seq_lengths
+          self.file_index
+          self.file_start
+          self.file_seq_start
+        Use load_seqs() to load the actual data.
+        :type filename: str
+        """
+        fin = h5py.File(filename, "r")
+        decode = lambda s: s if isinstance(s, str) else s.decode("utf-8")
+        if "targets" in fin:
+            self.labels = {
+                k: [decode(item).split("\0")[0] for item in fin["targets/labels"][k][...].tolist()]
+                for k in fin["targets/labels"]
+            }
+        if not self.labels:
+            labels = [decode(item).split("\0")[0] for item in fin["labels"][...].tolist()]
+            """:type: list[str]"""
+            self.labels = {"classes": labels}
+            assert len(self.labels["classes"]) == len(labels), (
+                "expected " + str(len(self.labels["classes"])) + " got " + str(len(labels))
+            )
+        tags = [decode(item).split("\0")[0] for item in fin["seqTags"][...].tolist()]
+        """ :type: list[str] """
+        self.files.append(filename)
+        if "times" in fin:
+            if self.timestamps is None:
+                self.timestamps = fin[attr_times][...]
+            else:
+                self.timestamps = numpy.concatenate(
+                    [self.timestamps, fin[attr_times][...]], axis=0
+                )  # .extend(fin[attr_times][...].tolist())
+        seq_lengths = fin[attr_seqLengths][...]
+        if "targets" in fin:
+            self.target_keys = sorted(fin["targets/labels"].keys())
+        else:
+            self.target_keys = ["classes"]
+
+        if len(seq_lengths.shape) == 1:
+            seq_lengths = numpy.array(zip(*[seq_lengths.tolist() for i in range(len(self.target_keys) + 1)]))
+
+        seq_start = [numpy.zeros((seq_lengths.shape[1],), "int64")]
+        if not self._seq_start:
+            self._seq_start = [numpy.zeros((seq_lengths.shape[1],), "int64")]
+        for l in seq_lengths:
+            self._seq_lengths.append(numpy.array(l))
+            seq_start.append(seq_start[-1] + l)
+        self.tags += tags
+        self.file_seq_start.append(seq_start)
+        nseqs = len(seq_start) - 1
+        for i in range(nseqs):
+            self.tag_idx[tags[i]] = i + self._num_seqs
+        self._num_seqs += nseqs
+        self.file_index.extend([len(self.files) - 1] * nseqs)
+        self.file_start.append(self.file_start[-1] + nseqs)
+        self._num_timesteps += sum([s[0] for s in seq_lengths])
+        if self._num_codesteps is None:
+            self._num_codesteps = [0 for i in range(1, len(seq_lengths[0]))]
+        for i in range(1, len(seq_lengths[0])):
+            self._num_codesteps[i - 1] += sum([s[i] for s in seq_lengths])
+        if len(fin["inputs"].shape) == 1:  # sparse
+            num_inputs = [fin.attrs[attr_inputPattSize], 1]
+        else:
+            num_inputs = [fin["inputs"].shape[1], len(fin["inputs"].shape)]  # fin.attrs[attr_inputPattSize]
+        if self.num_inputs == 0:
+            self.num_inputs = num_inputs[0]
+        assert self.num_inputs == num_inputs[0], "wrong input dimension in file %s (expected %s got %s)" % (
+            filename,
+            self.num_inputs,
+            num_inputs[0],
+        )
+        if "targets/size" in fin:
+            num_outputs = {
+                k: [fin["targets/size"].attrs[k], len(fin["targets/data"][k].shape)] for k in fin["targets/size"].attrs
+            }
+        else:
+            num_outputs = {"classes": fin.attrs["numLabels"]}
+        num_outputs["data"] = num_inputs
+        if not self.num_outputs:
+            self.num_outputs = num_outputs
+        assert self.num_outputs == num_outputs, "wrong dimensions in file %s (expected %s got %s)" % (
+            filename,
+            self.num_outputs,
+            num_outputs,
+        )
+        if "targets" in fin:
+            for name in fin["targets/data"]:
+                tdim = 1 if len(fin["targets/data"][name].shape) == 1 else fin["targets/data"][name].shape[1]
+                self.data_dtype[name] = str(fin["targets/data"][name].dtype) if tdim > 1 else "int32"
+                self.targets[name] = None
+        else:
+            self.targets = {"classes": numpy.zeros((self._num_timesteps,))}
+            self.data_dtype["classes"] = "int32"
+        self.data_dtype["data"] = fin["inputs"].dtype
+        assert len(self.target_keys) == len(self._seq_lengths[0]) - 1
+        fin.close()
+
+    def _load_seqs(self, start, end):
+        """
+        Load data sequences.
+        As a side effect, will modify / fill-up:
+          self.alloc_intervals
+          self.targets
+          self.chars
+
+        :param int start: start sorted seq idx
+        :param int end: end sorted seq idx
+        """
+        assert start < self.num_seqs
+        assert end <= self.num_seqs
+        selection = self.insert_alloc_interval(start, end)
+        assert len(selection) <= end - start, (
+            "DEBUG: more sequences requested (" + str(len(selection)) + ") as required (" + str(end - start) + ")"
+        )
+        file_info: List[List[int]] = [[] for l in range(len(self.files))]
+        # file_info[i] is (sorted seq idx from selection, real seq idx)
+        for idc in selection:
+            ids = self._seq_index[idc]
+            file_info[self.file_index[ids]].append((idc, ids))
+            self.preload_set.add(ids)
+        for i in range(len(self.files)):
+            if len(file_info[i]) == 0:
+                continue
+            print("loading file %d/%d" % (i + 1, len(self.files)), self.files[i], file=log.v4)
+            fin = h5py.File(self.files[i], "r")
+            for idc, ids in file_info[i]:
+                s = ids - self.file_start[i]
+                p = self.file_seq_start[i][s]
+                l = self._seq_lengths[ids]
+                if "targets" in fin:
+                    for k in fin["targets/data"]:
+                        if self.targets[k] is None:
+                            self.targets[k] = numpy.zeros((self._num_codesteps[self.target_keys.index(k)],)) - 1
+                        ldx = self.target_keys.index(k) + 1
+                        self.targets[k][self.get_seq_start(idc)[ldx] : self.get_seq_start(idc)[ldx] + l[ldx]] = fin[
+                            "targets/data/" + k
+                        ][p[ldx] : p[ldx] + l[ldx]]
+                self._set_alloc_intervals_data(idc, data=fin["inputs"][p[0] : p[0] + l[0]][...])
+            fin.close()
+        gc.collect()
+        assert self.is_cached(start, end)
+
+    def _get_seq_length_by_real_idx(self, real_seq_idx):
+        return self._seq_lengths[real_seq_idx]
+
+    def get_tag(self, sorted_seq_idx):
+        ids = self._seq_index[self._index_map[sorted_seq_idx]]
+        return self.tags[ids]
+
+    def is_data_sparse(self, key):
+        if key in self.num_outputs:
+            return self.num_outputs[key][1] == 1
+        if self.get_data_dtype(key).startswith("int"):
+            return True
+        return False
+
+    def get_data_dtype(self, key):
+        return self.data_dtype[key]
+
+    def len_info(self):
+        return ", ".join(["HDF dataset", "sequences: %i" % self.num_seqs, "frames: %i" % self.get_num_timesteps()])
+
+
 def dummy_iter_dataset(dataset: Dataset) -> int:
     """
     :param Dataset dataset:

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/tests/test_rf_array.py

@@ -33,6 +33,28 @@ def test_pack_padded():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_pack_padded_md():
+    batch_dim = Dim(3, name="batch")
+    hyps_dim = Dim(5, name="hyps")
+    hyps_spatial_dim = Dim(Tensor("hyps_spatial", [batch_dim, hyps_dim], dtype="int32"))
+    vocab_dim = Dim(1000, name="vocab")
+    extern_data = TensorDict(
+        {"labels": Tensor("labels", [batch_dim, hyps_dim, hyps_spatial_dim], dtype="int32", sparse_dim=vocab_dim)}
+    )
+    hyps_packed_spatial_dim = Dim(None, name="hyps_packed_spatial")
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        labels = extern_data["labels"]
+        assert hyps_spatial_dim in labels.dims
+        assert hyps_spatial_dim.dyn_size_ext.dims == (batch_dim, hyps_dim)
+        labels_, _ = rf.pack_padded(labels, dims=[hyps_dim, hyps_spatial_dim], out_dim=hyps_packed_spatial_dim)
+        assert hyps_packed_spatial_dim.dyn_size_ext.dims == (batch_dim,)
+        labels_.mark_as_default_output(shape=[batch_dim, hyps_packed_spatial_dim])
+
+    run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+
+
 def test_masked_select():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

{returnn-1.20240917.180949 → returnn-1.20240918.234323}/tools/dump-dataset.py

@@ -48,6 +48,11 @@ def dump_dataset(options):
     dataset.init_seq_order(epoch=options.epoch, seq_list=seq_list)
     print("Dataset keys:", dataset.get_data_keys(), file=log.v3)
     print("Dataset target keys:", dataset.get_target_list(), file=log.v3)
+    print(
+        "Dataset labels:",
+        ", ".join(f"{k!r}: {v[:3]}... len {len(v)}" for k, v in dataset.labels.items()) or "None",
+        file=log.v3,
+    )
     assert options.key in dataset.get_data_keys()
     max_seq_length = NumbersDict(options.max_seq_length)
     min_seq_length = NumbersDict(options.min_seq_length)

returnn-1.20240917.180949/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240917.180949'
-long_version = '1.20240917.180949+git.6964c19'