Implement tofile on tensors to reduce data write time by 40% #210
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.93% +0.09%
==========================================
Files 40 40
Lines 4922 4994 +72
Branches 980 998 +18
==========================================
+ Hits 3782 3842 +60
- Misses 856 864 +8
- Partials 284 288 +4 ☔ 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 |
|
justinchuby
changed the title
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]>
|
@titaiwangms @gramalingam this is ready for review, thanks. |
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
| file: A file-like object with a ``write`` method that accepts bytes, or has an ``fileno()`` method. | ||
| """ | ||
| if _supports_fileno(file) and isinstance(self._raw, np.ndarray): | ||
| # This is a duplication of tobytes() for handling special cases |
There was a problem hiding this comment.
Should we pack this to a private function, and document the reason we need it (I would say this is very technical knowledge)?
| the tensor is recommended if IO overhead and memory usage is a concern. | ||
| To obtain an array, call :meth:`numpy`. To obtain the bytes, | ||
| call :meth:`tobytes`. |
There was a problem hiding this comment.
Probably want to add tofiile()?
| if self._offset is not None: | ||
| src.seek(self._offset) | ||
| bytes_to_copy = self._length or self.nbytes | ||
| chunk_size = 1024 * 1024 # 1MB |
There was a problem hiding this comment.
I am wondering why do we know this is the most efficient chunk size? Do we randomly select it?
| """Return the bytes of the tensor.""" | ||
| return self._evaluate().tobytes() | ||
|
|
||
| def tofile(self, file) -> None: |
There was a problem hiding this comment.
I am wondering whether tofile() makes sense to LazyTensor. hmm
| data_file.write(b"\0" * (current_offset - file_size)) | ||
| data_file.write(raw_data) | ||
|
|
||
| if hasattr(tensor, "tofile"): |
There was a problem hiding this comment.
What tensors do not have tofile()? torch? Better document this
This PR introduces the
tofilemethod 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
TensorProtocols is maintained in the external data module by usingtofileonly when it is available in the class. TheTorchTensorclass 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.
Note
TensorProtocol is not updated because we do isinstance() checks on external implementations (PyTorch). Adding the method in the protocol will cause isinstance check to fail on those implementations that have not added the tofile method.
Reference: https://github.com/microsoft/onnxscript/pull/2241/files/b2381658492510a9bcc8c0a8574db7368e33bceb
Before:
After:
Fix #207