ONNX and ONNX Runtime

import torch
import torchvision.models as models

# Load pretrained model
model = models.mobilenet_v2(pretrained=True)
model.eval()

# Create dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# Export to ONNX
torch.onnx.export(
    model,
    dummy_input,
    "mobilenet_v2.onnx",
    opset_version=13,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)
print("Model exported to mobilenet_v2.onnx")

import onnxruntime as ort
import numpy as np

# Create inference session
session = ort.InferenceSession("mobilenet_v2.onnx")

# Prepare input
input_name = session.get_inputs()[0].name
image = np.random.randn(1, 3, 224, 224).astype(np.float32)

# Run inference
result = session.run(None, {input_name: image})
prediction = np.argmax(result[0])
print(f"Predicted class: {prediction}")

# ONNX Runtime is typically 2-3x faster than native PyTorch inference

import coremltools as ct

# Convert from PyTorch (via traced model)
model = models.mobilenet_v2(pretrained=True).eval()
traced = torch.jit.trace(model, torch.randn(1, 3, 224, 224))

mlmodel = ct.convert(
    traced,
    inputs=[ct.ImageType(shape=(1, 3, 224, 224))],
    classifier_config=ct.ClassifierConfig("imagenet_classes.txt")
)

mlmodel.save("MobileNetV2.mlpackage")