returnn 1.20241015.225231__tar.gz → 1.20241016.173036__tar.gz

{returnn-1.20241015.225231 → returnn-1.20241016.173036}/PKG-INFO

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

returnn-1.20241016.173036/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241016.173036'
+long_version = '1.20241016.173036+git.0035a76'

{returnn-1.20241015.225231 → returnn-1.20241016.173036}/returnn/datasets/audio.py

@@ -406,6 +406,21 @@ class OggZipDataset(CachedDataset2):
         self._lazy_init()
         return len(self._data)
 
+    def get_data_dim(self, key: str) -> int:
+        """:return: dim of data entry with `key`"""
+        if key == "data":
+            assert self.feature_extractor is not None
+            return self.feature_extractor.get_feature_dimension()
+        elif key == "classes":
+            assert self.targets is not None
+            return self.targets.num_labels
+        elif key == "raw":
+            return 0
+        elif key == "orth":
+            return 256
+        else:
+            raise ValueError(f"{self}: unknown data key: {key}")
+
     def get_data_dtype(self, key: str) -> str:
         """:return: dtype of data entry with `key`"""
         if key == "data":

{returnn-1.20241015.225231 → returnn-1.20241016.173036}/returnn/datasets/lm.py

@@ -7,7 +7,7 @@ and some related helpers.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Callable, List, Tuple, BinaryIO, cast
+from typing import Optional, Union, Callable, Iterator, List, Tuple, BinaryIO, cast
 import typing
 import os
 import sys
@@ -394,7 +394,16 @@ class LmDataset(CachedDataset2):
         """
         :rtype: list[str]
         """
-        return sorted(self.num_outputs.keys())
+
+        def _data_keys() -> Iterator[str]:
+            # return keys in alphabetically sorted
+            if self.add_delayed_seq_data:
+                yield "delayed"
+            yield "data"
+            for i in range(self.add_random_phone_seqs):
+                yield f"random{i}"
+
+        return list(_data_keys())
 
     def get_target_list(self):
         """
@@ -1626,6 +1635,12 @@ class TranslationDataset(CachedDataset2):
         """
         return True  # all is sparse
 
+    def get_data_dim(self, _key: str) -> int:
+        """
+        :return: the data dim of data entry `_key`
+        """
+        return self.num_inputs  # same dim for all keys
+
     def get_data_dtype(self, key):
         """
         :param str key:

{returnn-1.20241015.225231 → returnn-1.20241016.173036}/returnn/datasets/meta.py

@@ -1001,7 +1001,8 @@ class CombinedDataset(CachedDataset2):
         self.num_outputs = self.data_dims
 
         self.data_dtypes = {
-            data_key: _select_dtype(data_key, self.data_dims, data_dtypes) for data_key in self.data_keys
+            data_key: _select_dtype(data_key, dset_data_key, self.datasets[dset_key], data_dtypes)
+            for (dset_key, dset_data_key), data_key in data_map.items()
         }
 
         self.dataset_seq_idx_boundaries = None  # type: typing.Optional[typing.List[int]]
@@ -2086,13 +2087,16 @@ class AnythingDataset(Dataset):
         return self._data_keys[key].get("sparse", False)
 
 
-def _select_dtype(key, data_dims, data_dtypes):
-    if data_dtypes and key in data_dtypes:
-        v = data_dtypes[key]
+def _select_dtype(
+    data_key: str,
+    dataset_data_key: str,
+    dataset: Dataset,
+    existing_dtype_map: Optional[Dict[str, str]],
+):
+    if existing_dtype_map and data_key in existing_dtype_map:
+        v = existing_dtype_map[data_key]
         assert isinstance(v, str)  # e.g. "int32" or "float32"
         return v
-    assert key in data_dims
-    if data_dims[key][1] == 1:  # sparse
-        return "int32"  # standard for 1-of-k
-    else:
-        return "float32"  # standard otherwise
+    if dataset.is_data_sparse(dataset_data_key):
+        return "int32"  # standard 1-of-k
+    return "float32"  # default float32 otherwise

{returnn-1.20241015.225231 → returnn-1.20241016.173036}/returnn/torch/frontend/_backend.py

@@ -648,24 +648,34 @@ class TorchBackend(Backend[torch.Tensor]):
         assert targets.sparse_dim and targets.sparse_dim.dimension <= logits.feature_dim.dimension
         # PyTorch expects the logits to be of shape (T, B, C) where T is the input spatial dim.
         batch_dims = logits.remaining_dims((input_spatial_dim, logits.feature_dim))
+        batch_shape = [d.get_dim_value() for d in batch_dims]
+        batch_n_elems = prod(batch_shape)
         logits = logits.copy_transpose([input_spatial_dim] + batch_dims + [logits.feature_dim])
         logits_raw: torch.Tensor = logits.raw_tensor
-        input_lengths = input_spatial_dim.dyn_size_ext.copy_transpose(batch_dims).raw_tensor
+        input_lengths: torch.Tensor = input_spatial_dim.dyn_size_ext.copy_compatible_to_dims_raw(batch_dims)
+        if input_lengths.numel() != batch_n_elems:
+            input_lengths = input_lengths.expand(batch_shape)
         logits_raw_shape = logits_raw.shape  # [T, B..., C]
         if len(batch_dims) != 1:
-            logits_raw = torch.reshape(logits_raw, logits_raw.shape[:1] + (-1,) + logits_raw.shape[-1:])  # [T, B', C]
-            input_lengths = torch.reshape(input_lengths, (-1,))  # [B']
+            logits_raw = torch.reshape(
+                logits_raw, logits_raw.shape[:1] + (batch_n_elems,) + logits_raw.shape[-1:]
+            )  # [T, B', C]
+            input_lengths = torch.reshape(input_lengths, (batch_n_elems,))  # [B']
         if logits_normalized:
             log_probs = logits_raw
         else:
             log_probs = torch.nn.functional.log_softmax(logits_raw, dim=-1)
         # PyTorch expects the targets to be of shape (B, S) where S is the targets spatial dim.
-        targets = targets.copy_transpose(batch_dims + [targets_spatial_dim])
-        targets_raw = targets.raw_tensor  # [B..., S]
-        targets_lengths = targets_spatial_dim.dyn_size_ext.copy_transpose(batch_dims).raw_tensor
+        targets_raw = targets.copy_compatible_to_dims_raw(batch_dims + [targets_spatial_dim])  # [B..., S]
+        targets_raw_shape = batch_shape + [targets_spatial_dim.get_dim_value()]
+        if targets_raw.numel() != prod(targets_raw_shape):
+            targets_raw = targets_raw.expand(targets_raw_shape)
+        targets_lengths = targets_spatial_dim.dyn_size_ext.copy_compatible_to_dims_raw(batch_dims)
+        if targets_lengths.numel() != batch_n_elems:
+            targets_lengths = targets_lengths.expand(batch_shape)
         if len(batch_dims) != 1:
-            targets_raw = torch.reshape(targets_raw, (-1, targets_raw.shape[-1]))  # [B', S]
-            targets_lengths = torch.reshape(targets_lengths, (-1,))  # [B']
+            targets_raw = torch.reshape(targets_raw, (batch_n_elems, targets_raw.shape[-1]))  # [B', S]
+            targets_lengths = torch.reshape(targets_lengths, (batch_n_elems,))  # [B']
         loss_raw = torch.nn.functional.ctc_loss(
             log_probs=log_probs,
             targets=targets_raw,

{returnn-1.20241015.225231 → returnn-1.20241016.173036}/returnn.egg-info/PKG-INFO

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

{returnn-1.20241015.225231 → returnn-1.20241016.173036}/tests/test_rf_base.py

@@ -586,6 +586,76 @@ def test_weight_noise():
     conv(rf.random_normal([time_dim, in_dim]), in_spatial_dim=time_dim)
 
 
+def test_ctc_loss():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    target_time_dim = Dim(Tensor("target_time", [batch_dim], dtype="int32"))
+    out_dim = Dim(11, name="classes")
+    out_wb_dim = out_dim + 1
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, out_wb_dim], dtype="float32", feature_dim=out_wb_dim),
+            "classes": Tensor("classes", [batch_dim, target_time_dim], dtype="int32", sparse_dim=out_dim),
+        }
+    )
+
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        logits = extern_data["data"]
+        targets = extern_data["classes"]
+        loss = rf.ctc_loss(
+            logits=logits,
+            targets=targets,
+            input_spatial_dim=time_dim,
+            targets_spatial_dim=target_time_dim,
+            blank_index=out_wb_dim.dimension - 1,
+        )
+        loss.mark_as_default_output()
+
+    run_model(
+        extern_data,
+        lambda **_kwargs: rf.Module(),
+        _forward_step,
+        dyn_dim_min_sizes={time_dim: 4, target_time_dim: 2},
+        dyn_dim_max_sizes={time_dim: 11, target_time_dim: 5},
+    )
+
+
+def test_ctc_loss_broadcast():
+    branch_dim = Dim(3, name="branch")
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    target_time_dim = Dim(Tensor("target_time", [batch_dim], dtype="int32"))
+    out_dim = Dim(11, name="classes")
+    out_wb_dim = out_dim + 1
+    extern_data = TensorDict(
+        {
+            "data": Tensor(
+                "data", [batch_dim, branch_dim, time_dim, out_wb_dim], dtype="float32", feature_dim=out_wb_dim
+            ),
+            "classes": Tensor("classes", [batch_dim, target_time_dim], dtype="int32", sparse_dim=out_dim),
+        }
+    )
+
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        logits = extern_data["data"]
+        targets = extern_data["classes"]
+        loss = rf.ctc_loss(
+            logits=logits,
+            targets=targets,
+            input_spatial_dim=time_dim,
+            targets_spatial_dim=target_time_dim,
+            blank_index=out_wb_dim.dimension - 1,
+        )
+        loss.mark_as_default_output()
+
+    run_model(
+        extern_data,
+        lambda **_kwargs: rf.Module(),
+        _forward_step,
+        test_tensorflow=False,
+        dyn_dim_min_sizes={time_dim: 4, target_time_dim: 2},
+        dyn_dim_max_sizes={time_dim: 11, target_time_dim: 5},
+    )
+
+
 def test_edit_distance():
     import numpy
     import torch

returnn-1.20241015.225231/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241015.225231'
-long_version = '1.20241015.225231+git.61705b1'