Model Extras

In addition to interface and onnx-related objects, there are a few objects that may be of use when interactions with models.

---

HuggingFaceORTModel

The HuggingFaceORTModel is an enum that allows you to specify an ORT type for a HuggingFaceModel. Refer to the source code below for the available options.

from opsml import HuggingFaceORTModel, HuggingFaceOnnxArgs

HuggingFaceOnnxArgs(
    ort_type=HuggingFaceORTModel.ORT_SEQUENCE_CLASSIFICATION.value,
    provider="CPUExecutionProvider",
    quantize=False,
    )

opsml.HuggingFaceORTModel

Bases: str, Enum

Source code in opsml/types/huggingface.py

class HuggingFaceORTModel(str, Enum):
    ORT_AUDIO_CLASSIFICATION = "ORTModelForAudioClassification"
    ORT_AUDIO_FRAME_CLASSIFICATION = "ORTModelForAudioFrameClassification"
    ORT_AUDIO_XVECTOR = "ORTModelForAudioXVector"
    ORT_CUSTOM_TASKS = "ORTModelForCustomTasks"
    ORT_CTC = "ORTModelForCTC"
    ORT_FEATURE_EXTRACTION = "ORTModelForFeatureExtraction"
    ORT_IMAGE_CLASSIFICATION = "ORTModelForImageClassification"
    ORT_MASKED_LM = "ORTModelForMaskedLM"
    ORT_MULTIPLE_CHOICE = "ORTModelForMultipleChoice"
    ORT_QUESTION_ANSWERING = "ORTModelForQuestionAnswering"
    ORT_SEMANTIC_SEGMENTATION = "ORTModelForSemanticSegmentation"
    ORT_SEQUENCE_CLASSIFICATION = "ORTModelForSequenceClassification"
    ORT_TOKEN_CLASSIFICATION = "ORTModelForTokenClassification"
    ORT_SEQ2SEQ_LM = "ORTModelForSeq2SeqLM"
    ORT_SPEECH_SEQ2SEQ = "ORTModelForSpeechSeq2Seq"
    ORT_VISION2SEQ = "ORTModelForVision2Seq"
    ORT_PIX2STRUCT = "ORTModelForPix2Struct"
    ORT_CAUSAL_LM = "ORTModelForCausalLM"
    ORT_OPTIMIZER = "ORTOptimizer"
    ORT_QUANTIZER = "ORTQuantizer"
    ORT_TRAINER = "ORTTrainer"
    ORT_SEQ2SEQ_TRAINER = "ORTSeq2SeqTrainer"
    ORT_TRAINING_ARGUMENTS = "ORTTrainingArguments"
    ORT_SEQ2SEQ_TRAINING_ARGUMENTS = "ORTSeq2SeqTrainingArguments"
    ORT_STABLE_DIFFUSION_PIPELINE = "ORTStableDiffusionPipeline"
    ORT_STABLE_DIFFUSION_IMG2IMG_PIPELINE = "ORTStableDiffusionImg2ImgPipeline"
    ORT_STABLE_DIFFUSION_INPAINT_PIPELINE = "ORTStableDiffusionInpaintPipeline"
    ORT_STABLE_DIFFUSION_XL_PIPELINE = "ORTStableDiffusionXLPipeline"
    ORT_STABLE_DIFFUSION_XL_IMG2IMG_PIPELINE = "ORTStableDiffusionXLImg2ImgPipeline"

---

HuggingFaceTask

HuggingFaceTask is an enum that allows you to specify a task type for a HuggingFaceModel. Refer to the source code below for the available options.

from opsml import HuggingFaceTask, HuggingFaceModel

HuggingFaceModel(
    model=model,
    task_type=HuggingFaceTask.SEQUENCE_CLASSIFICATION.value,
    )

opsml.HuggingFaceTask

Bases: str, Enum

Source code in opsml/types/huggingface.py

class HuggingFaceTask(str, Enum):
    AUDIO_CLASSIFICATION = "audio-classification"
    AUTOMATIC_SPEECH_RECOGNITION = "automatic-speech-recognition"
    CONVERSATIONAL = "conversational"
    DEPTH_ESTIMATION = "depth-estimation"
    DOCUMENT_QUESTION_ANSWERING = "document-question-answering"
    FEATURE_EXTRACTION = "feature-extraction"
    FILL_MASK = "fill-mask"
    IMAGE_CLASSIFICATION = "image-classification"
    IMAGE_SEGMENTATION = "image-segmentation"
    IMAGE_TO_IMAGE = "image-to-image"
    IMAGE_TO_TEXT = "image-to-text"
    MASK_GENERATION = "mask-generation"
    OBJECT_DETECTION = "object-detection"
    QUESTION_ANSWERING = "question-answering"
    SUMMARIZATION = "summarization"
    TABLE_QUESTION_ANSWERING = "table-question-answering"
    TEXT2TEXT_GENERATION = "text2text-generation"
    TEXT_CLASSIFICATION = "text-classification"
    TEXT_GENERATION = "text-generation"
    TEXT_TO_AUDIO = "text-to-audio"
    TOKEN_CLASSIFICATION = "token-classification"
    TRANSLATION = "translation"
    TRANSLATION_XX_TO_YY = "translation_xx_to_yy"
    VIDEO_CLASSIFICATION = "video-classification"
    VISUAL_QUESTION_ANSWERING = "visual-question-answering"
    ZERO_SHOT_CLASSIFICATION = "zero-shot-classification"
    ZERO_SHOT_IMAGE_CLASSIFICATION = "zero-shot-image-classification"
    ZERO_SHOT_AUDIO_CLASSIFICATION = "zero-shot-audio-classification"
    ZERO_SHOT_OBJECT_DETECTION = "zero-shot-object-detection"

---

TorchSaveArgs

Optional TorchModel arguments for saving a TorchModel object. Only as_state_dict is currently supported. If True, the TorchModel model object's state dict will be

opsml.types.TorchSaveArgs

Bases: BaseModel

Torch save arguments.

Parameters:

Name	Type	Description	Default
as_state_dict		Indicates to save the torch model in state_dict format. If True, the model architecture will need to be provided at load time.	required

Source code in opsml/types/model.py

class TorchSaveArgs(BaseModel):
    """Torch save arguments.

    Args:
        as_state_dict:
            Indicates to save the torch model in state_dict format. If True, the model
            architecture will need to be provided at load time.
    """

    as_state_dict: bool = False

---

OnnxModel

OnnxModel is a pydantic class that is used to store converted onnx models. In the case of a BYO onnx model, you will need to supply an OnnxModel object to the ModelInterface class.

from opsml import OnnxModel, TorchModel
import onnx
import onnxruntime as ort

# Super Resolution model definition in PyTorch
import torch.nn as nn
import torch.nn.init as init
import torch.onnx
import torch.utils.model_zoo as model_zoo
from torch import nn

class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor**2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

        self._initialize_weights()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x

    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv2.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv3.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv4.weight)

# Create the super-resolution model by using the above model definition.
torch_model = SuperResolutionNet(upscale_factor=3)

# Load pretrained model weights
model_url = "https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth"
batch_size = 1  # just a random number

# Initialize model with the pretrained weights
def map_location(storage, loc):
    return storage

if torch.cuda.is_available():
    map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))

# set the model to inference mode
torch_model.eval()

# Input to the model
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
torch_model(x)

with tempfile.TemporaryDirectory() as tmpdir:
    onnx_path = f"{tmpdir}/super_resolution.onnx"
    # Export the model
    torch.onnx.export(
        torch_model,  # model being run
        x,  # model input (or a tuple for multiple inputs)
        onnx_path,  # where to save the model (can be a file or file-like object)
        export_params=True,  # store the trained parameter weights inside the model file
        opset_version=10,  # the ONNX version to export the model to
        do_constant_folding=True,  # whether to execute constant folding for optimization
        input_names=["input"],  # the model's input names
        output_names=["output"],  # the model's output names
        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},  # variable length axes
    )

    onnx_model = onnx.load(onnx_path)

    ort_sess = ort.InferenceSession(onnx_model.SerializeToString())

onnx_model = OnnxModel(onnx_version="1.14.0", sess=ort_sess)

interface = TorchModel(
    model=torch_model,
    sample_data=x,
    onnx_model=onnx_model,
    save_args={"as_state_dict": True},
)

opsml.types.OnnxModel

Bases: BaseModel

Source code in opsml/types/model.py

class OnnxModel(BaseModel):
    onnx_version: str = Field(..., description="Version of onnx model used to create proto")
    sess: Union[OnnxInferenceSession, ORTModel, Pipeline] = Field(default=None, description="Onnx model session")  # type: ignore

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def sess_to_path(self, path: Path) -> None:
        """Helper method for taking existing onnx model session and saving to path

        Args:
            path:
                Path to save onnx model
        """

        if not isinstance(self.sess, rt.InferenceSession):
            return

        logger.info("Saving existing onnx model")

        if self.sess._model_bytes:  # pylint: disable=protected-access
            path.write_bytes(self.sess._model_bytes)  # pylint: disable=protected-access

        else:
            # copy model path to new path
            from fsspec.implementations.local import LocalFileSystem

            file_sys = LocalFileSystem()
            lpath = self.sess._model_path  # pylint: disable=protected-access
            file_sys.copy(lpath, str(path))

        return

sess_to_path(path)

Helper method for taking existing onnx model session and saving to path

Parameters:

Name	Type	Description	Default
path	Path	Path to save onnx model	required

Source code in opsml/types/model.py

def sess_to_path(self, path: Path) -> None:
    """Helper method for taking existing onnx model session and saving to path

    Args:
        path:
            Path to save onnx model
    """

    if not isinstance(self.sess, rt.InferenceSession):
        return

    logger.info("Saving existing onnx model")

    if self.sess._model_bytes:  # pylint: disable=protected-access
        path.write_bytes(self.sess._model_bytes)  # pylint: disable=protected-access

    else:
        # copy model path to new path
        from fsspec.implementations.local import LocalFileSystem

        file_sys = LocalFileSystem()
        lpath = self.sess._model_path  # pylint: disable=protected-access
        file_sys.copy(lpath, str(path))

    return