Repository: rwth-i6/returnn

Issue:
Datasets are not pickleable
This is e.g. needed for PyTorch with a dataloader with num_workers > 1 (https://github.com/rwth-i6/returnn/issues/1120), or also other cases.



Discussion:
The datasets in question so far are ExternSprintDataset, HDFDataset and OggZipDataset
I would still strongly prefer to have the option in the code to use the dataloader with `num_workers=1`, so that we can continue testing with the current master until we have solutions.
I don't exactly understand what you want to test? There is nothing to test currently.
> I don't exactly understand what you want to test? There is nothing to test currently

Surely there is. From what I see the PyTorch backend is nearly feature-complete in the sense that you can run pure PyTorch models with it. Jingjing and I have the task to train a model with that and exporting it to ONNX, and I will also re-implement my current CTC aligner using PyTorch.
> PyTorch backend is nearly feature-complete in the sense that you can run pure PyTorch models with it.

No, not at all. We just had a simple proof-of-concept / demo ready, nothing else. We are actually far away from running pure PyTorch modules. Our current design, all our recent discussion is all only about running RETURNN frontend code. Yes, it is planned that pure PyTorch code can be easily integrated there, but so far we did not even start discussions on that. It's probably not difficult, but this needs to be discussed, and thought about, how we want to design the API for it. I just started an initial comment on this (https://github.com/rwth-i6/returnn/issues/1120#issuecomment-1481794007). You cannot really do any serious work with PyTorch in the current state. You need to expect that everything will totally break, maybe on a daily basis. Nothing is stable yet, nothing is really final. Except the things we discussed, which is mostly about the RETURNN frontend. Anyway, this discussion is for #1120. We can also discuss this in the next meeting.

Really, the dataloader is the least concern here. It doesn't make sense to make this part of the code more complex than needed.

Also, I still don't understand, you want to add workaround code here, using two different data loaders. Why? Why not just fix the datasets? This will be much easier and also makes much more sense. This should only be about 2 lines of code or so you need to add. Then they are pickleable, and then you don't need any hacks elsewhere.

This is all kind of off topic here, as this issue here is just about making the datasets pickleable, nothing else. I consider it a bug that they are not pickleable currently.


Resolving patch:
diff --git a/returnn/datasets/basic.py b/returnn/datasets/basic.py
--- a/returnn/datasets/basic.py
+++ b/returnn/datasets/basic.py
@@ -140,6 +140,7 @@ def __init__(
         self.random_seed_offset = random_seed_offset
         self.partition_epoch = partition_epoch or 1
         self.repeat_epoch = repeat_epoch or 1
+        self._seq_list_filter_file = seq_list_filter_file
         self.seq_tags_filter = set(self._load_seq_list_file(seq_list_filter_file)) if seq_list_filter_file else None
         self.unique_seq_tags = unique_seq_tags
         self._seq_order_seq_lens_file = seq_order_seq_lens_file
@@ -155,7 +156,9 @@ def __init__(
         self._estimated_num_seqs = estimated_num_seqs
         self.min_chunk_size = NumbersDict(min_chunk_size)
         self.chunking_variance = chunking_variance
+        self._chunking = chunking
         self.chunk_size, self.chunk_step, self.custom_chunking_func = self._parse_chunking(chunking)
+        self._context_window = context_window
         if isinstance(context_window, (tuple, list)):
             assert len(context_window) == 2
             for elem in context_window:
@@ -184,6 +187,7 @@ def __init__(
         assert isinstance(self.ctx_right, NumbersDict)
         self.shuffle_frames_of_nseqs = shuffle_frames_of_nseqs
         self.epoch = None
+        self.zpad = None
 
     def __repr__(self):
         return "<%s %r epoch=%s>" % (
@@ -192,6 +196,39 @@ def __repr__(self):
             getattr(self, "epoch", "<unknown>"),
         )
 
+    _getnewargs_exclude_attrs = set()  # type: typing.Set[str]
+
+    @staticmethod
+    def _create_from_reduce(cls, kwargs, state) -> Dataset:
+        """
+        :param type cls:
+        :param dict[str] kwargs:
+        :param dict[str] state:
+        :rtype: Dataset
+        """
+        assert issubclass(cls, Dataset)
+        ds = cls(**kwargs)
+        for attr, value in state.items():
+            setattr(ds, attr, value)
+        return ds
+
+    def __reduce__(self):
+        import inspect
+
+        kwargs = {}
+        for cls in self.__class__.__mro__:
+            if isinstance(cls, type) and issubclass(cls, Dataset):
+                for arg in inspect.getargs(cls.__init__.__code__).args[1:]:
+                    if arg in self._getnewargs_exclude_attrs:
+                        continue
+                    if hasattr(self, "_" + arg):
+                        kwargs[arg] = getattr(self, "_" + arg)
+                    else:
+                        kwargs[arg] = getattr(self, arg)
+
+        state = {attr: getattr(self, attr) for attr in ["epoch", "zpad"]}
+        return Dataset._create_from_reduce, (self.__class__, kwargs, state)
+
     @staticmethod
     def _get_default_random_seed_offset():
         """
diff --git a/returnn/datasets/cached.py b/returnn/datasets/cached.py
--- a/returnn/datasets/cached.py
+++ b/returnn/datasets/cached.py
@@ -11,11 +11,17 @@
 
 
 class CachedDataset(Dataset):
+    """
+    Base class for datasets with caching. This is only used for the :class:`HDFDataset`.
+    Also see :class:`CachedDataset2`.
+    """
+
     def __init__(self, cache_byte_size=0, **kwargs):
         """
         :param int cache_byte_size:
         """
         super(CachedDataset, self).__init__(**kwargs)
+        self._cache_byte_size = cache_byte_size
         self.cache_byte_size_total_limit = cache_byte_size
         if cache_byte_size == -1:
             self.cache_byte_size_limit_at_start = 1024**4
diff --git a/returnn/datasets/generating.py b/returnn/datasets/generating.py
--- a/returnn/datasets/generating.py
+++ b/returnn/datasets/generating.py
@@ -36,6 +36,8 @@ def __init__(self, input_dim, output_dim, num_seqs=float("inf"), **kwargs):
         super(GeneratingDataset, self).__init__(**kwargs)
         assert self.shuffle_frames_of_nseqs == 0
 
+        self._input_dim = input_dim
+        self._output_dim = output_dim
         self.num_inputs = input_dim
         output_dim = convert_data_dims(output_dim, leave_dict_as_is=False)
         if "data" not in output_dim and input_dim is not None:
@@ -273,6 +275,7 @@ class Task12AXDataset(GeneratingDataset):
 
     _input_classes = "123ABCXYZ"  # noqa
     _output_classes = "LR"
+    _getnewargs_exclude_attrs = Dataset._getnewargs_exclude_attrs.union(("input_dim", "output_dim"))
 
     def __init__(self, **kwargs):
         super(Task12AXDataset, self).__init__(
@@ -368,6 +371,7 @@ class TaskEpisodicCopyDataset(GeneratingDataset):
     # Blank, delimiter and some chars.
     _input_classes = " .01234567"
     _output_classes = _input_classes
+    _getnewargs_exclude_attrs = Dataset._getnewargs_exclude_attrs.union(("input_dim", "output_dim"))
 
     def __init__(self, **kwargs):
         super(TaskEpisodicCopyDataset, self).__init__(
@@ -440,6 +444,7 @@ class TaskXmlModelingDataset(GeneratingDataset):
     # Blank, XML-tags and some chars.
     _input_classes = " <>/abcdefgh"  # noqa
     _output_classes = _input_classes
+    _getnewargs_exclude_attrs = Dataset._getnewargs_exclude_attrs.union(("input_dim", "output_dim"))
 
     def __init__(self, limit_stack_depth=4, **kwargs):
         super(TaskXmlModelingDataset, self).__init__(
@@ -535,6 +540,7 @@ class TaskVariableAssignmentDataset(GeneratingDataset):
     # Blank/Delim/End, Store/Query, and some chars for key/value.
     _input_classes = " ,.SQ()abcdefgh"  # noqa
     _output_classes = _input_classes
+    _getnewargs_exclude_attrs = Dataset._getnewargs_exclude_attrs.union(("input_dim", "output_dim"))
 
     def __init__(self, **kwargs):
         super(TaskVariableAssignmentDataset, self).__init__(
@@ -662,6 +668,8 @@ class TaskNumberBaseConvertDataset(GeneratingDataset):
     Or basically convert some number from some base into another base.
     """
 
+    _getnewargs_exclude_attrs = Dataset._getnewargs_exclude_attrs.union(("input_dim", "output_dim"))
+
     def __init__(self, input_base=8, output_base=2, min_input_seq_len=1, max_input_seq_len=8, **kwargs):
         """
         :param int input_base:
@@ -846,6 +854,8 @@ class DummyDatasetMultipleDataKeys(DummyDataset):
     sequence length for each data point.
     """
 
+    _getnewargs_exclude_attrs = Dataset._getnewargs_exclude_attrs.union(("input_dim",))
+
     def __init__(
         self,
         output_dim,
@@ -1123,6 +1133,8 @@ class CopyTaskDataset(GeneratingDataset):
     Input/output is exactly the same random sequence of sparse labels.
     """
 
+    _getnewargs_exclude_attrs = Dataset._getnewargs_exclude_attrs.union(("input_dim", "output_dim"))
+
     def __init__(self, nsymbols, minlen=0, maxlen=0, minlen_epoch_factor=0, maxlen_epoch_factor=0, **kwargs):
         """
         :param int nsymbols:


Test patch:
diff --git a/tests/test_HDFDataset.py b/tests/test_HDFDataset.py
--- a/tests/test_HDFDataset.py
+++ b/tests/test_HDFDataset.py
@@ -303,23 +303,31 @@ def test_read_simple_hdf():
     seq_lens = [2, 3]
     raw_gzipped = (
         b"\x1f\x8b\x08\x00\x80\xc8f\\\x02\xff\xed\x9a=l\xd3@\x14\xc7\xefl'XQ\x8a\xd2vhTT\xf0\x84X\x90\xc2V1\xb8TJ\xa0C"
-        b'\x81\x8a\x04\xa9\x03BM\x85i"5Q\x82]$`h%\x18@\xea\xc6\x02\x1b\x03\x03\x1f\x0b\x0b\x1bm\x10S\xa7NLL\xb01\x96\x1d'
+        b'\x81\x8a\x04\xa9\x03BM\x85i"5Q\x82]$`h%\x18@'
+        b"\xea\xc6\x02\x1b\x03\x03\x1f\x0b\x0b\x1bm\x10S\xa7NLL\xb01\x96\x1d"
         b"\xa9\xd8\xbewI}\x8d\x9d\xb1i\xf9\xff\x86\\\xee\xf2\x9e\xef\xc3\xff\xbb{"
         b"\xf6\xe5\xc5\\\xf1\xeaHf2\xc3\x02L\x93\x19,"
-        b"\xc7\x0e\xb2O|\xb7\xa3y\xf9\xfb\x12\xa5\x9c\xd2\xe7\x94\xbe\xd3d\xb9\x19\xfe\x96\xa7\xf2\x1c]\xdf\xd2E\xbeF\x8e"
+        b"\xc7\x0e\xb2O|\xb7\xa3y\xf9\xfb\x12\xa5\x9c"
+        b"\xd2\xe7\x94\xbe\xd3d\xb9\x19\xfe\x96\xa7\xf2\x1c]\xdf\xd2E\xbeF\x8e"
         b"\x95[\xa5R`\xbd\xaf \xeby\x95\x12\xe9\x05\x06\xfeG\xe6J\xb3\x0bA\xbaH\xf9\x02\xa5;Z\xd4n\xb5\xba\xec\xac\xba"
         b"\x8c\xb9N{\xdei\xaex5W\x94\xd7\x9b\xad5O\x94W\xaa+nW\xaf\x83\xf44JzUu\x9de\xd3\xfe\\\t\x14{"
         b"\xda\xffn\x8a\xeb/T=\xaf\\\x7f\xec\x04:7\xfd\xe9\x14Z^\x89\xcc\x0f\x92\xbd\xefS "
-        b"\x7f3\xf4o\xae5\x8a\xf5\x86\x1b\xeb\xc7\x99ZoF\xfa\xcdS\x97\xc5\xfc\x1aX\xaf\xf4\x1f\x91\xfe\x95z\xc3q=\xa7\xe5"
-        b"&\xf9\xa7\xe2\xdb]v\xda\xf1\xed\xee\xdd\x9e|\xe28s\x96\x16>\\\xe6\x85=\xe7\xbc\xaf\xbdN\xeb\xca8\x17mKQ^\xd3"
-        b'\xb4\xd0\xc1$\x7f\x9d\xab+\x89`\x8cZ\x1b\x18o\xec\xdc\xbf\xd3\xbb\xc3\xc3A\xf9\xc6\xcd"\xf7G\xda\x92:\x1dK\xb67'
-        b"\xe5\xfak$\xdb\xc9\xd5\xfdg:\xd9N\xce\x8b\xd6\xd4\xf1^7\x0e\xebJ\xf4\x8c\x0b\x99t\xf5\xa9)\xfb\x9d\xd6\xd5Y.t\r"
-        b"\xf4-\x86\xd6\xa2\xf9@z\xd3YTo\x9a\xbew\x8a|\rqY\xa3\xbf\xdeZG2\x1e\xc1>\xcb\xfb\xed\xb3V\xb2\xdf\xe6Y\xa5\xc0"
+        b"\x7f3\xf4o\xae5\x8a\xf5\x86\x1b\xeb\xc7\x99ZoF\xfa\xcdS"
+        b"\x97\xc5\xfc\x1aX\xaf\xf4\x1f\x91\xfe\x95z\xc3q=\xa7\xe5"
+        b"&\xf9\xa7\xe2\xdb]v\xda\xf1\xed\xee\xdd\x9e|\xe28s\x96\x16"
+        b">\\\xe6\x85=\xe7\xbc\xaf\xbdN\xeb\xca8\x17mKQ^\xd3"
+        b"\xb4\xd0\xc1$\x7f\x9d\xab+\x89`\x8cZ\x1b\x18o\xec\xdc\xbf\xd3"
+        b'\xbb\xc3\xc3A\xf9\xc6\xcd"\xf7G\xda\x92:\x1dK\xb67'
+        b"\xe5\xfak$\xdb\xc9\xd5\xfdg:\xd9N\xce\x8b\xd6\xd4\xf1^7\x0e\xeb"
+        b"J\xf4\x8c\x0b\x99t\xf5\xa9)\xfb\x9d\xd6\xd5Y.t\r"
+        b"\xf4-\x86\xd6\xa2\xf9@z\xd3YTo\x9a\xbew\x8a|\rqY\xa3\xbf\xde"
+        b"ZG2\x1e\xc1>\xcb\xfb\xed\xb3V\xb2\xdf\xe6Y\xa5\xc0"
         b"\x88\x8e\x9b\x81-\n\x00\x00\x00\x00\xe0D1("
         b"\x8eN)\xcf\x99j|\xa9\xfb\xf1q`9j\x9d\xeb\xc6\xd1\x13&\x9bX\xef>_\xc6\xc6\xd3\xd3\xe3\xbdPS\x8f\x8f\xa7kG\x10G"
         b"\xeb\x87\xfa\x99\x1f\xe0\xb7iG\x9f\x86\xb5\x18\xbb\xb7\x8a]\\|\xfd\xc5\x8e\xe6\xd3\xca}@\\\x0e\x00\x00\x00"
         b"\x000lqt\xf4\x9cC}\x1f}\xf0\x9c#\xcd\x92\xce96\xe8\rm6R\xdf\xb0\x9fs\xb8N\xfbb!\xfc\xbc\x14\xf6Y\x06\xf9:\xa4"
-        b"\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x1cK>\xde>\xdf\xb9\xf6\xc3\xec\xfc"
+        b"\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+        b"\x00\x00\x00\x00\x1cK>\xde>\xdf\xb9\xf6\xc3\xec\xfc"
         b"}\xf0\xc7\xbe7\xf5m\xe6\xf2\xe7\xa7\xf6^\xee\xc9\xd6\xefG\xb3\x9d\x87\xcf\x9c\xed_\x1f&\xbf\xbe\xd4\xcet\xae"
         b'\xaf\xbf\x99\xc9~\xda\xdd\xaa/\xbf\xb7\xef.\xeen\xbf^2"\xff\x82\xfd\x07\xf7a\x8c\xa2\xd4>\x00\x00'
     )
@@ -468,9 +476,10 @@ def test_SimpleHDFWriter_swmr():
     writer.close()  # Should not matter.
 
 
-def dummy_iter_dataset(dataset):
+def dummy_iter_dataset(dataset: Dataset) -> int:
     """
     :param Dataset dataset:
+    :return: num seqs
     """
     dataset.init_seq_order(epoch=1)
     data_keys = dataset.get_data_keys()
@@ -482,6 +491,7 @@ def dummy_iter_dataset(dataset):
             dataset.get_tag(seq_idx)
         seq_idx += 1
     print("Iterated through %r, num seqs %i" % (dataset, seq_idx))
+    return seq_idx
 
 
 def test_hdf_simple_iter():
@@ -715,6 +725,36 @@ def test_HDFDataset_no_cache_efficiency():
     # TODO... check alloc intervals etc
 
 
+def test_HDFDataset_pickle():
+    hdf_fn = generate_hdf_from_other({"class": "Task12AXDataset", "num_seqs": 23})
+    ds = HDFDataset(files=[hdf_fn], cache_byte_size=0)
+    ds.initialize()
+    ds.init_seq_order(epoch=1)
+
+    import pickle
+
+    ds_copy = pickle.loads(pickle.dumps(ds))
+    assert ds_copy.epoch == 1
+    num_seqs = dummy_iter_dataset(ds)
+    num_seqs_copy = dummy_iter_dataset(ds_copy)
+    assert num_seqs == num_seqs_copy
+
+
+def test_HDFDataset_deepcopy():
+    hdf_fn = generate_hdf_from_other({"class": "Task12AXDataset", "num_seqs": 23})
+    ds = HDFDataset(files=[hdf_fn], cache_byte_size=0)
+    ds.initialize()
+    ds.init_seq_order(epoch=1)
+
+    import copy
+
+    ds_copy = copy.deepcopy(ds)
+    assert ds_copy.epoch == 1
+    num_seqs = dummy_iter_dataset(ds)
+    num_seqs_copy = dummy_iter_dataset(ds_copy)
+    assert num_seqs == num_seqs_copy
+
+
 def test_siamese_triplet_sampling():
     datasets_path = generate_dummy_hdf(3)
     dataset = SiameseHDFDataset(input_stream_name="features", seq_label_stream="classes", files=datasets_path)
diff --git a/tests/test_PTDataset.py b/tests/test_PTDataset.py
--- a/tests/test_PTDataset.py
+++ b/tests/test_PTDataset.py
@@ -7,19 +7,17 @@
 import torch.utils.data.datapipes as dp
 from torchdata.dataloader2 import DataLoader2
 
+from returnn.datasets.basic import Dataset
 from returnn.datasets.generating import Task12AXDataset
 from returnn.torch.data import pipeline as data_pipeline
 from returnn.torch.data import returnn_dataset_wrapper
 from returnn.util import better_exchook
 
 
-def test_pipeline_serialization():
-    dataset = Task12AXDataset(num_seqs=1000)
-
+def get_loader_from_returnn_dataset(dataset: Dataset, mp_manager: torch.multiprocessing.Manager) -> DataLoader2:
     # Follow mostly similar logic as in the PT engine.
 
-    _mp_manager = torch.multiprocessing.Manager()
-    epoch_mp_shared = _mp_manager.Value("i", 0)
+    epoch_mp_shared = mp_manager.Value("i", 0)
     epoch_mp_shared.value = 1
     reset_callback = returnn_dataset_wrapper.ReturnnDatasetResetMpSharedEpochCallback(
         dataset=dataset, epoch_mp_shared=epoch_mp_shared
@@ -46,9 +44,36 @@ def test_pipeline_serialization():
 
     pickle.loads(pickle.dumps(batches_dataset))
 
-    # Now test the actual usage of the dataset.
-
     loader = DataLoader2(batches_dataset)
+    return loader
+
+
+def test_pipeline_serialization():
+    dataset = Task12AXDataset(num_seqs=1000)
+
+    mp_manager = torch.multiprocessing.Manager()
+    loader = get_loader_from_returnn_dataset(dataset, mp_manager)
+
+    c = 0
+    n = 3
+    for batch in loader:
+        print(batch)
+        c += 1
+        if c >= n:
+            break
+
+    assert c == n
+
+
+def test_HDFDataset():
+    # https://github.com/rwth-i6/returnn/issues/1281
+    from test_HDFDataset import generate_hdf_from_other, HDFDataset
+
+    hdf_fn = generate_hdf_from_other({"class": "Task12AXDataset", "num_seqs": 23})
+    hdf_dataset = HDFDataset(files=[hdf_fn], cache_byte_size=0)
+
+    mp_manager = torch.multiprocessing.Manager()
+    loader = get_loader_from_returnn_dataset(hdf_dataset, mp_manager)
     c = 0
     n = 3
     for batch in loader: