Extract and analyze the structure of a mlProgram CoreML models with CoreML7 spec version. This script outputs the layer names, their corresponding offsets, and tensor dimensions in a readable format. It's particularly useful for debugging and identifying unexpected sizes in exported CoreML models.

coreml2text.py

# Copyright (c) 2024 Andrew Grigorev <andrew@ei-grad.ru>
#
# This software is licensed under the MIT License. Redistribution and use are permitted
# provided that the license terms are met. For details, see the MIT License:
# https://opensource.org/licenses/MIT

import sys

from functools import reduce
from operator import mul

import coremltools


type_sizes = {
    "FLOAT16": 2,
    "FLOAT32": 4,
    "INT32": 4,
    "UINT16": 2,
}

model = coremltools.models.MLModel(sys.argv[1])
spec = model.get_spec()
main_function = spec.mlProgram.functions["main"]
for i in main_function.block_specializations['CoreML7'].operations:
    op_name = i.attributes['name'].immediateValue.tensor.strings.values[0]
    v = i.attributes["val"]
    t = v.type.tensorType
    t_name = coremltools.proto.MIL_pb2.DataType.Name(t.dataType)
    match v.WhichOneof("value"):
        case "immediateValue":
            iv = v.immediateValue
            assert iv.WhichOneof("value") == "tensor"
            rv = iv.tensor
            rv_type = iv.tensor.WhichOneof("value")
            rv_values = repr(getattr(iv.tensor, rv_type).values)
            t_str = f"inline {rv_type}({rv_values})"
        case "blobFileValue":
            shape = [i.constant.size for i in t.dimensions]
            assert t_name in type_sizes
            byte_size = reduce(mul, shape) * type_sizes[t_name]
            t_str = f"shape={'x'.join(str(i) for i in shape)} offset={v.blobFileValue.offset} byte_size={byte_size}"
        case _:
            t_str = "unknown"
    print(
        op_name,
        i.type,
        t_name,
        t_str.replace('\n', ''),
    )