10000 initial export and load logic for PT2 archives by rbavery · Pull Request #87 · stac-extensions/mlm · GitHub
[go: up one dir, main page]
Skip to content

initial export and load logic for PT2 archives #87

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

Jump to bottom
Closed
wants to merge 1 commit into from
Closed

initial export and load logic for PT2 archives #87

wants to merge 1 commit into from

Conversation

rbavery
Copy link
Collaborator
@rbavery rbavery commented May 6, 2025 •
edited
Loading

Description

This is a WIP PR to implement export and loading of a PT2 model archive containing MLM metadata and optionally containing transforms that can be exported with Pytorch and loaded back to run on CPU or GPU prior to model inference.

It will use a new pt2 export API that will be available in a future release of Pytorch for exporting one or more models with (optional) compiled AOTI model artifacts and model metadata files (MLM).

package_pt2(
    f: FileLike,
    *,
    exported_programs: Optional[Union[ExportedProgram, dict[str, ExportedProgram]]] = None,
    aoti_files: Optional[Union[list[str], dict[str, list[str]]]] = None,
    extra_files: Optional[dict[str, Any]] = None,
) -> FileLike

@dataclass
class PT2ArchiveContents:
    exported_programs: dict[str, ExportedProgram]
    aoti_runners: dict[str, AOTICompiledModel]
    extra_files: dict[str, Any]

This will help standardize how to store multiple models and model metadatas in a PT2 archive, which we can then generalize to other framework archive types.

Initially I'm thinking that our opinionated STAC MLM + PT2 archive spec would require/suggest

  • one primary model file per pt2 archive, referred to by the model_name in the MLM metadata
  • optionally one transform model file per pt2 archive, referred to by a processing expression with a new format: pt2-transform

So in our instance of a PT2 Archive for a eurosat_model, the PT2ArchiveContents would be

exported_programs: {'eurosat_model': ExportedProgram,
                                  'transforms': ExportedProgram}
aoti_runners: {'eurosat_model': AOTICompiledModel,
                       'transforms': AOTICompiledMode
D42B
l}
extra_files: {'mlm': mlm.json} # mlm.json refers to eurosat_model and has a ProcessingExpression referring to `transforms`

I'm still working out an example where we use TorchGeo to collate and export the mlm metadata tot his archive.

Related Issue

Type of Change

  • 📚 Examples, docs, tutorials or dependencies update;
  • 🔧 Bug fix (non-breaking change which fixes an issue);
  • 🥂 Improvement (non-breaking change which improves an existing feature);
  • 🚀 New feature (non-breaking change which adds functionality);
  • 💥 Breaking change (fix or feature that would cause existing functionality to change);
  • 🔐 Security fix.

Checklist

  • I've read the CONTRIBUTING.md guide;
  • I've updated the code style using make check;
  • I've written tests for all new methods and classes that I created;
  • I've written the docstring in Google format for all the methods and classes that I used.

@rbavery rbavery changed the title initial export and load logic initial export and load logic for PT2 archives May 6, 2025
@fmigneault fmigneault self-requested a review May 28, 2025 19:26
@rbavery
Copy link
Collaborator Author
rbavery commented Jul 3, 2025

closing in favor of a follow up PR once torch 2.8 is released and the export spec for pt2 is solidified.

@rbavery rbavery closed this Jul 3, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@fmigneault fmigneault Awaiting requested review from fmigneault

Assignees

@rbavery rbavery

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@rbavery
0