Torch/TorchGeo Export & Packaging

.github/workflows/stac-model.yml

@@ -8,6 +8,8 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
+    env:
+      COVERAGE_FILE: .coverage.${{ matrix.python-version }}
 
     steps:
     - uses: actions/checkout@v4
@@ -24,12 +26,18 @@ jobs:
       with:
         path: .venv
         key: venv-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('uv.lock') }}
+
     - name: Install dependencies
       run: make install-dev
+      if: ${{ matrix.python-version == '3.10' }}
+
+    - name: Install dependencies with extras
+      run: make install-dev-extras
+      if: ${{ matrix.python-version != '3.10' }}
 
     - name: Run checks
-      run: |
-        make lint-all
+      if: ${{ matrix.python-version != '3.10' }}
+      run: make lint-all
 
     - name: Run tests
       run: |

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
       - id: check-yaml
       - id: end-of-file-fixer
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.1.12'  # Use the latest version of ruff-pre-commit
+    rev: 'v0.12.4'
     hooks:
       - id: ruff
         pass_filenames: false

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add torch export and packaging utilities for combining a model, transforms, and MLM schema
+  compliant metadata into a single `.pt2` archive.
 - Add `cpu` device type to `mlm:accelerator` and corresponding `stac_model.runtime.AcceleratorEnum`.
 - Add `pre_processing_function` and `post_processing_function` support as JSON array 
   of [Processing Expression](README.md#processing-expression) definitions respectively

Makefile

@@ -47,6 +47,11 @@ install-dev: setup
 	$(UV_COMMAND) export --format requirements-txt -o requirements-dev.txt
 	$(UV_COMMAND) pip install --python "$(UV_PYTHON_ROOT)" -r requirements-dev.txt
 
+.PHONY: install-dev-extras
+install-dev-extras: setup
+	$(UV_COMMAND) export --format requirements-txt -o requirements-dev.txt
+	$(UV_COMMAND) pip install --python "$(UV_PYTHON_ROOT)" -e .[torch] -r requirements-dev.txt
+
 .PHONY: pre-commit-install
 pre-commit-install: setup
 	$(UV_COMMAND) run --no-sync --python "$(UV_PYTHON_ROOT)" pre-commit install
@@ -59,11 +64,12 @@ codestyle: setup
 .PHONY: format
 format: codestyle
 
-#* Linting
+#* Testing
 .PHONY: test
 test: setup
-	$(UV_COMMAND) run --no-sync --python "$(UV_PYTHON_ROOT)" pytest -c pyproject.toml -v --cov-report=html --cov=stac_model tests/
+	$(UV_COMMAND) run --no-sync --python "$(UV_PYTHON_ROOT)" pytest -c pyproject.toml -v --cov-report=html --cov=stac_model --cov-config pyproject.toml tests/
 
+#* Linting
 .PHONY: check
 check: check-examples check-markdown check-lint check-mypy check-safety check-citation
 

README_STAC_MODEL.md

@@ -65,6 +65,62 @@ with open("examples/mlm-metadata.yaml", "r", encoding="utf-8") as f:
 MLModelProperties.model_validate(metadata["properties"])  
 ```
 
+## Exporting and Packaging PyTorch Models, Transforms, and Model Metadata
+
+As of PyTorch 2.8, and stac_model 1.5.0, you can now export and package PyTorch models, transforms,
+and model metadata using functions in `stac_model.torch.export`. Below is an example of exporting a
+U-Net model pretrained on the [Fields of The World (FTW) dataset](https://fieldsofthe.world/) for
+field boundary segmentation in Sentinel-2 satellite imagery using the [TorchGeo](https://github.com/microsoft/torchgeo) library.
+
+> 📝 **Note:** To customize the metadata for your model you can use this [example](./tests/torch/metadata.yaml) as a template.
+
+```python
+import torch
+import torchvision.transforms.v2 as T
+from torchgeo.models import Unet_Weights, unet
+from stac_model.torch.export import save
+
+weights = Unet_Weights.SENTINEL2_3CLASS_FTW
+transforms = torch.nn.Sequential(
+  T.Resize((256, 256)),
+  T.Normalize(mean=[0.0], std=[3000.0])
+)
+model = unet(weights=weights)
+
+save(
+    output_file="ftw.pt2",
+    model=model,  # Must be an nn.Module
+    transforms=transforms,  # Must be an nn.Module
+    metadata_path="metadata.yaml",  # Can be a metadata yaml or MLModelProperties object
+    input_shape=[-1, 8, -1, -1],  # -1 indicates a dynamic shaped dimension
+    device="cpu",
+    dtype=torch.float32,
+    aoti_compile_and_package=False,  # True for AOTInductor compile otherwise use torch.export
+)
+```
+
+The model, transforms, and metadata can then be loaded into an environment with only torch and stac_model as required dependencies like below:
+
+```python
+import yaml
+from torch.export.pt2_archive._package import load_pt2
+
+pt2 = load_pt2(archive_path)
+metadata = yaml.safe_load(pt2.extra_files["mlm-metadata"])
+
+# If exported with aoti_compile_and_package=True
+model = pt2.aoti_runners["model"]
+transforms = pt2.aoti_runners["transforms"]
+
+# If exported with aoti_compile_and_package=False
+model = pt2.exported_programs["model"].module()
+transforms = pt2.exported_programs["transforms"].module()
+
+# Inference
+batch = ...  # An input batch tensor
+outputs = model(transforms(batch))
+```
+
 ## 📈 Releases
 
 You can see the list of available releases on the [GitHub Releases][github-releases] page.