returnn 1.20240118.15942__tar.gz → 1.20240118.110229__tar.gz

{returnn-1.20240118.15942/returnn.egg-info → returnn-1.20240118.110229}/PKG-INFO

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

returnn-1.20240118.110229/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240118.110229'
+long_version = '1.20240118.110229+git.5d7c8d6'

{returnn-1.20240118.15942 → returnn-1.20240118.110229}/returnn/datasets/audio.py

@@ -89,7 +89,6 @@ class OggZipDataset(CachedDataset2):
             # Special case (mostly for debugging) to directly access the filesystem, not via zip-file.
             self.paths = [os.path.dirname(path)]
             self._names = [os.path.basename(path)]
-            self._zip_files = None
             self._use_zip_files = False
             assert not use_cache_manager, "cache manager only for zip file"
         else:
@@ -250,6 +249,17 @@ class OggZipDataset(CachedDataset2):
 
         self._data = data
 
+    def finish_epoch(self, *, free_resources: bool = False):
+        """finish epoch"""
+        super().finish_epoch()
+        self._seq_order = None
+        self._num_seqs = 0
+        if free_resources:
+            # Basically undo the _lazy_init, such that _lazy_init would init again next time.
+            self._data = None
+            self.segments = None
+            self._zip_files = None
+
     def _read_segment_list(self, segment_file):
         """
         read a list of segment names in either plain text or gzip

{returnn-1.20240118.15942 → returnn-1.20240118.110229}/returnn/datasets/basic.py

@@ -640,7 +640,7 @@ class Dataset(object):
         self.rnd_seq_drop = Random(self._get_random_seed_for_epoch(epoch=epoch))
         return False
 
-    def finish_epoch(self):
+    def finish_epoch(self, *, free_resources: bool = False):
         """
         This would get called at the end of the epoch (currently optional only).
         After this, further calls to :func:`get_data` or :func:`load_seqs` are invalid,

{returnn-1.20240118.15942 → returnn-1.20240118.110229}/returnn/datasets/meta.py

@@ -462,14 +462,14 @@ class MetaDataset(CachedDataset2):
         """
         return self.num_total_seqs
 
-    def finish_epoch(self):
+    def finish_epoch(self, *, free_resources: bool = False):
         """
         This would get called at the end of the epoch.
         """
-        super(MetaDataset, self).finish_epoch()
+        super(MetaDataset, self).finish_epoch(free_resources=free_resources)
         for _, dataset in self.datasets.items():
             assert isinstance(dataset, Dataset)
-            dataset.finish_epoch()
+            dataset.finish_epoch(free_resources=free_resources)
 
     def _load_seqs(self, start, end):
         """

{returnn-1.20240118.15942 → returnn-1.20240118.110229}/returnn/datasets/multi_proc.py

@@ -5,6 +5,7 @@ Multi-processing dataset
 from __future__ import annotations
 from typing import Optional, Any, Dict, List
 import sys
+import gc
 from .basic import init_dataset, Dataset, DatasetSeq
 from .cached2 import CachedDataset2
 from returnn.util.basic import try_run
@@ -144,13 +145,6 @@ class MultiProcDataset(CachedDataset2):
             assert msg == "labels"
 
     def __del__(self):
-        if self._seq_order_proc:
-            # noinspection PyBroadException
-            try:
-                self._seq_order_proc_parent_conn.send(("exit", {}))
-                self._seq_order_proc.join()
-            except Exception:
-                pass
         if self._worker_procs:
             got_exception = False
             for worker_parent_conn in self._worker_parent_conns:
@@ -281,6 +275,15 @@ class MultiProcDataset(CachedDataset2):
                     got_init_seq_order = True
                     next_seq_idx = 0
                     cache[:] = []
+                elif msg == "finish_epoch":
+                    got_init_seq_order = False
+                    next_seq_idx = 0
+                    cache[:] = []
+                    if dataset:
+                        dataset.finish_epoch(**kwargs)
+                    if kwargs["free_resources"]:
+                        dataset = None
+                        gc.collect()
                 else:
                     raise Exception(f"unknown msg {msg!r}")
         except KeyboardInterrupt:  # when parent dies
@@ -336,6 +339,12 @@ class MultiProcDataset(CachedDataset2):
             return self._total_num_seqs
         raise NotImplementedError
 
+    def finish_epoch(self, *, free_resources: bool = False):
+        """finish epoch"""
+        super().finish_epoch(free_resources=free_resources)
+        for worker_parent_conn in self._worker_parent_conns:
+            worker_parent_conn.send(("finish_epoch", {"free_resources": free_resources}))
+
 
 class _SetupProcPreInit:
     def __init__(self):

{returnn-1.20240118.15942 → returnn-1.20240118.110229}/returnn/datasets/sprint.py

@@ -715,7 +715,7 @@ class ExternSprintDataset(SprintDatasetBase):
     def _partitionEpoch(self):
         return self.partition_epoch
 
-    def finish_epoch(self):
+    def finish_epoch(self, *, free_resources: bool = False):
         """
         Called at the end of the epoch.
         """
@@ -724,7 +724,7 @@ class ExternSprintDataset(SprintDatasetBase):
             super(ExternSprintDataset, self).init_seq_order(epoch=None, seq_list=None)
         # Exit child, before we overwrite anything, such as new epoch or seq_list.
         self._exit_child(wait_thread=True)
-        super(ExternSprintDataset, self).finish_epoch()
+        super(ExternSprintDataset, self).finish_epoch(free_resources=free_resources)
 
     def _exit_handler(self):
         """

{returnn-1.20240118.15942 → returnn-1.20240118.110229}/returnn/torch/engine.py

@@ -603,7 +603,18 @@ class Engine(EngineBase):
 
         loader_opts = self.config.typed_value("torch_dataloader_opts") or {}
         assert isinstance(loader_opts, dict), f"config torch_dataloader_opts, expected dict, got {type(loader_opts)}"
-        return data_pipeline.create_data_loader_from_batches(batches_dataset, loader_opts)
+
+        data_loader = data_pipeline.create_data_loader_from_batches(batches_dataset, loader_opts)
+
+        if data_loader.num_workers > 0:  # uses multi processing
+            # We are not using the dataset anymore here in the main proc,
+            # so free all resources as much as we can.
+            # https://github.com/rwth-i6/returnn/issues/1443
+            # Do this after creating the data loader - so in case it used the fork start method,
+            # it would still potentially have resources ready to use.
+            dataset.finish_epoch(free_resources=True)
+
+        return data_loader
 
     def _run_step(
         self, extern_data: TensorDict, *, train_flag: bool = False, train_func: bool, _inside_wrapped: bool = False

{returnn-1.20240118.15942 → returnn-1.20240118.110229/returnn.egg-info}/PKG-INFO

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

returnn-1.20240118.15942/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240118.015942'
-long_version = '1.20240118.015942+git.c1bf6a0'