tensorflow / tensorflow/datasets
Please support prefetch with python datasets
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 4.6k
- Forks
- 1.6k
- Avg merge
- 3h 54m
- Merged PRs (30d)
- 1
Description
Is your feature request related to a problem? Please describe.
There's a tremendous performance difference between datasets which are fully tensor end-to-end and datasets where some data wrangling happens in Python.
I was hoping to use "prefetch" to prepare data with the CPU while the GPU does work, but unfortunately, this only works if the data preparation is fully tensorflow-ish (? not sure the right term here)
Python IO operations are often exponentially slower and act as a bottleneck and prevent computers from keeping accelerators working at capacity.
Describe the solution you'd like
I wish tf.data.Dataset prefetch was more broadly compatible with prefetching of non-tensorflow vanilla-python data preparations.
Would it be possible for prefetch to use some performant C++ to sidestep Python GIL issues and juggle python data wrangling CPU processes alongside GPU training / inference without depending on such python CPU work happening in the main python driver process? I just want to be able to prefetch custom python datasets. Often there's some prep involved, not every dataset is tensorflow end-to-end.
i.e. instead of (python)->(gpu) what if it were
(python)->(cpp)
(cpp)->(python_prefetch)
(cpp)->(gpu/tpu accelerator)
Since C++ lacks a GIL, it could just run the python generator in a process which is isolated from the main python driver process' GIL, you'd still have a GIL per generator, but that's an easy fix, just run more python processes with different random seeds etc.
Describe alternatives you've considered
Torch DataLoader could be an option but it also seems to be a python-driven solution and therefore not super performant. I tried Threading, but pickling and unpickling overhead in python can be pretty bad. I think maybe C++ could run a background python process to prefetch python data in a more performant way than python could.
Additional context
Broader accessibility of custom dataset prefetching could enable new use-cases for tf.data especially in prototyping or infinite search spaces where it might not make sense to convert entire datasets to tensors in advance.
Apologies if I misunderstand the intricacies involved. I just want to prefetch datasets built from generators. I tried doing this last week but it didn't work, so hopefully I'm not mentioning an issue which is fixed and I missed a way to pull it off. It's hard to provide an example since the code in question is closed-source and quite extensive anyway. For a good example of when this might be handy, consider RL gym envs or datasets which involve making GET requests.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by examining the tf.data.Dataset prefetch and generator-related entry points; the issue does not name repository files or tests. Reproduce the reported behavior with a Python generator and prefetch, then determine the supported behavior and constraints. Done would require an agreed design and tests showing custom Python dataset preparation can be prefetched without blocking accelerator work.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, tensorflow
- Domain
- data, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100