returnn 1.20240704.121454__tar.gz → 1.20240705.144031__tar.gz

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

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

returnn-1.20240705.144031/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240705.144031'
+long_version = '1.20240705.144031+git.9ae3453'

{returnn-1.20240704.121454 → returnn-1.20240705.144031}/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.20240705.144031}/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
 

{returnn-1.20240704.121454 → returnn-1.20240705.144031}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240704.121454 → returnn-1.20240705.144031}/tests/test_Dataset.py

@@ -637,10 +637,30 @@ def test_DistributeFilesDataset_distribute_evenly_by_size():
         res_ = [[file_sizes[fn] for fn in sub_epoch] for sub_epoch in res]
         assert res_ == expected
 
+    _test([1, 10, 1, 1], 3, [[1], [10], [1, 1]])
     _test([1, 1, 78, 120], 4, [[1], [1], [78], [120]])
     _test([1, 1, 1, 56, 141], 5, [[1], [1], [1], [56], [141]])
     _test([1, 1, 1, 56, 141], 4, [[1, 1], [1], [56], [141]])
-    _test([5, 5] + [10] * 7, 5, [[5, 5, 10], [10, 10], [10, 10], [10], [10]])
+    _test([5, 5] + [10] * 7, 5, [[5, 5, 10], [10, 10], [10], [10], [10, 10]])
+    _test(
+        [1] * 29,
+        8,
+        [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]],
+    )
+
+    def _test_stddev(sizes: List[int], partition_epoch: int, max_stddev_percent: float):
+        avg_per_bin = sum(sizes) / partition_epoch
+        files = [f"file-{i}" for i in range(len(sizes))]
+        file_sizes = {f: s for f, s in zip(files, sizes)}
+        res = DistributeFilesDataset._distribute_evenly_by_size(
+            num_bins=partition_epoch, file_sizes=file_sizes, files_order=files
+        )
+        sizes = [sum([file_sizes[s] for s in v]) for v in res]
+        assert numpy.std(sizes) <= max_stddev_percent * avg_per_bin
+
+    # This test verifies the algorithm distributes evenly in easy cases
+    rng = numpy.random.RandomState(42)
+    _test_stddev(rng.uniform(low=97, high=103, size=(100,)), 20, 0.03)
 
 
 def test_DistributeFilesDataset():

returnn-1.20240704.121454/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240704.121454'
-long_version = '1.20240704.121454+git.2c0b75c'