diff --git a/README.md b/README.md index b590d3a..b42ccfa 100644 --- a/README.md +++ b/README.md @@ -28,11 +28,12 @@ see [Screen capture](#screen-capture). place (`C:\Program Files\NVIDIA\CUDNN`) is found without changing PATH. The OCR window says what, if anything, is missing. - **On macOS**, labelling and OCR run on the GPU through Core ML, from Core ML copies of the models - made once with `tools/export_siglip_coreml.py` (SigLIP 2's image model, about ten times as fast as - the CPU) and `tools/export_ppocr_coreml.py` (PP-OCRv5's detector and recognizer, about twice as - fast; run it with `--model mobile` and `--model server`; their first lines say how to set up the - Python they need). Without them, each runs on the + **On macOS**, labelling, OCR and segmentation run on the GPU through Core ML, from Core ML copies + of the models made once with `tools/export_siglip_coreml.py` (SigLIP 2's image model, about ten + times as fast as the CPU), `tools/export_ppocr_coreml.py` (PP-OCRv5's detector and recognizer, + about twice as fast; run it with `--model mobile` and `--model server`) and + `tools/export_mask2former_coreml.py` (Mask2Former, about four times as fast; after + `tools/export_mask2former.py`). Their first lines say how to set up the Python they need. Without them, each runs on the CPU. The first scan after making them is slower: Core ML prepares each image size the first time it meets it, and macOS keeps that for later runs. diff --git a/check_targets.jai b/check_targets.jai index 8746d58..42a0990 100644 --- a/check_targets.jai +++ b/check_targets.jai @@ -49,7 +49,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; GIT_BRANCH :: "check"; GIT_HASH :: "-"; -VERSION :: "0.1.31"; +VERSION :: "0.1.32"; BUILD_DATE :: "-"; DEBUG :: true; COMPILE_TESTS :: true; diff --git a/first.jai b/first.jai index 67d8011..4f728b0 100644 --- a/first.jai +++ b/first.jai @@ -8,7 +8,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; APP_BUNDLE_IDENTIFIER :: "com.example.mcap"; -VERSION :: "0.1.31"; +VERSION :: "0.1.32"; // Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated // builds never write over the ones you run by hand. build_directory := "bin"; @@ -407,6 +407,7 @@ copy_assets :: () { copy_if_changed("modules/Segment/models.config", tprint("%/segment-models/models.config", build_directory)); if file_exists("modules/Segment/models") { for file_list("modules/Segment/models") copy_if_changed(it, tprint("%/segment-models/%", build_directory, path_filename(it))); + #if OS == .MACOS clone_packages("modules/Segment/models", tprint("%/segment-models", build_directory)); } else { log("Segmentation: the model has not been made; see modules/Segment/README.md."); } diff --git a/modules/Core_ML/module.jai b/modules/Core_ML/module.jai index 08798e7..632df92 100644 --- a/modules/Core_ML/module.jai +++ b/modules/Core_ML/module.jai @@ -85,9 +85,18 @@ Coreml_Output :: struct { copy: *float; // when data is a copy: the default allocator's } -// Runs the model on input, a float32 tensor of shape. ok is false, with why, when it fails. +// Runs the model on input, a float32 tensor of shape: its output (model.output_name). ok is false, +// with why, when it fails. coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Coreml_Output, ok: bool, error: string { - result: Coreml_Output; + outs: [1] Coreml_Output; + ok, error := coreml_predict_outputs(model, input, shape, .[model.output_name], outs); + return outs[0], ok, error; +} + +// As coreml_predict, for a model with several outputs: those named, into outs (one each, in order), +// each released with coreml_release_output. On failure none is set. +coreml_predict_outputs :: (model: *Coreml_Model, input: [] float, shape: [] s64, names: [] string, outs: [] Coreml_Output) -> ok: bool, error: string { + assert(names.count == outs.count); init_objective_c(); pool := objc_new(NSAutoreleasePool); defer NSAutoreleasePool.drain(pool); @@ -96,13 +105,13 @@ coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Core strides := NewArray(shape.count, s64,, temp); step: s64 = 1; for < shape { strides[it_index] = step; step *= it; } - if step != input.count return result, false, tprint("the input has % values, not % for its shape %", input.count, step, shape); + if step != input.count return false, tprint("the input has % values, not % for its shape %", input.count, step, shape); error: *NSError; func_array_init :: #type (*void, Selector, *void, *NSObject, s64, *NSObject, *void, **NSError) -> *NSObject #c_call; array := (cast(func_array_init) objc_msgSend)(objc_alloc(objc_getClass("MLMultiArray")), selector("initWithDataPointer:shape:dataType:strides:deallocator:error:"), input.data, numbers(shape), MULTIARRAY_FLOAT32, numbers(strides), null, *error); - if !array return result, false, tprint("Core ML could not take the input: %", error_text(error)); + if !array return false, tprint("Core ML could not take the input: %", error_text(error)); autorelease(array); func_value :: #type (*void, Selector, *NSObject) -> *NSObject #c_call; @@ -112,18 +121,51 @@ coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Core dictionary := (cast(func_dictionary) objc_msgSend)(objc_getClass("NSDictionary"), selector("dictionaryWithObject:forKey:"), value, key); func_provider :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call; features := (cast(func_provider) objc_msgSend)(objc_alloc(objc_getClass("MLDictionaryFeatureProvider")), selector("initWithDictionary:error:"), dictionary, *error); - if !features return result, false, tprint("Core ML could not take the input: %", error_text(error)); + if !features return false, tprint("Core ML could not take the input: %", error_text(error)); autorelease(features); func_predict :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call; lock(*model.mutex); prediction := (cast(func_predict) objc_msgSend)(model.model, selector("predictionFromFeatures:error:"), features, *error); unlock(*model.mutex); - if !prediction return result, false, tprint("Core ML could not run the model: %", error_text(error)); + if !prediction return false, tprint("Core ML could not run the model: %", error_text(error)); - out_value := (cast(func_value) objc_msgSend)(prediction, selector("featureValueForName:"), NSString.getTempString(model.output_name)); + for names { + out, ok, read_error := read_output(prediction, it); + if !ok { + for 0..it_index - 1 coreml_release_output(*outs[it]); + return false, read_error; + } + outs[it_index] = out; + } + return true, ""; +} + +coreml_release_output :: (out: *Coreml_Output) { + if out.array release(out.array); + if out.copy free(out.copy,, context.default_allocator); + out.* = .{}; +} + +// coreml_predict, its output copied into output (which must hold exactly as many values). +coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string { + out, ok, error := coreml_predict(model, input, shape); + if !ok return false, error; + defer coreml_release_output(*out); + if out.count != output.count return false, tprint("the model gives % values, not %", out.count, output.count); + memcpy(output.data, out.data, out.count * size_of(float)); + return true, ""; +} + +#scope_file + +// The output named name of a prediction (an MLFeatureProvider), as float32 values in C order. +read_output :: (prediction: *NSObject, name: string) -> Coreml_Output, ok: bool, error: string { + result: Coreml_Output; + func_value :: #type (*void, Selector, *NSObject) -> *NSObject #c_call; + out_value := (cast(func_value) objc_msgSend)(prediction, selector("featureValueForName:"), NSString.getTempString(name)); out_array := ifx out_value then cast(*NSObject) objc_msgSend_typed(out_value, selector("multiArrayValue")) else null; - if !out_array return result, false, tprint("the model gives no %", model.output_name); + if !out_array return result, false, tprint("the model gives no %", name); out_shape := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("shape"))); out_strides := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("strides"))); @@ -159,23 +201,6 @@ coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Core return result, true, ""; } -coreml_release_output :: (out: *Coreml_Output) { - if out.array release(out.array); - if out.copy free(out.copy,, context.default_allocator); - out.* = .{}; -} - -// coreml_predict, its output copied into output (which must hold exactly as many values). -coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string { - out, ok, error := coreml_predict(model, input, shape); - if !ok return false, error; - defer coreml_release_output(*out); - if out.count != output.count return false, tprint("the model gives % values, not %", out.count, output.count); - memcpy(output.data, out.data, out.count * size_of(float)); - return true, ""; -} - -#scope_file MULTIARRAY_FLOAT32 :: 0x10000 | 32; // MLMultiArrayDataTypeFloat32 MULTIARRAY_FLOAT16 :: 0x10000 | 16; // MLMultiArrayDataTypeFloat16 diff --git a/modules/Segment/README.md b/modules/Segment/README.md index a199c00..d7a4130 100644 --- a/modules/Segment/README.md +++ b/modules/Segment/README.md @@ -18,7 +18,10 @@ conda run -n mcap-seg python tools/export_mask2former.py --check FOLDER_OF_IMAGE ``` This writes `models/mask2former-swin-base-coco-panoptic.onnx` (462 MB) and its class list, then checks -the .onnx against PyTorch on the images. Building the app copies `models.config` and the models into +the .onnx against PyTorch on the images. On macOS, `tools/export_mask2former_coreml.py --check +FOLDER_OF_IMAGES` then makes `models/mask2former-swin-base-coco-panoptic.mlpackage`, which runs on the +GPU through Core ML (`modules/Core_ML`), one image a run, about four times as fast as the CPU with the +same segments. Building the app copies `models.config` and the models into `bin/segment-models`. The models folder is not in git. ## Using it diff --git a/modules/Segment/models.config b/modules/Segment/models.config index 4da4183..8fdb411 100644 --- a/modules/Segment/models.config +++ b/modules/Segment/models.config @@ -6,6 +6,7 @@ [model mask2former-swin-base-coco-panoptic] title: "Mask2Former Swin-B, COCO panoptic (133 classes)" model: "mask2former-swin-base-coco-panoptic.onnx" +model_coreml: "mask2former-swin-base-coco-panoptic.mlpackage" # macOS: on the GPU (tools/export_mask2former_coreml.py makes it) classes: "mask2former-swin-base-coco-panoptic.classes.txt" image_size: 384 # squashed to 384 x 384, not keeping its shape, as the checkpoint's processor does mean: 0.485 0.456 0.406 diff --git a/modules/Segment/module.jai b/modules/Segment/module.jai index 10e4dac..e31975b 100644 --- a/modules/Segment/module.jai +++ b/modules/Segment/module.jai @@ -11,7 +11,7 @@ // transformers' post_process_panoptic_segmentation does (tools/export_mask2former.py checks that // they match). -Segment_Provider :: enum { CPU; CUDA; } +Segment_Provider :: enum { CPU; CUDA; COREML; } // Tightly packed or not: stride is the bytes from one row to the next. Segment_Image :: struct { @@ -23,6 +23,7 @@ Segment_Model_Info :: struct { name: string; title: string; model: string; // the .onnx, in the models folder + model_coreml: string; // macOS: the model as a Core ML package (tools/export_mask2former_coreml.py), run on the GPU; "" for none classes: string; // the class list, in the models folder image_size: s32 = 384; // the input's width and height: images are resized to it, not keeping their shape mean := float.[0.485, 0.456, 0.406]; // ImageNet's, per channel, after scaling to 0..1 @@ -50,6 +51,7 @@ Segment_Engine :: struct { env: *OrtEnv; memory_info: *OrtMemoryInfo; session: Ort_Session; + #if OS == .MACOS coreml: *Coreml_Model; // the model, when Core ML runs it (one image a run) cancel: *bool; // when set, and true, segment_images stops between batches: the rest fail as "cancelled" stats: Segment_Stats; // since it was made } @@ -114,6 +116,7 @@ segment_list_models :: (models_folder: string) -> [] Segment_Model_Info, error: case "title"; model.title = copy_string(value); case "model"; model.model = copy_string(value); case "classes"; model.classes = copy_string(value); + case "model_coreml"; model.model_coreml = copy_string(value); case "image_size"; if number_ok model.image_size = cast(s32) number; case "mean"; segment_parse_three(value, *model.mean, path, line_number); case "std"; segment_parse_three(value, *model.std, path, line_number); @@ -155,11 +158,29 @@ segment_create :: (models_folder: string, model_name: string, prefer_gpu := true HINT :: "make it with tools/export_mask2former.py (see modules/Segment/README.md) and rebuild"; path := tprint("%/%", models_folder, model.model); use_gpu := false; - if prefer_gpu { + #if OS == .MACOS if prefer_gpu { + // Core ML on the GPU; on any failure, onnxruntime on the CPU below, with the reason. + coreml_path := tprint("%/%", models_folder, model.model_coreml); + if !model.model_coreml || !file_exists(coreml_path) { + engine.gpu_explanation = sprint("Segmentation runs on the CPU: % is not there; tools/export_mask2former_coreml.py makes it.", ifx model.model_coreml then model.model_coreml else "its Core ML model"); + } else { + engine.coreml, ok, error = coreml_load(coreml_path, "pixel_values", "class_queries_logits"); + if ok { + engine.provider = .COREML; + ok, error = segment_warm_up(engine); + if ok return engine, true, ""; + coreml_free(engine.coreml); + engine.coreml = null; + engine.provider = .CPU; + } + engine.gpu_explanation = sprint("Segmentation runs on the CPU: Core ML could not run % on the GPU:\n%", model.model_coreml, error); + } + } + if prefer_gpu && OS != .MACOS { available, explanation := ort_cuda_check("Segmentation"); use_gpu = available; if !available engine.gpu_explanation = copy_string(explanation); - } else { + } else if !prefer_gpu { engine.gpu_explanation = copy_string("Segmentation runs on the CPU, as it was asked to."); } if use_gpu { @@ -190,6 +211,7 @@ segment_destroy :: (engine: *Segment_Engine) { if engine.provider == .CUDA ort_gpu_lock(); ort_release_session(*engine.session); if engine.provider == .CUDA ort_gpu_unlock(); + #if OS == .MACOS coreml_free(engine.coreml); if engine.memory_info ort.api.ReleaseMemoryInfo(engine.memory_info); if engine.env ort.api.ReleaseEnv(engine.env); for engine.classes { free(it.name); free(it.model_name); } @@ -198,6 +220,24 @@ segment_destroy :: (engine: *Segment_Engine) { free(engine); } +// Images a run: Core ML's model takes one. +segment_batch :: (engine: *Segment_Engine) -> s32 { + if engine.provider == { + case .CUDA; return engine.model.batch_gpu; + case .COREML; return 1; + } + return engine.model.batch_cpu; +} + +// What runs the model, for people: "CUDA", "Core ML" or "CPU". +segment_provider_name :: (engine: *Segment_Engine) -> string { + if engine.provider == { + case .CUDA; return "CUDA"; + case .COREML; return "Core ML"; + } + return "CPU"; +} + // The map size for an image of width x height: its shape, the longest side at most max_side. segment_map_size :: (width: s32, height: s32, max_side: s32) -> s32, s32 { scale := min(cast(float64) 1, cast(float64) max_side / max(width, height)); @@ -210,7 +250,7 @@ segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator results := NewArray(images.count, Segment_Result,, allocator); size := engine.model.image_size; plane := cast(s64) size * size; - batch := cast(s64) max(ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu, 1); + batch := cast(s64) max(segment_batch(engine), 1); input := NewArray(plane * 3 * batch, float,, context.default_allocator); defer array_free(input,, context.default_allocator); @@ -224,10 +264,10 @@ segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator start := current_time_monotonic(); for 0..count - 1 segment_prepare_image(engine.model, images[first + it], input.data + plane * 3 * it); shape := s64.[count, 3, size, size]; - outs: [2] Ort_Output; + outs: [2] Segment_Output; ran := current_time_monotonic(); engine.stats.prepare_us += to_microseconds(ran - start); - ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, array_view(input, 0, plane * 3 * count), shape, .["class_queries_logits", "masks_queries_logits"], outs); + ok, error := segment_run(engine, array_view(input, 0, plane * 3 * count), shape, outs); posted := current_time_monotonic(); engine.stats.model_us += to_microseconds(posted - ran); engine.stats.images += count; @@ -249,8 +289,7 @@ segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator cast(s32) queries, cast(s32) labels, cast(s32) mask_w, cast(s32) mask_h, image.width, image.height, allocator); } } - ort_release_output(*outs[0]); - ort_release_output(*outs[1]); + segment_release_outputs(engine, outs); } first += count; } @@ -524,14 +563,48 @@ segment_warm_up :: (engine: *Segment_Engine) -> ok: bool, error: string { size := engine.model.image_size; input := NewArray(3 * size * size, float, initialized = true,, context.default_allocator); defer array_free(input,, context.default_allocator); - outs: [2] Ort_Output; - ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, input, s64.[1, 3, size, size], .["class_queries_logits", "masks_queries_logits"], outs); + outs: [2] Segment_Output; + ok, error := segment_run(engine, input, s64.[1, 3, size, size], outs); if !ok return false, tprint("The first run on the GPU failed: %", error); - ort_release_output(*outs[0]); - ort_release_output(*outs[1]); + segment_release_outputs(engine, outs); return true, ""; } +// The two outputs, class and mask logits, as onnxruntime or Core ML gives them. +Segment_Output :: struct { + data: *float; + dims: [4] s64; + dim_count: s64; + ort: Ort_Output; + #if OS == .MACOS coreml: Coreml_Output; +} + +segment_run :: (engine: *Segment_Engine, input: [] float, shape: [] s64, outs: [] Segment_Output) -> ok: bool, error: string { + NAMES :: string.["class_queries_logits", "masks_queries_logits"]; + #if OS == .MACOS if engine.coreml { + coreml_outs: [2] Coreml_Output; + ok, error := coreml_predict_outputs(engine.coreml, input, shape, NAMES, coreml_outs); + if !ok return false, error; + for * outs { + it.* = .{data = coreml_outs[it_index].data, dim_count = min(coreml_outs[it_index].dim_count, 4), coreml = coreml_outs[it_index]}; + for d: 0..it.dim_count - 1 it.dims[d] = coreml_outs[it_index].dims[d]; + } + return true, ""; + } + ort_outs: [2] Ort_Output; + ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, input, shape, NAMES, ort_outs); + if !ok return false, error; + for * outs it.* = .{data = ort_outs[it_index].data, dims = ort_outs[it_index].dims, dim_count = ort_outs[it_index].dim_count, ort = ort_outs[it_index]}; + return true, ""; +} + +segment_release_outputs :: (engine: *Segment_Engine, outs: [] Segment_Output) { + for * outs { + ort_release_output(*it.ort); + #if OS == .MACOS coreml_release_output(*it.coreml); + } +} + // The class list: `class: INDEX thing|stuff "model's name" "name"`, a line each, in index order. segment_load_classes :: (engine: *Segment_Engine, path: string) -> ok: bool, error: string { text, ok := read_entire_file(path, log_errors = false,, temp); @@ -597,9 +670,11 @@ segment_strip_comment :: (line: string) -> string { #import "Basic"; #import "String"; #import "File"; +#import "File_Utilities"; #import "Math"; #import "System"; #import "Hash"; #import "Sort"; #import "Capture_File"; // segment_mcap_chunks #import "onnxruntime"; +#if OS == .MACOS #import "Core_ML"; diff --git a/src/app/segment_ui.jai b/src/app/segment_ui.jai index b6e3baf..d32df32 100644 --- a/src/app/segment_ui.jai +++ b/src/app/segment_ui.jai @@ -198,7 +198,7 @@ segment_ui_scan_task :: () { return; } segment_ui.engine = engine; - log("[Segment] % on % in % ms.%", SEGMENT_MODEL, engine.provider, to_milliseconds(current_time_monotonic() - start), ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else ""); + log("[Segment] % on % in % ms.%", SEGMENT_MODEL, segment_provider_name(engine), to_milliseconds(current_time_monotonic() - start), ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else ""); } engine := segment_ui.engine; // Cancelling (or closing the app) stops it between batches of the engine's, rather than waiting the @@ -335,7 +335,7 @@ segment_window_contents :: () { using ImGui; // The model. model_file := tprint("%/%.onnx", segment_models_folder(), SEGMENT_MODEL); if !busy && segment_ui.engine { - Text("% on %", SEGMENT_MODEL, ifx segment_ui.engine.provider == .CUDA then "the GPU" else "the CPU"); + Text("% on %", SEGMENT_MODEL, ifx segment_ui.engine.provider != .CPU then tprint("the GPU (%)", segment_provider_name(segment_ui.engine)) else "the CPU"); if segment_ui.engine.gpu_explanation TextColored(app_theme.palette.dim, "%", segment_ui.engine.gpu_explanation); } else if !busy && segment_ui.engine_error { TextColored(app_theme.palette.warning, "% could not start: %", SEGMENT_MODEL, segment_ui.engine_error); diff --git a/tools/export_mask2former_coreml.py b/tools/export_mask2former_coreml.py new file mode 100644 index 0000000..2a8a796 --- /dev/null +++ b/tools/export_mask2former_coreml.py @@ -0,0 +1,173 @@ +# Converts Mask2Former (facebook/mask2former-swin-base-coco-panoptic) to a Core ML model for +# modules/Segment on macOS, then checks it against the ONNX model on a folder of images, segment by +# segment. Run once, on a Mac, after tools/export_mask2former.py; mcap only uses the .mlpackage it writes. +# +# python3 -m venv ~/mcap-coreml +# ~/mcap-coreml/bin/pip install torch==2.7.0 transformers coremltools==9.0 onnxruntime pillow scipy +# ~/mcap-coreml/bin/pip install --no-deps timm +# ~/mcap-coreml/bin/python tools/export_mask2former_coreml.py --check FOLDER_OF_IMAGES +# +# The model as tools/export_mask2former.py exports it (pixel_values [1, 3, 384, 384] to +# class_queries_logits [1, 100, 134] and masks_queries_logits [1, 100, 96, 96]), as an ML Program in +# float32, one image a run (a batch of 4 did not convert). Three things coremltools 9.0 cannot take +# are worked around for the conversion only, none changing what the model computes: +# - it casts a one-element tensor with int(array), which NumPy 2 refuses; +# - its meshgrid refuses the inputs Mask2Former's traced linspaces come out as; +# - Core ML takes tensors of at most 5 dimensions, and the multi-scale deformable attention makes 6 +# (batch, queries, heads, levels, points, xy): it is run a level at a time instead, which gives +# the same output in PyTorch (checked below) with at most 5. +# Measured on an M5 Pro: 42 ms an image on the GPU against 166 ms for onnxruntime on the CPU, with +# the same segments (IoU 1.0) on every test image. + +import argparse, glob, os, sys, time +import numpy as np +import torch +from torch import nn +from PIL import Image +import coremltools as ct +from coremltools.converters.mil import Builder as mb +from coremltools.converters.mil.frontend.torch import ops as ct_ops +from coremltools.converters.mil.frontend.torch.torch_op_registry import register_torch_op +from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation +from transformers.utils import logging +import transformers.models.mask2former.modeling_mask2former as modeling + +logging.set_verbosity_error() +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from export_mask2former import Logits + +CHECKPOINT = "facebook/mask2former-swin-base-coco-panoptic" +MODELS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "modules", "Segment", "models") +ONNX = os.path.join(MODELS, "mask2former-swin-base-coco-panoptic.onnx") +OUTPUT = os.path.join(MODELS, "mask2former-swin-base-coco-panoptic.mlpackage") +SIZE = 384 + + +_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 + + +@register_torch_op(override=True) +def meshgrid(context, node): + """torch.meshgrid, from its inputs flattened to 1-d first.""" + inputs = ct_ops._get_inputs(context, node, expected=[1, 2]) + indexing = inputs[1].val if len(inputs) > 1 else "ij" + indexing = ct_ops._get_kwinputs(context, node, "indexing", default=[indexing])[0] + if hasattr(indexing, "val"): + indexing = indexing.val + tensors = [mb.reshape(x=t, shape=[-1], name=f"{node.name}_flat_{i}") for i, t in enumerate(inputs[0])] + sizes = [t.shape[0] for t in tensors] + grids = [] + for i, t in enumerate(tensors): + view = [1] * len(tensors); view[i] = -1 + reps = list(sizes); reps[i] = 1 + grid = mb.tile(x=mb.reshape(x=t, shape=view, name=f"{node.name}_view_{i}"), reps=reps, name=f"{node.name}_grid_{i}") + if indexing == "xy": + grid = mb.transpose(x=grid, perm=[1, 0] + list(range(2, len(tensors))), name=f"{node.name}_xy_{i}") + grids.append(grid) + context.add(tuple(grids), node.name) + + +def deformable_attention_rank5(self, hidden_states, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, + position_embeddings=None, reference_points=None, spatial_shapes_list=None, level_start_index=None, + output_attentions=False): + """Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention.forward, a level at a time.""" + if position_embeddings is not None: + hidden_states = self.with_pos_embed(hidden_states, position_embeddings) + if reference_points.shape[-1] != 2: + raise ValueError("only 2-d reference points are rewritten") + B, Q, _ = hidden_states.shape + _, S, _ = encoder_hidden_states.shape + H, L, P, D = self.n_heads, self.n_levels, self.n_points, self.d_model // self.n_heads + value = self.value_proj(encoder_hidden_states) + if attention_mask is not None: + value = value.masked_fill(attention_mask[..., None], float(0)) + value = value.view(B, S, H, D) + offsets = self.sampling_offsets(hidden_states).view(B, Q, H, L, P * 2) + weights = nn.functional.softmax(self.attention_weights(hidden_states).view(B, Q, H, L * P), -1) + weights = weights.transpose(1, 2).reshape(B * H, 1, Q, L * P) + values = value.split([h * w for h, w in spatial_shapes_list], dim=1) + sampled = [] + for level, (h, w) in enumerate(spatial_shapes_list): + value_l = values[level].flatten(2).transpose(1, 2).reshape(B * H, D, h, w) + offset_l = offsets[:, :, :, level].reshape(B, Q, H, P, 2) + location = reference_points[:, :, None, level, None, :] + offset_l / torch.tensor([w, h], dtype=offset_l.dtype) + grid = (2 * location - 1).transpose(1, 2).flatten(0, 1) + sampled.append(nn.functional.grid_sample(value_l, grid, mode="bilinear", padding_mode="zeros", align_corners=False)) + output = (torch.stack(sampled, dim=-2).flatten(-2) * weights).sum(-1).view(B, H * D, Q).transpose(1, 2).contiguous() + return self.output_proj(output), None + + +def segments_of(processor, class_logits, mask_logits): + holder = type("Outputs", (), {"class_queries_logits": torch.from_numpy(class_logits), "masks_queries_logits": torch.from_numpy(mask_logits)})() + return processor.post_process_panoptic_segmentation(holder, target_sizes=[(SIZE, SIZE)])[0] + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--check", metavar="FOLDER", help="images to check the Core ML model against the ONNX one on") + args = parser.parse_args() + + processor = AutoImageProcessor.from_pretrained(CHECKPOINT) + model = Mask2FormerForUniversalSegmentation.from_pretrained(CHECKPOINT).eval() + example = torch.zeros(1, 3, SIZE, SIZE) + with torch.no_grad(): + original = Logits(model)(torch.randn(1, 3, SIZE, SIZE, generator=torch.Generator().manual_seed(0))) + modeling.Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention.forward = deformable_attention_rank5 + rewritten = Logits(model)(torch.randn(1, 3, SIZE, SIZE, generator=torch.Generator().manual_seed(0))) + difference = max((a - b).abs().max().item() for a, b in zip(original, rewritten)) + print(f"The attention a level at a time, against the original in PyTorch: max difference {difference:.2e}") + if difference > 1e-4: + sys.exit("The rewritten attention does not give the same output.") + traced = torch.jit.trace(Logits(model).eval(), example) + start = time.perf_counter() + ml = ct.convert(traced, inputs=[ct.TensorType(name="pixel_values", shape=(1, 3, SIZE, SIZE))], + outputs=[ct.TensorType(name="class_queries_logits"), ct.TensorType(name="masks_queries_logits")], + convert_to="mlprogram", compute_precision=ct.precision.FLOAT32, minimum_deployment_target=ct.target.macOS14) + ml.short_description = f"{CHECKPOINT}, class and mask logits, float32, for mcap" + ml.save(OUTPUT) + print(f"Wrote {OUTPUT} in {time.perf_counter() - start:.0f} s.") + + if not args.check: + return + import onnxruntime as ort + paths = sorted(p for p in glob.glob(os.path.join(args.check, "**", "*"), recursive=True) + if p.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))) + if not paths: + sys.exit(f"No images in {args.check}.") + session = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"]) + coreml = ct.models.MLModel(OUTPUT, compute_units=ct.ComputeUnit.CPU_AND_GPU) + total, worst_iou, unmatched, times = 0, 1.0, 0, [] + for path in paths: + pixels = processor(images=Image.open(path).convert("RGB"), return_tensors="np")["pixel_values"].astype(np.float32) + ref = session.run(["class_queries_logits", "masks_queries_logits"], {"pixel_values": pixels}) + start = time.perf_counter() + got = coreml.predict({"pixel_values": pixels}) + times.append(time.perf_counter() - start) + a = segments_of(processor, ref[0], ref[1]) + b = segments_of(processor, got["class_queries_logits"].astype(np.float32), got["masks_queries_logits"].astype(np.float32)) + for sa in a["segments_info"]: + total += 1 + best = 0.0 + for sb in b["segments_info"]: + if sb["label_id"] == sa["label_id"]: + ma, mbk = (a["segmentation"] == sa["id"]).numpy(), (b["segmentation"] == sb["id"]).numpy() + best = max(best, (ma & mbk).sum() / max((ma | mbk).sum(), 1)) + if best == 0: + unmatched += 1 + else: + worst_iou = min(worst_iou, best) + print(f"{len(paths)} images, {total} segments: worst IoU against the ONNX model {worst_iou:.4f}, unmatched {unmatched}; " + f"{np.median(times[1:] or times) * 1000:.0f} ms an image on the GPU") + if unmatched or worst_iou < 0.98: + sys.exit("The Core ML model segments differently from the ONNX one.") + + +if __name__ == "__main__": + main() diff --git a/tools/segment_probe.jai b/tools/segment_probe.jai index ae599b0..1e82d0a 100644 --- a/tools/segment_probe.jai +++ b/tools/segment_probe.jai @@ -33,7 +33,7 @@ main :: () { engine, ok, error := segment_create(models_folder, model_name, prefer_gpu); if !ok { print("Cannot start %: %\n", model_name, error); exit(1); } defer segment_destroy(engine); - print("% on %, % classes (made in % ms)%\n", model_name, engine.provider, engine.classes.count, ms(start), + print("% on %, % classes (made in % ms)%\n", model_name, segment_provider_name(engine), engine.classes.count, ms(start), ifx engine.gpu_explanation then tprint(": %", engine.gpu_explanation) else ""); failures := 0; @@ -129,7 +129,7 @@ main :: () { print("\n% images: decoding % ms each, segmenting % ms each (model and post-processing, batch %)\n", done, formatFloat(decode_ms / max(done, 1), trailing_width = 0), formatFloat(model_ms / max(done, 1), trailing_width = 0), - ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu); + segment_batch(engine)); per := max(engine.stats.images, 1) * 1000.0; print(" of which: preparing % ms, the model % ms, making segments % ms, per image\n", formatFloat(engine.stats.prepare_us / per, trailing_width = 1), formatFloat(engine.stats.model_us / per, trailing_width = 1), formatFloat(engine.stats.post_us / per, trailing_width = 1));