-
Notifications
You must be signed in to change notification settings - Fork 14
Implement tofile on tensors to reduce data write time by 40% #210
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
base: main
Are you sure you want to change the base?
Conversation
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #210 +/- ##
==========================================
- Coverage 76.83% 76.57% -0.27%
==========================================
Files 40 40
Lines 4922 4965 +43
Branches 980 989 +9
==========================================
+ Hits 3782 3802 +20
- Misses 856 873 +17
- Partials 284 290 +6 ☔ View full report in Codecov by Sentry. |
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
cc @iksnagreb |
|
# Implement tobytes to support native PyTorch types so we can use types like bloat16 | ||
# Reading from memory directly is also more efficient because | ||
# it avoids copying to a NumPy array | ||
_, address = self._get_data_chunk() |
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.
Rename variable to be more accurate. This should be a list of some py objects.
TODO: tests
This PR introduces the
tofile
method on tensors (similarly named as the one on numpy arrays), which allows for faster write and lower memory usage on external data by bypassing tobytes().Compatibility with existing
TensorProtocol
s is maintained in the external data module by usingtofile
only when it is available in the class. TheTorchTensor
class in PyTorch exporter should be updated accordingly to leverage the new logic when saving.Note that io time to disk is reduced by 40% below.
Reference: https://github.com/microsoft/onnxscript/pull/2241/files/b2381658492510a9bcc8c0a8574db7368e33bceb
Before:
After:
Fix #207