fgnt / lazy_dataset Goto Github PK

In [1] it is demonstrated to address the issue of the memory consumption, when multiprocessing is used. Although, we don't use multiprocessing (it's implemented, but threads are usually faster for audio data and avoid the fork issues mentioned in [1]), the idea can be integrated in our dataset implementation, with a small improvement for the memory consumption of large dataset. We have already a pickle based serialization, so there will ne no additional overhead.

Code from [1]:

class NumpySerializedList:
  def __init__(self, lst: list[Any]):
    lst = [np.frombuffer(pickle.dumps(x), dtype=np.uint8) for x in lst]
    self._addr = np.cumsum([len(x) for x in lst])
    self._lst = np.concatenate(lst)

  def __len__(self):
    return len(self._addr)

  def __getitem__(self, idx: int):
    start = 0 if idx == 0 else self._addr[idx - 1]
    end = self._addr[idx]
    return pickle.loads(memoryview(self._lst[start:end]))

[1] https://ppwwyyxx.com/blog/2022/Demystify-RAM-Usage-in-Multiprocess-DataLoader

The current version is a bit opaque in the sense that I do not know how to use .shuffle() in combination with .prefetch() or .map(..., num_workers=10).

Please provide some documentation of best practices.

In some cases, items should work (each entry in the dataset has a key and a value), but the dataset does not have keys (because, e.g., the length of the dataset is unknown). This currently gives a NotImplementedError:

import lazy_dataset
lazy_dataset.new({'a': 1, 'b': 2}).filter(lambda x: True).items()

NotImplementedError: keys is not implemented for <class 'lazy_dataset.core.FilterDataset'>.
self: 
    DictDataset(len=2)
  MapDataset(_pickle.loads)
FilterDataset(<function <lambda> at 0x7f5c36b3a3a0>)

It would be logical if lazy_dataset.new would accept another dataset, like this:

import lazy_dataset
lazy_dataset.new(lazy_dataset.new([1, 2, 3, 4]))

But it is unclear how new should handle datasets. It could

• return the input dataset
• make a copy of the dataset (return dataset.copy(freeze=??))
• cache and transform into ListDataset or DictDataset, i.e., lazy_dataset.new(list(lazy_dataset.new([1, 2, 3, 4]))) or lazy_dataset.new(dict(lazy_dataset.new({'a': 1, 'b': 2})))

The discussion in pull request #10 lead to this issue: The API should support method and function calls in the style of numpy, i.e.

def op(self, (d1, d2, ,...), *args):
    ...

and

def op((d1, d2, ...), *args):
    ...

This is different to the current API, which looks like

def op(self, *datasets):
    ....

and

def op(*datasets):
    ...

respectively.

At the moment .prefetch only supports datasets at the input that are indexable.
When the number of workers is 2 or more, this is necessary.
In case of one worker indexable is not necessary, iterable would be enough.

I opened this issue to collect use cases as motivation to implement this.
At the moment there is no use case.