Repository: rwth-i6/returnn

Issue:
HdfDataset breaks for partition_epoch > 1
When I specify a `HdfDataset` using

    train = {
        "class": "HDFDataset",
        "files": ["mnist-data/mnist.train.h5"],
        "partition_epoch": 2,
    }

it runs fine for the first sub-epoch, but then aborts with a `KeyError`:

```
epoch 1 score: 1.0397028114570483 error: 0.27679999899111274 elapsed: 0:00:02
Reinitialize dataset seq order for epoch 2.
EXCEPTION
Traceback (most recent call last):
  File "rnn.py", line 655, in <module>
    line: main(sys.argv)
    locals:
      main = <local> <function main at 0x7fea9ee9b400>
      sys = <local> <module 'sys' (built-in)>
      sys.argv = <local> ['rnn.py', 'mnist.config', '++num_epochs', '20', '++random_seed', '42'], len = 6, _[0]: {len = 6}
  File "rnn.py", line 643, in main
    line: execute_main_task()
    locals:
      execute_main_task = <global> <function execute_main_task at 0x7fea9ee9b2f0>
  File "rnn.py", line 453, in execute_main_task
    line: engine.train()
    locals:
      engine = <global> <TFEngine.Engine object at 0x7fea3399d198>
      engine.train = <global> <bound method Engine.train of <TFEngine.Engine object at 0x7fea3399d198>>
  File "/home/christopher/Documents/returnn-current/TFEngine.py", line 1216, in train
    line: if self.train_data.init_seq_order(epoch=self.epoch):
    locals:
      self = <local> <TFEngine.Engine object at 0x7fea3399d198>
      self.train_data = <local> <HDFDataset 'train' epoch=2>
      self.train_data.init_seq_order = <local> <bound method CachedDataset.init_seq_order of <HDFDataset 'train' epoch=2>>
      epoch = <local> 2
      self.epoch = <local> 2
  File "/home/christopher/Documents/returnn-current/CachedDataset.py", line 108, in init_seq_order
    line: self._index_map = [self._seq_index_inv[i] for i in seq_index]  # sorted seq idx -> seq_index idx
    locals:
      self = <local> <HDFDataset 'train' epoch=2>
      self._index_map = <local> range(0, 2500), len = 2500
      self._seq_index_inv = <local> {6795: 0, 21588: 1, 48095: 2, 14907: 3, 6883: 4, 29592: 5, 48984: 6, 12277: 7, 21455: 8, 36598: 9, 24935: 10, 8745: 11, 7927: 12, 36357: 13, 16825: 14, 43859: 15, 39488: 16, 14135: 17, 15715: 18, 46115: 19, 26501: 20, 47059: 21, 22178: 22, 23250: 23, 18267: 24, 38287: 25, 13426: 26, 7187: 27, 312..., len = 2500
      i = <not found>
      seq_index = <local> [14301, 40929, 6175, 26822, 21473, 9281, 1025, 9092, 34920, 16578, 12840, 23479, 40027, 43267, 35285, 44475, 23986, 22221, 23289, 10292, 37317, 26857, 21719, 10678, 12913, 11157, 42460, 10067, 22173, 32143, 11908, 16189, 16298, 24477, 28341, 29458, 40448, 32751, 3289, 41541, 11726, 17672, 36568, ..., len = 2500
  File "/home/christopher/Documents/returnn-current/CachedDataset.py", line 108, in <listcomp>
    line: self._index_map = [self._seq_index_inv[i] for i in seq_index]  # sorted seq idx -> seq_index idx
    locals:
      self = <local> <HDFDataset 'train' epoch=2>
      self._index_map = <local> range(0, 2500), len = 2500
      self._seq_index_inv = <local> {6795: 0, 21588: 1, 48095: 2, 14907: 3, 6883: 4, 29592: 5, 48984: 6, 12277: 7, 21455: 8, 36598: 9, 24935: 10, 8745: 11, 7927: 12, 36357: 13, 16825: 14, 43859: 15, 39488: 16, 14135: 17, 15715: 18, 46115: 19, 26501: 20, 47059: 21, 22178: 22, 23250: 23, 18267: 24, 38287: 25, 13426: 26, 7187: 27, 312..., len = 2500
      i = <local> 14301
      seq_index = <not found>
KeyError: 14301
```

I looked into it and noticed that in the [CachedDataset](https://github.com/rwth-i6/returnn/blob/master/CachedDataset.py#L97), `self._seq_index` still contains the sequence ids of the previous epoch, when `seq_index` already changed to those of the next epoch.

https://github.com/rwth-i6/returnn/blob/master/CachedDataset.py#L97

But I am not sure what the expected behavior is.


Discussion:
I just added a new test `test_HDFDataset_partition_epoch`, and it works just fine.
I assume thought that you might use `cache_size` (`cache_byte_size` option for `HDFDataset`) (if you do not set it explicitly, it will be used by default)? That behavior is more complicated, and likely buggy. Also, I would not recommend it anyway (there is already buffering elsewhere implemented, and this cache is worse).
If that is indeed the problem, my suggestion for now is that we make an `assert partition_epoch == 1` if the cache is enabled.
Indeed, if I set `cache_byte_size = 0`, it works.
I'll create a PR for this assertion.

Resolving patch:
diff --git a/HDFDataset.py b/HDFDataset.py
--- a/HDFDataset.py
+++ b/HDFDataset.py
@@ -36,6 +36,8 @@ def __init__(self, files=None, use_cache_manager=False, **kwargs):
     :param bool use_cache_manager: uses :func:`Util.cf` for files
     """
     super(HDFDataset, self).__init__(**kwargs)
+    assert self.partition_epoch == 1 or self.cache_byte_size_total_limit == 0, \
+      "To use partition_epoch in HDFDatasets, disable caching by setting cache_byte_size=0"
     self._use_cache_manager = use_cache_manager
     self.files = []  # type: typing.List[str]  # file names
     self.h5_files = []  # type: typing.List[h5py.File]


Test patch: