Update Prefetcher and Implement PinMemory IterDataPipe#1014
Update Prefetcher and Implement PinMemory IterDataPipe#1014ejguan wants to merge 8 commits intometa-pytorch:mainfrom
Conversation
f42a19d to
38fe408
Compare
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
1 similar comment
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
I need to rebase onto #1003 after it's landed. |
d3863ed to
956748c
Compare
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
| import torch | ||
|
|
||
|
|
||
| def pin_memory_fn(data, device=None): |
There was a problem hiding this comment.
Non-blocking: A short docstring and add this to torchdata.datapipes.utils.rst will be nice.
956748c to
1beb49f
Compare
|
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
NivekT
left a comment
There was a problem hiding this comment.
It will be nice to mention pin memory and what it does to users in tutorial at some point, but this is beyond the scope of this PR.
…1014) Summary: Fixes meta-pytorch#1013 ## Changes - Simplify the control flow of prefetcher - Delay Exception raised from thread worker to main thread in `__iter__` - Stop prefetching whenever Exception is received - As long as `stop_iteration` is not turned on or `buffer` is not empty, continue yielding data from `__iter__`. - Add serialization test - Add `PinMemory` DataPipe - `is_replciable() -> False` to keep it in the main process - Add unit tests - Update `test_proto_multi_rs.py` to `test_mprs.py` Pull Request resolved: meta-pytorch#1014 Reviewed By: NivekT Differential Revision: D43329696 Pulled By: ejguan fbshipit-source-id: da4326dbe2388f4e23b9a1a3a5c43da09d29185a
Summary: Fixes #1013 ## Changes - Simplify the control flow of prefetcher - Delay Exception raised from thread worker to main thread in `__iter__` - Stop prefetching whenever Exception is received - As long as `stop_iteration` is not turned on or `buffer` is not empty, continue yielding data from `__iter__`. - Add serialization test - Add `PinMemory` DataPipe - `is_replciable() -> False` to keep it in the main process - Add unit tests - Update `test_proto_multi_rs.py` to `test_mprs.py` Pull Request resolved: #1014 Reviewed By: NivekT Differential Revision: D43329696 Pulled By: ejguan fbshipit-source-id: da4326dbe2388f4e23b9a1a3a5c43da09d29185a
Fixes #1013
Changes
__iter__stop_iterationis not turned on orbufferis not empty, continue yielding data from__iter__.PinMemoryDataPipeis_replciable() -> Falseto keep it in the main processtest_proto_multi_rs.pytotest_mprs.py