Machine Learning Model Extension Specification

• Title: Machine Learning Model Extension
• Identifier: https://schemas.stacspec.org/2.0.0.alpha.0/extensions/ml-model/json-schema/schema.json
• Field Name Prefix: mlm
• Scope: Item, Collection
• Extension Maturity Classification: Proposal
• Owner:
  • @fmigneault
  • @rbavery
  • @ymoisan
  • @sfoucher

The STAC Machine Learning Model (MLM) Extension provides a standard set of fields to describe machine learning models trained on overhead imagery and enable running model inference.

The main objectives of the extension are:

1. to enable building model collections that can be searched alongside associated STAC datasets
2. record all necessary bands, parameters, modeling artifact locations, and high-level processing steps to deploy an inference service.

Specifically, this extension records the following information to make ML models searchable and reusable:

1. Sensor band specifications
2. Model input transforms including resize and normalization
3. Model output shape, data type, and its semantic interpretation
4. An optional, flexible description of the runtime environment to be able to run the model
5. Scientific references

The MLM specification is biased towards providing metadata fields for supervised machine learning models. However, fields that relate to supervised ML are optional and users can use the fields they need for different tasks.

See Best Practices for guidance on what other STAC extensions you should use in conjunction with this extension. The Machine Learning Model Extension purposely omits and delegates some definitions to other STAC extensions to favor reusability and avoid metadata duplication whenever possible. A properly defined MLM STAC Item/Collection should almost never have the Machine Learning Model Extension exclusively in stac_extensions.

Check the original technical report for an earlier version of the Model Extension, formerly known as the Deep Learning Model Extension (DLM), here for more details. The DLM was renamed to the current MLM Extension and refactored to form a cohesive definition across all machine learning approaches, regardless of whether the approach constitutes a deep neural network or other statistical approach.

• Examples:
  • Example with a ??? trained with torchgeo TODO update example
  • Collection example: Shows the basic usage of the extension in a STAC Collection
• JSON Schema TODO update
• Changelog

Item Properties and Collection Fields

In addition, fields from the following extensions must be imported in the item:

• Scientific Extension Specification to describe relevant publications.
• Version Extension Specification to define version tags.

Model Input Object

Parameters Object

The parameters field can either be specified in the Model Input Object if they are associated with a specific input or as an Item or Collection field if the parameters are supplied without relation to a specific model input.

Bands and Statistics

We use the STAC 1.1 Bands Object for representing bands information, including the nodata value, data type, and common band names. Only bands used to train or fine tune the model should be included in this bands field.

A deviation from the STAC 1.1 Bands Object is that we do not include the Statistics object at the band object level, but at the Model Input level. This is because in machine learning, it is common to only need overall statistics for the dataset used to train the model to normalize all bands.

Array Object

Note: It is common in the machine learning, computer vision, and remote sensing communities to refer to rasters that are inputs to a model as arrays or tensors. Array Objects are distinct from the JSON array type used to represent lists of values.

Runtime Object

Accelerator Enum

It is recommended to define accelerator with one of the following values:

• amd64 models compatible with AMD or Intel CPUs (no hardware specific optimizations)
• cuda models compatible with NVIDIA GPUs
• xla models compiled with XLA. models trained on TPUs are typically compiled with XLA.
• amd-rocm models trained on AMD GPUs
• intel-ipex-cpu for models optimized with IPEX for Intel CPUs
• intel-ipex-gpu for models optimized with IPEX for Intel GPUs
• macos-arm for models trained on Apple Silicon

Container Object

If you're unsure how to containerize your model, we suggest starting from the latest official container image for your framework that works with your model and pinning the container version.

Examples: Pytorch Dockerhub Pytorch Docker Run Example

Tensorflow Dockerhub Tensorflow Docker Run Example

Using a base image for a framework looks like

# In your Dockerfile, pull the latest base image with all framework dependencies including accelerator drivers
FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime

### Your specific environment setup to run your model
RUN pip install my_package

You can also use other base images. Pytorch and Tensorflow offer docker images for serving models for inference.

• Torchserve
• TFServing

Model Output Object

While only task is a required field, all fields are recommended for supervised tasks that produce a fixed shape tensor and have output classes. image-captioning, multi-modal, and generative tasks may not return fixed shape tensors or classes.

Task Enum

It is recommended to define task with one of the following values for each Model Output Object:

• regression
• classification
• object detection
• semantic segmentation
• instance segmentation
• panoptic segmentation
• multi-modal
• similarity search
• image captioning
• generative
• super resolution

If the task falls within the category of supervised machine learning and uses labels during training, this should align with the label:tasks values defined in STAC Label Extension for relevant STAC Collections and Items published with the model described by this extension.

Result Array Object

Class Object

See the documentation for the Class Object. We don't use the Bit Field Object since inputs and outputs to machine learning models don't typically use bit fields.

Relation types

The following types should be used as applicable rel types in the Link Object of STAC Items describing Band Assets used with a model.

Contributing

All contributions are subject to the STAC Specification Code of Conduct. For contributions, please follow the STAC specification contributing guide Instructions for running tests are copied here for convenience.

Running tests

The same checks that run as checks on PRs are part of the repository and can be run locally to verify that changes are valid. To run tests locally, you'll need npm, which is a standard part of any node.js installation.

First, install everything with npm once. Navigate to the root of this repository and on your command line run:

npm install

Then to check Markdown formatting and test the examples against the JSON schema, you can run:

npm test

This will spit out the same texts that you see online, and you can then go and fix your markdown or examples.

If the tests reveal formatting problems with the examples, you can fix them with:

npm run format-examples