| |
| """Inspection helpers for the SAM 2 recipe (sam2-hiera-tiny.py). CPU only, read-only on inputs. |
| |
| _sam2.py ops A.pte [B.pte] every delegated op call with its argument types/shapes; with B, |
| the op-signature set difference (ops/arity/ranks in A not in B) |
| _sam2.py consts A.pte B.pte how many delegate constants of A are byte-identical in B |
| _sam2.py io A.pte method names, program inputs/outputs, non-delegated kernels |
| """ |
|
|
| import collections |
| import hashlib |
| import json |
| import sys |
|
|
|
|
| def _program(data): |
| from executorch.exir._serialize._program import deserialize_pte_binary |
|
|
| pte = deserialize_pte_binary(data) |
| return pte, (pte.program if hasattr(pte, "program") else pte) |
|
|
|
|
| def _blobs(data): |
| pte, prog = _program(data) |
| seg_base = int.from_bytes(data[24:32], "little") |
| out = [] |
| for plan in prog.execution_plan: |
| for d in plan.delegates: |
| proc = d.processed |
| if proc.location.name == "INLINE": |
| out.append( |
| ( |
| plan.name, |
| d.id, |
| bytes(prog.backend_delegate_data[proc.index].data), |
| ) |
| ) |
| else: |
| seg = prog.segments[proc.index] |
| out.append( |
| ( |
| plan.name, |
| d.id, |
| data[seg_base + seg.offset : seg_base + seg.offset + seg.size], |
| ) |
| ) |
| nd = getattr(pte, "named_data", None) |
| resolve = ( |
| ( |
| lambda k: bytes(nd.buffers[nd.pte_data[k].buffer_index]) |
| if k in nd.pte_data |
| else None |
| ) |
| if nd |
| else (lambda k: None) |
| ) |
| return out, resolve |
|
|
|
|
| def vk_graph(blob): |
| from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( |
| extract_vk_flatbuffer, |
| flatbuffer_to_vk_graph, |
| ) |
|
|
| return flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob)) |
|
|
|
|
| def describe(g, vid): |
| v = g.values[vid].value |
| n = type(v).__name__ |
| if n == "VkTensor": |
| kind = "const" if v.constant_id >= 0 else "tensor" |
| return f"{kind}{list(v.dims)}:{v.datatype.name if hasattr(v.datatype, 'name') else v.datatype}" |
| if n == "Int": |
| return f"Int={v.int_val}" |
| if n == "Double": |
| return f"Double={v.double_val:g}" |
| if n == "Bool": |
| return f"Bool={v.bool_val}" |
| if n in ("IntList", "DoubleList", "BoolList"): |
| return f"{n}={list(v.items)}" |
| if n == "ValueList": |
| return "ValueList[" + ", ".join(describe(g, i) for i in v.items) + "]" |
| if n == "String": |
| return f"String={v.string_val!r}" |
| return n |
|
|
|
|
| def ops(path): |
| data = open(path, "rb").read() |
| blobs, _ = _blobs(data) |
| calls = [] |
| for plan, did, blob in blobs: |
| g = vk_graph(blob) |
| for c in g.chain: |
| calls.append( |
| {"plan": plan, "op": c.name, "args": [describe(g, a) for a in c.args]} |
| ) |
| return calls |
|
|
|
|
| def signature(call): |
| """Op name + arity + per-argument kind/rank (shape-agnostic) for set comparisons.""" |
| parts = [] |
| for a in call["args"]: |
| if a.startswith(("tensor[", "const[")): |
| dims = json.loads(a[a.index("[") : a.index("]") + 1]) |
| parts.append(f"{a.split('[')[0]}{len(dims)}d:{a.rsplit(':', 1)[1]}") |
| else: |
| parts.append(a.split("=")[0]) |
| return f"{call['op']}({', '.join(parts)})" |
|
|
|
|
| def consts(path): |
| from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( |
| VulkanDelegateHeader, |
| ) |
|
|
| data = open(path, "rb").read() |
| blobs, resolve = _blobs(data) |
| out = [] |
| for plan, did, blob in blobs: |
| h = VulkanDelegateHeader.from_bytes( |
| blob[: VulkanDelegateHeader.EXPECTED_LENGTH] |
| ) |
| g = vk_graph(blob) |
| raw = blob[h.bytes_offset : h.bytes_offset + h.bytes_size] |
| for v in g.values: |
| t = v.value |
| if type(t).__name__ != "VkTensor" or t.constant_id < 0: |
| continue |
| c = g.constants[t.constant_id] |
| b = ( |
| resolve(c.named_key) |
| if c.named_key |
| else raw[c.offset : c.offset + c.length] |
| ) |
| out.append( |
| { |
| "dims": list(t.dims), |
| "sha": hashlib.sha256(b).hexdigest(), |
| "bytes": len(b), |
| } |
| ) |
| return out |
|
|
|
|
| def io(path): |
| data = open(path, "rb").read() |
| _, prog = _program(data) |
| res = {"file": path, "bytes": len(data), "methods": {}} |
| for plan in prog.execution_plan: |
|
|
| def d(i): |
| v = plan.values[i].val |
| return ( |
| f"{list(v.sizes)}:{v.scalar_type.name}" |
| if hasattr(v, "sizes") |
| else type(v).__name__ |
| ) |
|
|
| names = [f"{o.name}.{o.overload}" for o in plan.operators] |
| kernels = collections.Counter( |
| names[ins.instr_args.op_index] |
| for ch in plan.chains |
| for ins in ch.instructions |
| if type(ins.instr_args).__name__ == "KernelCall" |
| ) |
| res["methods"][plan.name] = { |
| "inputs": [d(i) for i in plan.inputs], |
| "outputs": [d(i) for i in plan.outputs], |
| "delegates": [x.id for x in plan.delegates], |
| "non_delegated_kernels": dict(kernels), |
| } |
| return res |
|
|
|
|
| OFFICIAL = { |
| "sam2.1": ("facebook/sam2.1-hiera-tiny", "de431c4043854a71d8101e17995dfe596bf101a5", "sam2.1_hiera_tiny.pt", |
| "7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69"), |
| } |
|
|
|
|
| def verify_official(weights="sam2.1"): |
| import importlib.util |
| import pathlib |
| import tempfile |
|
|
| import torch |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
|
|
| here = pathlib.Path(__file__).resolve().parent |
| spec = importlib.util.spec_from_file_location("sam2_recipe", here / "sam2-hiera-tiny.py") |
| recipe = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(recipe) |
| snap, digest = recipe.snapshot(weights) |
| hf = load_file(str(snap / "model.safetensors")) |
| repo, rev, name, want = OFFICIAL[weights] |
| with tempfile.TemporaryDirectory() as tmp: |
| p = hf_hub_download(repo, name, revision=rev, local_dir=tmp) |
| got = hashlib.sha256(open(p, "rb").read()).hexdigest() |
| if got != want: |
| sys.exit(f"{name}: sha256 {got} != {want}") |
| orig = torch.load(p, map_location="cpu", weights_only=True)["model"] |
| h = lambda t: hashlib.sha256(t.detach().contiguous().numpy().tobytes()).hexdigest() |
| pool = {h(v): k for k, v in orig.items()} |
| matched, unmatched = 0, [] |
| for k, v in hf.items(): |
| if h(v) in pool: |
| matched += 1 |
| else: |
| unmatched.append(k) |
| notes = {} |
| if "prompt_encoder.point_embed.weight" in unmatched: |
| pe = hf["prompt_encoder.point_embed.weight"] |
| if all(torch.equal(pe[i], orig[f"sam_prompt_encoder.point_embeddings.{i}.weight"][0]) for i in range(pe.shape[0])): |
| unmatched.remove("prompt_encoder.point_embed.weight") |
| matched += 1 |
| notes["prompt_encoder.point_embed.weight"] = "rows == sam_prompt_encoder.point_embeddings.{0..3}.weight" |
| return {"hf": f"{repo}@{rev}/model.safetensors", "hf_sha256": digest, "official": f"{repo}@{rev}/{name}", |
| "official_sha256": want, "hf_tensors": len(hf), "official_tensors": len(orig), |
| "hf_tensors_bit_identical_to_official": matched, "unmatched": unmatched, "notes": notes} |
|
|
|
|
| def main(): |
| cmd, args = sys.argv[1], sys.argv[2:] |
| if cmd == "io": |
| print(json.dumps(io(args[0]), indent=1)) |
| elif cmd == "ops": |
| a = ops(args[0]) |
| if len(args) == 1: |
| for c in a: |
| print(c["op"], ", ".join(c["args"])) |
| return |
| b = ops(args[1]) |
| sa, sb = ( |
| collections.Counter(map(signature, a)), |
| collections.Counter(map(signature, b)), |
| ) |
| print( |
| json.dumps( |
| { |
| "ops_count": [sum(sa.values()), sum(sb.values())], |
| "op_types_only_in_a": sorted( |
| {c["op"] for c in a} - {c["op"] for c in b} |
| ), |
| "op_types_only_in_b": sorted( |
| {c["op"] for c in b} - {c["op"] for c in a} |
| ), |
| "signatures_only_in_a": { |
| k: v for k, v in sorted(sa.items()) if k not in sb |
| }, |
| "signatures_only_in_b": { |
| k: v for k, v in sorted(sb.items()) if k not in sa |
| }, |
| }, |
| indent=1, |
| ) |
| ) |
| elif cmd == "verify-official": |
| print(json.dumps(verify_official(*(args or ["sam2.1"])), indent=1)) |
| elif cmd == "consts": |
| a, b = consts(args[0]), consts(args[1]) |
| hb, ha = ( |
| collections.Counter(c["sha"] for c in b), |
| collections.Counter(c["sha"] for c in a), |
| ) |
| print( |
| json.dumps( |
| { |
| "a_constants": len(a), |
| "b_constants": len(b), |
| "a_found_in_b": sum(c["sha"] in hb for c in a), |
| "b_found_in_a": sum(c["sha"] in ha for c in b), |
| "a_missing": [c["dims"] for c in a if c["sha"] not in hb], |
| "b_missing": [c["dims"] for c in b if c["sha"] not in ha], |
| }, |
| indent=1, |
| ) |
| ) |
| else: |
| sys.exit(__doc__) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
