re-ai-model

AI model reverse engineering (ONNX / PyTorch / Safetensors)

When to use / when not to use

Use: when you have model files such as .onnx / .pt / .pth / .safetensors / .tflite and need to recover the network structure and extract weights
Use:file-level watermarkdetection (where things are hidden: weight pattern / metadata / tensor hash / embedding anomalies—when you suspect a model was copied/fine-tuned from the original, check the file side first)
Use: when the model file itself is the payload—data hidden in weights, malicious pickle packaged by torch.save, backdoored/poisoned models (download-and-execute type samples)
Do not use:behavior-level watermark(what the model exhibits: trigger-triggered responses / query responses / black-box fingerprints) and API behavior-layer attacks (see [[re-ai-attack]])
Do not use: pure inference scripts/training code (that is source code, see [[re-script-deob]])
Do not use: when the model is packaged into an executable (PyInstaller/pyarmor etc.)—first unpack with [[re-binary-core]], then bring the extracted model files back to this skill
Boundary: this skill's positioning = model file parsing / structure analysis / weight analysis / file-level watermarking—malicious model determination, poisoning/backdoor behavior side, and attribution forensics belong to the forensics domain (to be handled by a future independent re-ai-malware skill; currently such needs temporarily fall within this skill's scope and are handled under the security boundary (pitfall 2 pickle isolation), with the behavior side redirected to [[re-ai-attack]])
Note:security tip—do not directly torch.load unknown pkl files(pickle deserialization can execute arbitrary code, see pitfall 2); all loading of unknown pkl defaults to an isolated environment ([[re-analyze/platform-tips]] sandbox highest principle), read first then run; model parsing/weight extraction are static steps and can skip the sandbox

Tool preparation

Refer to [[re-analyze/platform-tips]]—model files are commonly GB-scale; for static analysis follow the "static first (large samples)" approach: first format identification and structure parsing, extract weights as needed, without loading entirely into memory (pitfall 1).

python3 — the basis for all parsing scripts

Linux: apt install python3 / dnf install python3 / pacman -S python
macOS: brew install python3；Windows: choco install python
Verify: python3 --version(all scripts in this skill are Python 3)

onnx (pip, Python 3.10+) — the mainstay of ONNX parsing

pip install onnx(official PyPI; currently onnx requires Python 3.10+, specific version subject to PyPI Requires-Python; comes with protobuf dependency and onnx.proto3 type definitions)
Verify: python3 -c "import onnx; print(onnx.__version__)"

netron (pip, Python 3) — model visualization

pip install netron(official PyPI, no upper bound on Python version); standalone desktop version optional: macOS brew install --cask netron、Windows winget install netron、Linux snap snap install netron
Verify: netron --help has output (pip show netron check version)
Usage: netron model.onnx(start an http service locally and open a browser for visualization;--no-browser headless mode)

torch (pip, Python >=3.10) — PyTorch model loading

Linux/Windows: pip install torch(the default PyPI wheel is the full CUDA-enabled package, several GB; for CPU-only analysis use pip install torch --index-url https://download.pytorch.org/whl/cpu）
Python version lower bound: currently torch 2.x requires Python >=3.10（PyPI requires-python; torchvision synchronized to >=3.10) — to install in a 3.9 environment you must pin to an older version that still supports 3.9 (2.8.x and earlier). For the specific combinations of torch / torchvision / Python, refer to the official compatibility matrix; don't generalize based on "3.9+"
macOS: pip install torch(official wheels are CPU/arm64)
Verify: python3 -c "import torch; print(torch.__version__)"
Security note: torch.load is based on pickle — do not directly load unknown pkl files; PyTorch 2.6+ defaults to weights_only=True, older versions/explicit weights_only=False still carry arbitrary code execution risk; for unknown models first unzip -l/xxd rough inspect (pitfall 2), load in an isolated environment with weights_only=True and convert to safetensors if possible

safetensors (pip, Python 3) — safe format reading

pip install safetensors; verify: python3 -c "import safetensors; print(safetensors.__version__)"
is designed to have no code execution (pure data + JSON header), and is an alternative analysis entry point for unknown pkl (one of the countermeasures for pitfall 2)

protobuf / protoc (onnx is proto) — underlying format

Python bindings: pip install protobuf(currently 7.x requires Python >=3.10, the lower bound has been rising since 6.x; 3.8/3.9 need to pin the corresponding older version. onnx already bundles its dependencies, usually no need to install separately)
protoc compiler tool: Debian/Ubuntu apt install protobuf-compiler、Fedora dnf install protobuf-compiler、Arch pacman -S protobuf、macOS brew install protobuf
Verify: protoc --version；python3 -c "import google.protobuf; print(google.protobuf.__version__)"

Operation steps

Execute in order, archive each step's artifacts (model hash/structure summary/weight list) with sha256 + path ([[re-ioc]] evidence chain).

Model format identification (onnx / safetensors / pytorch pkl / tflite)：
```
file model.bin && sha256sum model.bin > model.sha256 && xxd model.bin | head -2
```
- ONNX: no fixed magic number — a protobuf stream, the first field is ir_version (tag 0x08 + varint value;the value increases with the ONNX version: IR 8 corresponds to ONNX 1.10, current releases are already at IR 13 — so the second byte is not the constant 08, you can't use 08 08 12 as an identification feature); producer_name (tag 0x12, length-prefixed) per the spec only SHOULD appear, it is not guaranteed to immediately follow;onnx.checker.check_model can verify validity
- Safetensors: first 8 bytes = little-endian u64 header length (struct.unpack("<Q", data[:8])), followed by a JSON header (tensor name/shape/dtype/offset)
- PyTorch：file shows Zip archive (PK\x03\x04 header) —unzip -l model.pt look at entries (state_dict contains data.pkl; torch.jit.script contains data.pkl/constants.pkl/bytecode.pkl); old-style pure pkl is a bare pickle stream (no PK header) —do not directly load, first do a rough inspection with xxd/strings (pitfall 2)
- TFLite: flatbuffers stream — no fixed magic number, but bytes 4–7 are the file identifier TFL3(corresponding to the schema's __model_identifier field), visible via xxd/strings; for structure parsing use Netron (supports .tflite visualization), the weight extraction branch follows the same idea as onnx (flatbuffers parsing; when beyond the depth of this skill, annotate "structured dump as authoritative")
- After identification, follow the corresponding branch by format; if unsure, first do [[re-triage]] preliminary survey (entropy/strings features)

Structure parsing (graph/layers/operators)：

netron model.onnx --no-browser          # 可视化（有图形界面再开浏览器）
python3 - <<'PY'
import onnx
m = onnx.load("model.onnx")
print("producer:", m.producer_name, m.producer_version)   # 框架/优化器指纹（见坑 3）
g = m.graph
print("inputs:", [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in g.input])
print("nodes:", len(g.node), "initializers:", len(g.initializer), "outputs:", [o.name for o in g.output])
for n in g.node[:20]:
    print(n.op_type, n.name, list(n.input), "->", list(n.output))
PY

PyTorch side:torch.jit.load get ScriptModule can print model.graph(TorchScript structure);torch.load state_dict has only tensorswithout network structure— The structure is in the training/inference script and must be reconstructed together with the source code (see cross-domain [[re-script-deob]])
Focus points: operator sequence (structural fingerprints such as convolution/attention), input/output tensor shapes, constant node locations (where the weights are hidden)

Weight extraction (tensor dump)：

mkdir -p weights
# ONNX：逐 tensor 惰性落地（大模型见坑 1）
# 注意：这不是"流式"——to_array() 对使用 external data 的 tensor 会先把该张量【完整】读进 ndarray，
#       只是"一次一个张量"而不是"一次整个模型"；显存/内存峰值取决于最大的那个张量
python3 - <<'PY'
import onnx, numpy as np
m = onnx.load("model.onnx", load_external_data=False)     # 只载图结构，不载外部权重
for init in m.graph.initializer:
    # base_dir 默认为空字符串——外部权重不在 CWD 时必须显式指定其所在目录
    arr = onnx.numpy_helper.to_array(init, base_dir=".")   # 该张量在此处被完整 materialize
    np.save(f"weights/{init.name.replace('/', '_')}.npy", arr)
    print(init.name, arr.shape, arr.dtype)
    del arr                                                # 及时释放，为下一张量让出内存
PY
# Safetensors：惰性按张量读取（不整载内存）
python3 - <<'PY'
from safetensors import safe_open
with safe_open("model.safetensors", framework="numpy") as f:
    print(len(f.keys()), "tensors")
    for k in list(f.keys())[:10]:
        t = f.get_tensor(k); print(k, t.shape, t.dtype)
PY
# PyTorch state_dict（weights_only 安全加载，见坑 2）
python3 - <<'PY'
import torch
sd = torch.load("model.pth", weights_only=True)
for k, v in list(sd.items())[:10]:
    print(k, tuple(v.shape) if hasattr(v, "shape") else type(v))
PY

Output: weight inventory (tensor name/shape/dtype/value summary) + npy archive — this is the raw material for watermark detection and theft determination

Model watermark/fingerprint detection (embedded weights)：
- Weight level: per-tensor statistics (min/max/mean/std, histogram binning) compared with the suspected original model; watermarks are often embedded in specific layers (first-layer convolution bias, normalization scale, embedding matrix rows), mostly as low-bit perturbations — inspect the low-bit patterns and numerical distribution anomalies of key tensors, rather than exact equality (Pitfall 4)
- Fingerprint level: full-weight sha256 digest, per-layer tensor hash sequence; per-layer distance (L2/cosine) comparison with the candidate original, outputting a "per-layer distance heatmap"
- Behavior level: run inference with the same test input set on both models, compare logits and activation distributions — retraining/distillation thieves have different weights but similar behavior
- Output: similarity matrix + heatmap, with the conclusion stating the basis for the determination and the threshold
Model theft determination (architecture similarity)：
```
python3 - <<'PY'
import onnx
from difflib import SequenceMatcher
m1 = onnx.load("suspect.onnx"); m2 = onnx.load("original.onnx")
s1 = [n.op_type for n in m1.graph.node]; s2 = [n.op_type for n in m2.graph.node]
print("op-seq similarity:", SequenceMatcher(None, s1, s2).ratio())   # 算子序列相似度
# 层对齐后逐层权重余弦/L2 距离：见步骤 4 的每层距离矩阵
PY
```
- Dimensions: operator sequence (SequenceMatcher), input/output shapes, per-layer weight cosine/L2, output logits distance
- Determination discipline: a single dimension is weak evidence (same architecture with different training = normal); only when multiple dimensions agree and the weight distributions are highly close (e.g., per-layer cosine >0.99) can theft be claimed, with the values for each dimension noted in the conclusion
- Hidden payload inspection:m.metadata_props(strings/code hidden in ONNX metadata), extra files among zip entries (related to Pitfall 2), anomalous tensors with shapes/numerical distributions that are peculiar in the weights

Cross-domain correlation

[[re-managed]]: this gateway's "recognize runtime" fixedly invokes this skill after identifying an AI model file (a model is a managed-domain branch of "code inside data")
[[re-binary-core]]: code embedded in a model, a model packaged into an executable (an inference program packaged by PyInstaller/pyarmor) — first disassemble in the binary domain (format parsing/decompilation), and the extracted model file returns to this skill
[[re-sandbox]]: the default isolated environment for loading any unknown pkl ([[re-analyze/platform-tips]] highest principle)
[[re-malware]]: the behavior and intelligence side of malicious model payloads (pickle malicious code, backdoor weights, poisoned model distribution)
[[re-ioc]]: model fingerprints (sha256/tensor hash/watermark pattern) go into IOC; [[re-triage]]: entry point for initial survey of model files and hash archiving
[[re-script-deob]]: PyTorch inference/training script reconstruction (the completion path when state_dict has no structure)
Reference [[re-analyze/platform-tips]] static-first (large samples) and the sandbox highest-principle branch
[[re-ai-attack]]：File-level evidence → behavior-level verification(one-way data flow) — the file-side conclusions produced by this skill (structure/weights/file-level fingerprint) serve as input for re-ai-attack behavior consistency verification. Example: obtain suspect.pt → this skill extracts structure + weights to generate a fingerprint → hand off to re-ai-attack for behavior consistency verification (API-side comparison)
Evaluations with only an API (no file) go directly to [[re-ai-attack]], bypassing this skill

Common pitfalls and traps

Large model files are huge (GB-scale): Symptom—onnx.load/torch.load Memory fills up and it hangs, Netron times out when opening it,np.save and batch writes saturate the disk; cause—the weights are several GB, loaded entirely into memory all at once; countermeasure—before analysis, first du -sh/sha256sum archive it; for ONNX, use onnx.load(..., load_external_data=False) to load only the graph structure, thentensor by tensor retrieve values—note that numpy_helper.to_array() for this tensor reads it out completely(not streaming), so the peak depends on the largest tensor, and when using external data you mustexplicitly provide base_dir(the default is an empty string, and it fails directly when the weights are not in the CWD); for Safetensors, use safe_open for lazy per-tensor reads; for large PyTorch models torch.load(..., mmap=True); the object being processed is a "structure summary + targeted tensors", not the entire file
pkl deserialization risk (do not directly torch.load an unknown pkl): Symptom—after loading, the process bounces a shell/deletes files, or reports a bizarre AttributeError/ModuleNotFoundError; cause—the pickle protocol can inject arbitrary code (__reduce__/__setstate__), torch.load is fundamentally pickle, and a malicious model is a carrier for a poison payload; countermeasure—safety tip: never directly torch.loadan unknown model; first unzip -l/xxd/strings do a rough check (zip header PK vs. bare pickle, whether entries have suspicious module names), load with weights_only=True(the default in PyTorch 2.6+), and when full-featured loading is needed, execute it in an isolated environment ([[re-sandbox]]); samples that can be converted to safetensors should be converted directly (pure data, no code execution)
Graph-optimization obfuscation layer structure: Symptom—the operator sequence of a model after onnx-simplifier/TensorRT/onnxruntime optimization does not match the training state (Conv+BN fused into a single Conv, constant folding, all names changed), causing misjudgment of structural similarity; cause—the optimizer performs operator fusion/constant folding, the graph structure differs from the training state, and the producer field changes; countermeasure—first read m.producer_name/m.producer_version to identify the optimizer and version (the signature of fused ConvBN: the BN layer disappears and the scale is merged into the conv weights); before comparing architectures, first normalize the operator sequence (abstract by operator category, ignoring names and constant differences); you can use onnx-simplifier/onnxruntime graph_optimization_level to compare the diff before and after optimization and restore the original layers
Watermark robustness (still present after pruning): Symptom—concluding "no watermark" when an exact-value comparison does not hit, or misjudging "theft" when two unrelated models happen to have similar values in individual layers; cause—a robust watermark still survives pruning/quantization/retraining (by design), so exact matching is bound to miss; normal models with the same architecture and data have similar weight distributions, and a single-layer coincidence is a false positive; countermeasure—watermark detection uses "statistical anomalies" (low-bit perturbation/numerically distinctive layers) rather than "exact equality", and theft determination uses multi-dimensional evidence + thresholds (the discipline of step 5), with conclusions annotated with confidence and evidence strength