116 lines
5.3 KiB
Python
116 lines
5.3 KiB
Python
# Converts SigLIP 2's vision encoder to a Core ML model for modules/Labels on macOS, then checks it
|
|
# against the ONNX model mcap runs elsewhere, on a folder of images. Run once, on a Mac; mcap only uses
|
|
# the .mlpackage it writes (into modules/Labels/models, where the build copies it from).
|
|
#
|
|
# python3 -m venv ~/mcap-coreml
|
|
# ~/mcap-coreml/bin/pip install torch==2.7.0 transformers coremltools==9.0 onnxruntime pillow sentencepiece protobuf
|
|
# ~/mcap-coreml/bin/python tools/export_siglip_coreml.py --check FOLDER_OF_IMAGES
|
|
#
|
|
# torch 2.7.0 is the newest coremltools 9.0 was tested with. The model is the vision tower of
|
|
# google/siglip2-base-patch16-384, from pixel_values [1, 3, 384, 384] (RGB resized to 384 x 384,
|
|
# scaled to -1..1) to pooler_output [1, 768], as the ONNX one; an ML Program in float16, for the GPU.
|
|
# Measured on an M5 Pro: 6.7 ms an image on the GPU, against 67 ms for onnxruntime on the CPU, with
|
|
# vectors within cosine 0.99994 of the CPU's. The Neural Engine gets float16 badly wrong for this
|
|
# model (cosine 0.72), so mcap runs it on the CPU and GPU only.
|
|
|
|
import argparse, glob, os, sys, time
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
import coremltools as ct
|
|
from transformers import AutoModel
|
|
from transformers.utils import logging
|
|
|
|
logging.set_verbosity_error()
|
|
|
|
CHECKPOINT = "google/siglip2-base-patch16-384"
|
|
SIZE = 384
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
MODELS = os.path.join(HERE, "..", "modules", "Labels", "models")
|
|
OUTPUT = os.path.join(MODELS, "siglip2-base-patch16-384-vision.mlpackage")
|
|
ONNX = os.path.join(MODELS, "siglip2-base-patch16-384-vision.onnx")
|
|
|
|
|
|
# coremltools 9.0 casts a folded one-element tensor with int(array), which NumPy 2 refuses ("only
|
|
# 0-dimensional arrays can be converted"); the same cast, with the element taken out first.
|
|
from coremltools.converters.mil.frontend.torch import ops as ct_ops
|
|
_original_cast = ct_ops._cast
|
|
def _cast_one_element(context, node, dtype, dtype_name):
|
|
x = ct_ops._get_inputs(context, node, expected=1)[0]
|
|
if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
|
|
context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
|
|
return
|
|
_original_cast(context, node, dtype, dtype_name)
|
|
ct_ops._cast = _cast_one_element
|
|
|
|
|
|
class Pooled(torch.nn.Module):
|
|
"""The vision tower, from pixel values to the pooled vector, as the ONNX graph has it."""
|
|
|
|
def __init__(self, model):
|
|
super().__init__()
|
|
self.model = model
|
|
|
|
def forward(self, pixel_values):
|
|
return self.model(pixel_values=pixel_values).pooler_output
|
|
|
|
|
|
def load_images(folder, limit):
|
|
paths = sorted(p for p in glob.glob(os.path.join(folder, "**", "*"), recursive=True)
|
|
if p.lower().endswith((".png", ".jpg", ".jpeg", ".webp")))[:limit]
|
|
images = []
|
|
for p in paths:
|
|
image = Image.open(p).convert("RGB").resize((SIZE, SIZE), Image.BICUBIC)
|
|
a = np.asarray(image, dtype=np.float32) / 255.0
|
|
images.append(((a - 0.5) / 0.5).transpose(2, 0, 1))
|
|
return np.stack(images).astype(np.float32)
|
|
|
|
|
|
def worst_cosine(a, b):
|
|
a = a / np.linalg.norm(a, axis=-1, keepdims=True)
|
|
b = b / np.linalg.norm(b, axis=-1, keepdims=True)
|
|
return float(np.min(np.sum(a * b, axis=-1)))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--check", metavar="FOLDER", help="images to check the Core ML model against the ONNX one on")
|
|
parser.add_argument("--limit", type=int, default=32)
|
|
args = parser.parse_args()
|
|
|
|
model = Pooled(AutoModel.from_pretrained(CHECKPOINT, torch_dtype=torch.float32).vision_model.eval()).eval()
|
|
example = torch.zeros(1, 3, SIZE, SIZE)
|
|
with torch.no_grad():
|
|
traced = torch.jit.trace(model, example)
|
|
start = time.perf_counter()
|
|
ml = ct.convert(traced, inputs=[ct.TensorType(name="pixel_values", shape=(1, 3, SIZE, SIZE))],
|
|
outputs=[ct.TensorType(name="pooler_output")],
|
|
convert_to="mlprogram", compute_precision=ct.precision.FLOAT16,
|
|
minimum_deployment_target=ct.target.macOS14)
|
|
ml.short_description = f"{CHECKPOINT} vision encoder (pooler_output), float16, for mcap"
|
|
ml.save(OUTPUT)
|
|
print(f"Wrote {OUTPUT} in {time.perf_counter() - start:.1f} s.")
|
|
|
|
if not args.check:
|
|
return
|
|
import onnxruntime as ort
|
|
images = load_images(args.check, args.limit)
|
|
if not len(images):
|
|
sys.exit(f"No images in {args.check}.")
|
|
session = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"])
|
|
name = session.get_inputs()[0].name
|
|
reference = np.concatenate([session.run(["pooler_output"], {name: images[i:i + 1]})[0] for i in range(len(images))])
|
|
coreml = ct.models.MLModel(OUTPUT, compute_units=ct.ComputeUnit.CPU_AND_GPU)
|
|
coreml.predict({"pixel_values": images[:1]})
|
|
start = time.perf_counter()
|
|
got = np.concatenate([coreml.predict({"pixel_values": images[i:i + 1]})["pooler_output"] for i in range(len(images))])
|
|
per = (time.perf_counter() - start) / len(images)
|
|
cosine = worst_cosine(got.astype(np.float32), reference)
|
|
print(f"{len(images)} images: {per * 1000:.1f} ms an image on the GPU, worst cosine to the ONNX model on the CPU {cosine:.5f}.")
|
|
if cosine < 0.999:
|
|
sys.exit("The Core ML model is too far from the ONNX one.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|