[0.1.32] macOS phase 5: Mask2Former segmentation on the GPU through native Core ML
This commit is contained in:
parent
a1ab9f3a96
commit
91f6cf768f
11
README.md
11
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.
|
||||
|
||||
|
||||
@ -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;
|
||||
|
||||
@ -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.");
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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";
|
||||
|
||||
@ -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);
|
||||
|
||||
173
tools/export_mask2former_coreml.py
Normal file
173
tools/export_mask2former_coreml.py
Normal file
@ -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()
|
||||
@ -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));
|
||||
|
||||
Loading…
Reference in New Issue
Block a user