Support >2G model export - alternative implementation | torchlib(feat) #1004
Conversation
|
@BowenBao Speed seems to be reasonably good. There's errors: |
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1004 +/- ##
==========================================
+ Coverage 78.64% 78.65% +0.01%
==========================================
Files 118 118
Lines 15441 15435 -6
Branches 2424 2422 -2
==========================================
- Hits 12144 12141 -3
+ Misses 2899 2896 -3
Partials 398 398 ☔ View full report in Codecov by Sentry. |
|
Adding back infer_shapes and check_model made tests pass locally. |
|
Let me test again |
Nice, thanks! I updated this PR for review |
justinchuby
added
module: torchlib
|
Needs validation with real models |
|
@BowenBao should we move this to under a flag? |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| ) | ||
| tensor_proto = onnx.helper.make_tensor( | ||
| name=name, | ||
| data_type=_type_utils.JitScalarType.from_dtype(tensor.dtype).onnx_type(), |
Check failure
Code scanning / lintrunner
PYLINT/E0602
| ) | ||
| tensor_proto = onnx.helper.make_tensor( | ||
| name=name, | ||
| data_type=_type_utils.JitScalarType.from_dtype(tensor.dtype).onnx_type(), |
Check failure
Code scanning / lintrunner
MYPY/name-defined
| ) | ||
| tensor_proto = onnx.helper.make_tensor( | ||
| name=name, | ||
| data_type=_type_utils.JitScalarType.from_dtype(tensor.dtype).onnx_type(), |
Check failure
Code scanning / lintrunner
RUFF/F821
| @@ -378,12 +374,38 @@ def eval_function( # type: ignore[override] | |||
| return self._graph.add_function_call(function, inputs, attributes) | |||
|
|
|||
|
|
|||
| def _add_initializers( | |||
There was a problem hiding this comment.
fyi @kunal-vaishnavi creating onnx tensorproto from torch tensor dataptr
Support >2G model export - alternative implementation where we build initializers ourselves. This is a follow up of #1003. We now build the TensorProto ourselves directly from PyTorch tensors. This circumvents torchscript _export_onnx's limitation of 2G protobuf serialization and additional serialization, because we are now keeping everything in memory.
2G model is now no longer a special case because we add initializers in a separate step.
TODO