-
Notifications
You must be signed in to change notification settings - Fork 159
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Update Prefetcher and Implement PinMemory IterDataPipe #1014
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. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Overall, LGTM with a few notes/questions. The simplification of Prefetcher
is great!
Also, I don't think our CI runs with CUDA, can you confirm?
d3863ed
to
956748c
Compare
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM! Thanks!
torchdata/utils/pin_memory.py
Outdated
import torch | ||
|
||
|
||
def pin_memory_fn(data, device=None): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Non-blocking: A short docstring and add this to torchdata.datapipes.utils.rst
will be nice.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good point!
956748c
to
1beb49f
Compare
@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Summary: Fixes 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: 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_iteration
is not turned on orbuffer
is not empty, continue yielding data from__iter__
.PinMemory
DataPipeis_replciable() -> False
to keep it in the main processtest_proto_multi_rs.py
totest_mprs.py