to_pytorch: enable prefetching
#664
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 653 did not really enable prefetching. Prefetch was only implemented for map(), so the example gave me a false impression that the prefetching was working, but it was not. Now, to_pytorch uses AsyncMapper to prefetch the data. The number of workers is set to 2 by default, but it can be changed by setting the `prefetch` in the settings. For me, this dropped the time to load the data by 90%, from ~300s to now ~35s.
Deploying datachain-documentation with
|
| Latest commit: |
9bbecac
|
| Status: | ✅ Deploy successful! |
| Preview URL: | https://3ad70991.datachain-documentation.pages.dev |
| Branch Preview URL: | https://prefetch-pytorch.datachain-documentation.pages.dev |
skshetry
commented
Dec 5, 2024
| @@ -31,6 +32,8 @@ def label_to_int(value: str, classes: list) -> int: | |||
|
|
|||
|
|
|||
| class PytorchDataset(IterableDataset): | |||
| prefetch: int = 2 | |||
There was a problem hiding this comment.
Keeping it same as the default.
datachain/src/datachain/lib/udf.py
Line 293 in 911c22f
skshetry
commented
Dec 5, 2024
| total_rank, total_workers = self.get_rank_and_workers() | ||
| def _rows_iter(self, total_rank: int, total_workers: int): | ||
| catalog = self._get_catalog() | ||
| session = Session("PyTorch", catalog=catalog) |
There was a problem hiding this comment.
We need to create a new session and a catalog here, since the AsyncMapper runs this on a separate thread.
skshetry
changed the title
to_pytorch: enable prefetching
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #664 +/- ##
==========================================
- Coverage 87.32% 87.32% -0.01%
==========================================
Files 113 113
Lines 10717 10727 +10
Branches 1469 1469
==========================================
+ Hits 9359 9367 +8
- Misses 985 986 +1
- Partials 373 374 +1
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
#653 did not really enable prefetching. Prefetch was only implemented for
map(), so the example gave me a false impression that the prefetching was working, but it was not.Now,
to_pytorchusesAsyncMapperto prefetch the data. The number of workers is set to 2 by default, but it can be changed by setting theprefetchin the settings.For me, this dropped the time to download the data by 90%, from ~300s to now ~35s.
Closes #631.