Skip to content

Add a verbose mode to torch api for external data save #2643

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.

Sign up for GitHub

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

Merged
merged 5 commits into from
Oct 27, 2025
+29 −2
Merged

Add a verbose mode to torch api for external data save #2643

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
62855af
Create progress bar
justinchuby Oct 18, 2025
0a75c9f
Update onnxscript/_framework_apis/torch_2_5.py
justinchuby Oct 18, 2025
743069e
format
justinchuby Oct 18, 2025
57d5648
shortname
justinchuby Oct 18, 2025
3fd2986
pylint
justinchuby Oct 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 29 additions & 2 deletions onnxscript/_framework_apis/torch_2_5.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
]

import dataclasses
import importlib.util
import os
import pathlib
from typing import Callable
Expand Down Expand Up @@ -63,7 +64,9 @@ def check_model(model: ir.Model) -> None:
del model # Unused yet


def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike) -> None:
def save_model_with_external_data(
model: ir.Model, model_path: str | os.PathLike, verbose: bool = False
) -> None:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we duplicate the api and add verbose in a later version of framework_apis (eg 2_10) if we only want this after 2.10?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is easier and is done in a bc compatible way. But if you think duplicating makes it more obvious I can do that too

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will merge first. LMK if you would like to see a follow up, thanks!

"""Save the model with external data. The model is unchanged after saving."""

# TODO(#1835): Decide if we want to externalize large attributes as well
Expand All @@ -78,7 +81,31 @@ def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike
destination_path = pathlib.Path(model_path)
data_path = f"{destination_path.name}.data"

ir.save(model, model_path, external_data=data_path)
# Show a progress bar if verbose is True and tqdm is installed
use_tqdm = verbose and importlib.util.find_spec("tqdm") is not None

if use_tqdm:
import tqdm # pylint: disable=import-outside-toplevel

with tqdm.tqdm() as pbar:
total_set = False

def callback(
tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo
) -> None:
nonlocal total_set
if not total_set:
pbar.total = metadata.total
total_set = True

pbar.update()
pbar.set_description(
f"Saving {tensor.name} ({tensor.dtype.short_name()}, {tensor.shape}) at offset {metadata.offset}"
)

ir.save(model, model_path, external_data=data_path, callback=callback)
else:
ir.save(model, model_path, external_data=data_path)


def get_torchlib_ops() -> list[_OnnxFunctionMeta]:
Expand Down
Loading