returnn 1.20240704.121454__tar.gz → 1.20240709.122157__tar.gz

{returnn-1.20240704.121454 → returnn-1.20240709.122157}/PKG-INFO

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

returnn-1.20240709.122157/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240709.122157'
+long_version = '1.20240709.122157+git.a0f8be5'

{returnn-1.20240704.121454 → returnn-1.20240709.122157}/returnn/datasets/distrib_files.py

@@ -6,7 +6,7 @@ https://github.com/rwth-i6/returnn/issues/1519
 
 from __future__ import annotations
 
-from typing import Optional, Any, Dict, Tuple, List, Sequence, Callable, Union
+from typing import Union, Optional, Any, Callable, Sequence, Tuple, List, Dict
 import os
 import sys
 import numpy
@@ -355,13 +355,17 @@ class DistributeFilesDataset(CachedDataset2):
         assert len(files_per_bin) == num_bins
         bin_idx = 0
         size_taken = 0
+        total_size_taken = 0
         for i, f_tree in enumerate(files_order):
             size = file_sizes[_get_key_for_file_tree(f_tree)]
             num_remaining = len(files_order) - i
+            total_size_taken += size
+
             if num_remaining <= num_bins - bin_idx - 1:
                 # All remaining sub epochs must be filled.
                 assert files_per_bin[bin_idx]
                 bin_idx += 1
+                avg_size_per_sub_epoch = (total_size - total_size_taken) / (num_bins - bin_idx)
                 files_per_bin[bin_idx].append(f_tree)
                 size_taken = size
                 continue
@@ -370,7 +374,6 @@ class DistributeFilesDataset(CachedDataset2):
                 files_per_bin[bin_idx].append(f_tree)
                 size_taken += size
                 continue
-            assert size_taken <= avg_size_per_sub_epoch
             if size_taken + size <= avg_size_per_sub_epoch:
                 files_per_bin[bin_idx].append(f_tree)
                 size_taken += size
@@ -388,6 +391,7 @@ class DistributeFilesDataset(CachedDataset2):
                 files_per_bin[bin_idx + 1].append(f_tree)
                 size_taken = size
             bin_idx += 1
+            avg_size_per_sub_epoch = (total_size - total_size_taken) / (num_bins - bin_idx)
         assert all(files_per_bin)
         return files_per_bin
 

{returnn-1.20240704.121454 → returnn-1.20240709.122157}/returnn/datasets/meta.py

@@ -1391,6 +1391,7 @@ class ConcatSeqsDataset(CachedDataset2):
         seq_tag_delim=";",
         remove_in_between_postfix=None,
         repeat_in_between_last_frame_up_to_multiple_of=None,
+        pad_narrow_data_to_multiple_of_target_len=None,
         use_cache_manager=False,
         epoch_wise_filter=None,
         **kwargs,
@@ -1406,6 +1407,12 @@ class ConcatSeqsDataset(CachedDataset2):
           Now it could happen that ceildiv(data_len1 + data_len2, 6) < align_len1 + align_len2.
           This option would repeat intermediate ending frames such that data_len1 % 6 == 0,
           by setting it to {"data": 6}.
+        :param dict[str,(str,int)]|None pad_narrow_data_to_multiple_of_target_len: data_key -> (target_key, multiple).
+            Similar as repeat_in_between_last_frame_up_to_multiple_of, but works for more padding/alignment schemes.
+            Example: align_len == ceildiv(data_len - P, F) for all your sub-sequences, where P is a custom number,
+            repeat_in_between_last_frame_up_to_multiple_of would not work because align_len != ceildiv(data_len, F)
+            This option would pad/narrow so that align_len * F == data_len for all but the last sub-sequences
+            by setting it to {"data": ("classes", F)} to ensure concat_align_len == ceildiv(concat_data_len - P, F)
         :param bool use_cache_manager:
         :param dict[(int,int),dict] epoch_wise_filter: see :class:`EpochWiseFilter`
         """
@@ -1413,6 +1420,7 @@ class ConcatSeqsDataset(CachedDataset2):
         self.seq_tag_delim = seq_tag_delim
         self.remove_in_between_postfix = remove_in_between_postfix or {}
         self.repeat_in_between_last_frame_up_to_multiple_of = repeat_in_between_last_frame_up_to_multiple_of or {}
+        self.pad_narrow_data_to_multiple_of_target_len = pad_narrow_data_to_multiple_of_target_len or {}
         self.epoch_wise_filter = EpochWiseFilter(epoch_wise_filter) if epoch_wise_filter else None
         if isinstance(dataset, dict):
             dataset = dataset.copy()
@@ -1486,7 +1494,7 @@ class ConcatSeqsDataset(CachedDataset2):
             sub_seq_list.extend(sub_seq_tags)
         assert sub_seq_idx == len(sub_seq_list) and len(seq_list) == len(sub_seq_idxs)
         self.cur_sub_seq_idxs = sub_seq_idxs
-        return self.sub_dataset.init_seq_order(seq_list=sub_seq_list)
+        return self.sub_dataset.init_seq_order(epoch=epoch, seq_list=sub_seq_list)
 
     def supports_seq_order_sorting(self) -> bool:
         """supports sorting"""
@@ -1539,6 +1547,11 @@ class ConcatSeqsDataset(CachedDataset2):
                     key,
                     sub_dataset_keys,
                 )
+            for key in self.pad_narrow_data_to_multiple_of_target_len:
+                assert key in sub_dataset_keys, (
+                    f"{self}: pad_narrow_data_to_multiple_of_target_len key {key}"
+                    f" not in sub dataset data-keys {sub_dataset_keys}"
+                )
         for sub_seq_idx, sub_seq_tag in zip(sub_seq_idxs, sub_seq_tags):
             self.sub_dataset.load_seqs(sub_seq_idx, sub_seq_idx + 1)
             sub_dataset_tag = self.sub_dataset.get_tag(sub_seq_idx)
@@ -1562,6 +1575,17 @@ class ConcatSeqsDataset(CachedDataset2):
                     if data.shape[0] % multiple != 0:
                         data = numpy.concatenate([data] + [data[-1:]] * (multiple - data.shape[0] % multiple), axis=0)
                         assert data.shape[0] % multiple == 0
+                if key in self.pad_narrow_data_to_multiple_of_target_len and sub_seq_idx != sub_seq_idxs[-1]:
+                    target_key, multiple = self.pad_narrow_data_to_multiple_of_target_len[key]
+                    target_data = self.sub_dataset.get_data(sub_seq_idx, target_key)
+                    len_diff = data.shape[0] - target_data.shape[0] * multiple
+                    if len_diff > 0:
+                        # if data longer than ref_data * frame_rate, narrow the data
+                        data = data[:-len_diff]
+                    elif len_diff < 0:
+                        # if data shorter than ref_data * frame_rate, pad by repeating last frame
+                        data = numpy.concatenate([data] + [data[-1:]] * -len_diff, axis=0)
+                    assert data.shape[0] == target_data.shape[0] * multiple
                 features[key].append(data)
         features = {key: numpy.concatenate(values, axis=0) for (key, values) in features.items()}
         return DatasetSeq(seq_idx=seq_idx, seq_tag=seq_tag, features=features)
@@ -1606,6 +1630,10 @@ class ConcatSeqsDataset(CachedDataset2):
         """
         return self.sub_dataset.get_data_shape(key)
 
+    def get_total_num_seqs(self) -> int:
+        """total num seqs"""
+        return len(self.full_seq_list)
+
 
 class ChunkShuffleDataset(CachedDataset2):
     """

{returnn-1.20240704.121454 → returnn-1.20240709.122157}/returnn/torch/engine.py

@@ -929,7 +929,8 @@ class Engine(EngineBase):
 
         # keep only the last two optimizer states (two in case one file gets corrupted)
         clean_epoch = self.epoch - 2
-        if clean_epoch > 0:
+        apply_keep_logic = self.config.bool("apply_cleanup_old_models_to_optim_states", False)
+        if clean_epoch > 0 and not apply_keep_logic:
             filename = self.get_epoch_model_filename(epoch=clean_epoch) + ".opt" + util.get_model_filename_postfix()
             if os.path.isfile(filename):
                 os.unlink(filename)
@@ -1078,12 +1079,14 @@ class Engine(EngineBase):
         # They consist of a file with the extension ".pt" and potentially an optimizer state with extension ".opt.pt"
 
         count_bytes = 0
-        assert os.path.exists(filename + ".pt")
-        count_bytes += os.stat(filename + ".pt").st_size
-        os.remove(filename + ".pt")
-        if os.path.exists(filename + "opt.pt"):
-            count_bytes += os.stat(filename + "opt.pt").st_size
-            os.remove(filename + "opt.pt")
+        fname = filename + ".pt"
+        assert os.path.exists(fname)
+        count_bytes += os.stat(fname).st_size
+        os.remove(fname)
+        opt_fname = filename + ".opt.pt"
+        if os.path.exists(opt_fname):
+            count_bytes += os.stat(opt_fname).st_size
+            os.remove(opt_fname)
         assert count_bytes > 0
         return count_bytes