"""Script to export Deformable-DETR pretrained ONNX model(s). Deformable DETR (Deformable Transformers for End-to-End Object Detection) from SenseTime / fundamentalvision. Reference: https://github.com/fundamentalvision/Deformable-DETR Paper: https://arxiv.org/abs/2010.04159 Detection variants (Apache 2.0, COCO pretrained): deformable_detr_single_scale – 800×800, 34M params, AP50:95 39.4 deformable_detr_single_scale_dc5 – 800×800, 34M params, AP50:95 41.5 deformable_detr – 800×800, 40M params, AP50:95 44.5 deformable_detr_plus_iterative_bbox_refinement – 800×800, 41M params, AP50:95 46.2 deformable_detr_two_stage – 800×800, 41M params, AP50:95 46.9 Export methods (--method): torch (default) – clones the official GitHub repo, downloads weights from Google Drive via gdown. Requires internet access to drive.google.com (may be blocked on corporate proxies). optimum – downloads from HuggingFace Hub via the transformers library. Proxy-friendly, no CUDA compilation, no Google Drive. Uses HuggingFace model IDs under SenseTime/. Notes: - All variants use a ResNet-50 backbone, pre-trained on ImageNet. - DC5 variant is disabled: TIDL does not support dilated convolution in ResNet. Usage: python prepare_model.py python prepare_model.py --method optimum python prepare_model.py --model deformable_detr_single_scale python prepare_model.py --model deformable_detr_single_scale --method optimum python prepare_model.py --model deformable_detr deformable_detr_two_stage python prepare_model.py --model deformable_detr --shape 640 640 python prepare_model.py --model deformable_detr --weights /path/to/checkpoint.pth python prepare_model.py --model deformable_detr --opset 18 --output-dir ./exports python prepare_model.py --model deformable_detr --skip-simplify python prepare_model.py --model all python prepare_model.py --list-models """ from __future__ import annotations import argparse import importlib import math import os import subprocess import sys # ───────────────────────────────────────────── # Model catalogue # ───────────────────────────────────────────── # Each entry: variant_key → metadata dict # num_feature_levels : 1 (single scale) or 4 (multi-scale) # with_box_refine : iterative bounding box refinement # two_stage : two-stage proposal + detection # dilation : DC5 – dilation in ResNet's last block # gdrive_id : Google Drive file ID (used by --method torch) # hf_model_id : HuggingFace model ID (used by --method optimum) MODEL_CATALOG: dict[str, dict] = { "deformable_detr_single_scale": { "num_feature_levels": 1, "with_box_refine": False, "two_stage": False, "dilation": False, "shape": (800, 800), "params_m": 34, "ap50_95": 39.4, "flops_g": 78, "fps_v100": 27.0, "license": "Apache 2.0", "gdrive_id": "1WEjQ9_FgfI5sw5OZZ4ix-OKk-IJ_-SDU", "hf_model_id": "SenseTime/deformable-detr-single-scale", }, "deformable_detr_single_scale_dc5": { "num_feature_levels": 1, "with_box_refine": False, "two_stage": False, "dilation": True, "shape": (800, 800), "params_m": 34, "ap50_95": 41.5, "flops_g": 128, "fps_v100": 22.1, "license": "Apache 2.0", "gdrive_id": "1m_TgMjzH7D44fbA-c_jiBZ-xf-odxGdk", "hf_model_id": "SenseTime/deformable-detr-single-scale-dc5", }, "deformable_detr": { "num_feature_levels": 4, "with_box_refine": False, "two_stage": False, "dilation": False, "shape": (800, 800), "params_m": 40, "ap50_95": 44.5, "flops_g": 173, "fps_v100": 15.0, "license": "Apache 2.0", "gdrive_id": "1nDWZWHuRwtwGden77NLM9JoWe-YisJnA", "hf_model_id": "SenseTime/deformable-detr", }, "deformable_detr_plus_iterative_bbox_refinement": { "num_feature_levels": 4, "with_box_refine": True, "two_stage": False, "dilation": False, "shape": (800, 800), "params_m": 41, "ap50_95": 46.2, "flops_g": 173, "fps_v100": 15.0, "license": "Apache 2.0", "gdrive_id": "1JYKyRYzUH7uo9eVfDaVCiaIGZb5YTCuI", "hf_model_id": "SenseTime/deformable-detr-with-box-refine", }, "deformable_detr_two_stage": { "num_feature_levels": 4, "with_box_refine": True, "two_stage": True, "dilation": False, "shape": (800, 800), "params_m": 41, "ap50_95": 46.9, "flops_g": 173, "fps_v100": 14.5, "license": "Apache 2.0", "gdrive_id": "15I03A7hNTpwuLNdfuEmW9_taZMNVssEp", "hf_model_id": "SenseTime/deformable-detr-with-box-refine-two-stage", }, } DEFAULT_MODEL = "deformable_detr" _REPO_URL = "https://github.com/fundamentalvision/Deformable-DETR.git" _REPO_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "deformable_detr") # ───────────────────────────────────────────── # Dependency installer # ───────────────────────────────────────────── def _pip_install(*packages: str) -> None: """Install *packages* via pip, suppressing verbose output.""" print(f"[DEP] Installing: {', '.join(packages)} …") result = subprocess.run( [sys.executable, "-m", "pip", "install", *packages], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, ) if result.returncode != 0: print(f"[DEP] ERROR: pip install failed (exit code {result.returncode}).") if result.stderr: print(result.stderr.strip()) print("[DEP] Please install manually and re-run:") print(f" pip install {' '.join(packages)}") sys.exit(1) print("[DEP] Installation complete.\n") def ensure_dependencies() -> None: """Ensure torch, torchvision, onnx, scipy, gdown, and onnxsim are importable.""" required = [ ("torch", "torch>=1.12.0"), ("torchvision", "torchvision>=0.13.0"), ("onnx", "onnx>=1.14.0"), ("onnxsim", "onnx-simplifier"), ("scipy", "scipy"), ("gdown", "gdown>=5.2.0"), ] missing = [] for mod_name, pip_spec in required: try: importlib.import_module(mod_name) print(f"[DEP] ✔ {mod_name} is installed.") except ImportError: print(f"[DEP] ✘ {mod_name} not found.") missing.append(pip_spec) if missing: _pip_install(*missing) print() def ensure_hf_dependencies() -> None: """Ensure torch, torchvision, onnx, onnxsim, and transformers are importable (needed for --method optimum).""" required = [ ("torch", "torch>=1.12.0"), ("torchvision", "torchvision>=0.13.0"), ("onnx", "onnx>=1.14.0"), ("onnxsim", "onnx-simplifier"), ("transformers", "transformers>=4.30.0"), ] missing = [] for mod_name, pip_spec in required: try: importlib.import_module(mod_name) print(f"[DEP] ✔ {mod_name} is installed.") except ImportError: print(f"[DEP] ✘ {mod_name} not found.") missing.append(pip_spec) if missing: _pip_install(*missing) print() # ───────────────────────────────────────────── # Repo setup # ───────────────────────────────────────────── def setup_repo(force_reclone: bool = False) -> str: """Clone (or reuse) the Deformable-DETR repository. Returns the absolute path to the repository root. """ if os.path.isdir(_REPO_CACHE_DIR) and not force_reclone: print(f"[REPO] Using cached repo: {_REPO_CACHE_DIR}") return _REPO_CACHE_DIR if os.path.isdir(_REPO_CACHE_DIR): import shutil shutil.rmtree(_REPO_CACHE_DIR) os.makedirs(os.path.dirname(_REPO_CACHE_DIR), exist_ok=True) print(f"[REPO] Cloning Deformable-DETR into {_REPO_CACHE_DIR} …") result = subprocess.run( ["git", "clone", "--depth", "1", _REPO_URL, _REPO_CACHE_DIR], capture_output=True, text=True, ) if result.returncode != 0: print(f"[REPO] ERROR: git clone failed.\n{result.stderr.strip()}") sys.exit(1) print("[REPO] Clone complete.\n") return _REPO_CACHE_DIR # ───────────────────────────────────────────── # Torchvision compatibility shim # ───────────────────────────────────────────── def _patch_torchvision_compat() -> None: """Stub out removed torchvision symbols referenced by Deformable-DETR's util/misc.py. The repo does ``float(torchvision.__version__[:3]) < 0.5`` to gate old code. For torchvision >= 0.10 the string ``"0.15"[:3]`` is ``"0.1"`` → float 0.1 → the condition is True, triggering an import of ``_NewEmptyTensorOp`` that was removed in torchvision 0.9. We add a harmless stub so the import succeeds. The function util/misc.interpolate() falls back to ``torch.nn.functional.interpolate`` for non-empty tensors (all practical cases), so the stub is never called. """ import torch # noqa: PLC0415 import torchvision.ops.misc as _tvm # noqa: PLC0415 if hasattr(_tvm, "_NewEmptyTensorOp"): return # already present (old torchvision) — nothing to do class _NewEmptyTensorOp(torch.autograd.Function): @staticmethod def forward(ctx, x, new_size): return x.new_empty(new_size) @staticmethod def backward(ctx, grad): return grad, None _tvm._NewEmptyTensorOp = _NewEmptyTensorOp print("[COMPAT] Added _NewEmptyTensorOp stub to torchvision.ops.misc.\n") # ───────────────────────────────────────────── # Python fallback for deformable attention # ───────────────────────────────────────────── def _install_python_fallback(repo_root: str) -> None: """Set up a pure-Python replacement for the multi-scale deformable attention CUDA extension so that ONNX tracing works on CPU without requiring CUDA compilation. Strategy: 1. Patch torchvision.ops.misc to add the removed _NewEmptyTensorOp stub (required for util/misc.py to import on torchvision >= 0.10). 2. Register a placeholder 'MultiScaleDeformableAttention' module in sys.modules before any model code is imported (models/ops imports this at module level). 3. Import the pure-Python ms_deform_attn_core_pytorch function from the repo source. 4. Monkeypatch MSDeformAttn.forward to call ms_deform_attn_core_pytorch directly, bypassing the MSDeformAttnFunction custom autograd op (which has no ONNX symbolic and would break torch.onnx.export). """ import types import torch # noqa: PLC0415 if repo_root not in sys.path: sys.path.insert(0, repo_root) # Step 1 – Fix torchvision compatibility before importing any repo code. _patch_torchvision_compat() # Step 2 – Register a stub MSDA module so models/ops imports succeed. if "MultiScaleDeformableAttention" not in sys.modules: stub = types.ModuleType("MultiScaleDeformableAttention") sys.modules["MultiScaleDeformableAttention"] = stub print("[OPS] Registered stub MultiScaleDeformableAttention module.") # Step 3 – Import the Python-only reference implementation. from models.ops.functions.ms_deform_attn_func import ( # noqa: PLC0415 ms_deform_attn_core_pytorch, ) # Step 4 – Patch MSDeformAttn.forward to use ms_deform_attn_core_pytorch # directly instead of calling MSDeformAttnFunction.apply. # This is a verbatim rewrite of the original forward with only the final # output = MSDeformAttnFunction.apply(...) line replaced. import torch.nn.functional as F # noqa: PLC0415 import models.ops.modules.ms_deform_attn as _attn_mod # noqa: PLC0415 def _py_forward( self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None, ): N, Len_q, _ = query.shape N, Len_in, _ = input_flatten.shape assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in value = self.value_proj(input_flatten) if input_padding_mask is not None: value = value.masked_fill(input_padding_mask[..., None], float(0)) value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads) sampling_offsets = self.sampling_offsets(query).view( N, Len_q, self.n_heads, self.n_levels, self.n_points, 2 ) attention_weights = self.attention_weights(query).view( N, Len_q, self.n_heads, self.n_levels * self.n_points ) attention_weights = F.softmax(attention_weights, -1).view( N, Len_q, self.n_heads, self.n_levels, self.n_points ) if reference_points.shape[-1] == 2: offset_normalizer = torch.stack( [input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1 ) sampling_locations = ( reference_points[:, :, None, :, None, :] + sampling_offsets / offset_normalizer[None, None, None, :, None, :] ) elif reference_points.shape[-1] == 4: sampling_locations = ( reference_points[:, :, None, :, None, :2] + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 ) else: raise ValueError( f"Last dim of reference_points must be 2 or 4, " f"got {reference_points.shape[-1]}" ) output = ms_deform_attn_core_pytorch( value, input_spatial_shapes, sampling_locations, attention_weights ) output = self.output_proj(output) return output _attn_mod.MSDeformAttn.forward = _py_forward print("[OPS] Pure-Python fallback installed for MSDeformAttn (no CUDA required).\n") # ───────────────────────────────────────────── # ONNX simplification # ───────────────────────────────────────────── def simplify_onnx(src_path: str, force: bool = False) -> bool: """Run onnx-simplifier on *src_path* in-place. Simplification folds constants, removes dead nodes, and cleans up redundant ops produced by torch.onnx.export, making the graph smaller and easier to deploy. Falls back gracefully (copies as-is) if onnxsim is not installed. Args: src_path: Path to the .onnx file to simplify (modified in-place). force : Re-run even if the file was already simplified. Returns: True on success (or if simplification was skipped gracefully). """ print(f"\n[SIM] Running onnxsim on {os.path.basename(src_path)} …") try: import onnx # noqa: PLC0415 import onnxsim # noqa: PLC0415 except ImportError as exc: missing = str(exc).split("'")[1] if "'" in str(exc) else str(exc) print(f" [WARN] {missing} not installed – skipping simplification.") print(" Install with: pip install onnx-simplifier") return True try: model = onnx.load(src_path) except Exception as exc: print(f" [ERROR] Failed to load {src_path}: {exc}") return False try: model_sim, check = onnxsim.simplify(model) except Exception as exc: print(f" [WARN] onnxsim failed: {exc} – keeping unsimplified model.") return True if not check: print(" [WARN] onnxsim validation failed – keeping unsimplified model.") return True orig_nodes = len(model.graph.node) sim_nodes = len(model_sim.graph.node) delta = orig_nodes - sim_nodes print(f" Nodes: {orig_nodes} → {sim_nodes} (−{delta})") try: onnx.save(model_sim, src_path) except Exception as exc: print(f" [ERROR] Failed to save simplified model: {exc}") return False size_mb = os.path.getsize(src_path) / (1024 * 1024) print(f"[OK] Simplified model saved: {src_path} ({size_mb:.1f} MB)") return True # ───────────────────────────────────────────── # Weight download # ───────────────────────────────────────────── def download_weights(gdrive_id: str, weights_path: str, force: bool = False) -> bool: """Download a pretrained checkpoint from Google Drive using gdown. Args: gdrive_id : Google Drive file ID. weights_path: Local destination path for the .pth checkpoint. force : Re-download even if file already exists. Returns: True on success. """ import gdown # noqa: PLC0415 if os.path.exists(weights_path) and not force: size_mb = os.path.getsize(weights_path) / 1024 / 1024 print(f"[SKIP] Weights already exist ({size_mb:.1f} MB). " "Use --force to re-download.\n") return True url = f"https://drive.google.com/uc?id={gdrive_id}" print(f"[DOWN] Downloading pretrained weights from Google Drive …") print(f" File ID : {gdrive_id}") print(f" Dest : {weights_path}") # Pick up proxy settings from the environment (e.g. TI corporate proxy). proxy = ( os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy") or os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy") ) if proxy: print(f" Proxy : {proxy}") try: dl_kwargs: dict = {"quiet": False} if proxy: dl_kwargs["proxy"] = proxy gdown.download(url, weights_path, **dl_kwargs) except Exception as exc: print(f"[ERROR] gdown download failed: {exc}") _print_manual_download_hint(gdrive_id, weights_path) return False if not os.path.exists(weights_path): print("[ERROR] Download finished but file was not created.") _print_manual_download_hint(gdrive_id, weights_path) return False size_mb = os.path.getsize(weights_path) / 1024 / 1024 print(f"[OK] Weights saved: {weights_path} ({size_mb:.1f} MB)\n") return True def _print_manual_download_hint(gdrive_id: str, weights_path: str) -> None: """Print instructions for manually downloading a Google Drive checkpoint.""" url = f"https://drive.google.com/uc?id={gdrive_id}" print( f"\n[HINT] If you are behind a corporate proxy, download the checkpoint\n" f" manually using one of the following commands:\n" f"\n" f" # with gdown and explicit proxy:\n" f" gdown --proxy '{url}' -O '{weights_path}'\n" f"\n" f" # with curl:\n" f" curl -L -x '{url}' -o '{weights_path}'\n" f"\n" f" Then re-run with --weights to skip the download:\n" f" python prepare_model.py --model --weights '{weights_path}'\n" ) # ───────────────────────────────────────────── # Model catalogue helpers # ───────────────────────────────────────────── def print_model_table() -> None: """Print a formatted table of all available model variants.""" col = 48 header = ( f" {'Variant':<{col}} {'Shape':<10} {'Params(M)':<10} " f"{'AP50:95':<8} {'FLOPs(G)':<9} {'FPS(V100)':<10} {'License'}" ) sep = " " + "-" * (len(header) - 2) print("\n" + "=" * len(header)) print(" Available Deformable-DETR model variants") print("=" * len(header)) print(header) print(sep) for key, info in MODEL_CATALOG.items(): h, w = info["shape"] print( f" {key:<{col}} {h}×{w:<5} " f"{info['params_m']:<10} {info['ap50_95']:<8.1f} " f"{info['flops_g']:<9} {info['fps_v100']:<10.1f} " f"{info['license']}" ) print("=" * len(header) + "\n") print(" All variants use ResNet-50 backbone, trained on COCO 2017.") print(" AP50:95 measured on COCO val2017, inference speed on V100 GPU.\n") # ───────────────────────────────────────────── # Float64 removal # ───────────────────────────────────────────── def fix_float64_nodes(src_path: str) -> bool: """Remove Cast-to-DOUBLE nodes and fix float64 initializers/constants so the model is compatible with TIDL (which does not support float64). torch.onnx.export inserts Cast(to=DOUBLE) nodes when Python-level float literals (e.g. math.pi, which is float64) appear in position-encoding computations. These nodes propagate float64 through most of the graph. Strategy: 1. Find every Cast node with to=DOUBLE. 2. Re-wire each consumer of the Cast's output to use the Cast's input (the upstream float32 tensor) directly, then delete the Cast node. 3. Convert any float64 graph initializers to float32. 4. Fix any Constant/ConstantOfShape attribute tensors that are DOUBLE. 5. Validate with onnx.checker and save in-place. Args: src_path: Path to the .onnx file to fix (modified in-place). Returns: True on success or if no float64 tensors were found. """ import numpy as np # noqa: PLC0415 try: import onnx # noqa: PLC0415 from onnx import TensorProto, numpy_helper # noqa: PLC0415 except ImportError: print(" [WARN] onnx not installed – skipping float64 fix.") return True try: model = onnx.load(src_path) except Exception as exc: print(f" [ERROR] Failed to load {src_path}: {exc}") return False graph = model.graph # Step 1 – Remove Cast-to-DOUBLE by re-wiring consumers to use Cast input. consumers: dict = {} for node in graph.node: for inp in node.input: consumers.setdefault(inp, []).append(node) removed = 0 for node in list(graph.node): if node.op_type != "Cast": continue for attr in node.attribute: if attr.name == "to" and attr.i == TensorProto.DOUBLE: cast_in = node.input[0] cast_out = node.output[0] for consumer in consumers.get(cast_out, []): consumer.input[:] = [ cast_in if t == cast_out else t for t in consumer.input ] graph.node.remove(node) removed += 1 break # Step 2 – Convert float64 graph initializers to float32. init_fixed = 0 for init in graph.initializer: if init.data_type == TensorProto.DOUBLE: arr = numpy_helper.to_array(init).astype(np.float32) init.CopyFrom(numpy_helper.from_array(arr, name=init.name)) init_fixed += 1 # Step 3 – Fix Constant/ConstantOfShape nodes with float64 value tensors. # Use attr.type == TENSOR (the correct API) instead of attr.HasField("t"), # which is unreliable across protobuf versions. const_fixed = 0 for node in graph.node: for attr in node.attribute: if (attr.type == onnx.AttributeProto.TENSOR and attr.t.data_type == TensorProto.DOUBLE): arr = numpy_helper.to_array(attr.t).astype(np.float32) attr.t.CopyFrom(numpy_helper.from_array(arr)) const_fixed += 1 # Step 4 – Update stale float64 type annotations in value_info. # onnxsim stores intermediate tensor types in graph.value_info. When Cast- # to-DOUBLE nodes are removed the stored annotations become stale and still # say float64, which causes type-inference errors in TIDL and ONNX tools # even though the actual computation is now float32. vi_fixed = 0 for vi in list(graph.value_info) + list(graph.input) + list(graph.output): if (vi.type.HasField("tensor_type") and vi.type.tensor_type.elem_type == TensorProto.DOUBLE): vi.type.tensor_type.elem_type = TensorProto.FLOAT vi_fixed += 1 print( f"[F64] Cast-to-DOUBLE removed: {removed}, " f"initializers fixed: {init_fixed}, constants fixed: {const_fixed}, " f"type annotations fixed: {vi_fixed}" ) if removed == 0 and init_fixed == 0 and const_fixed == 0 and vi_fixed == 0: print("[F64] No float64 tensors found – model already clean.") return True try: onnx.checker.check_model(model) print("[F64] ONNX model validation passed after float64 fix.") except Exception as exc: print(f"[WARN] ONNX validation after float64 fix: {exc}") try: onnx.save(model, src_path) except Exception as exc: print(f" [ERROR] Failed to save fixed model: {exc}") return False size_mb = os.path.getsize(src_path) / (1024 * 1024) print(f"[OK] Float64-free model saved: {src_path} ({size_mb:.1f} MB)") return True # ───────────────────────────────────────────── # Core export # ───────────────────────────────────────────── def export_model( model_key: str, output_dir: str, shape: tuple[int, int] | None, opset: int, batch_size: int, verbose: bool, custom_weights: str | None, force: bool, force_reclone: bool, skip_simplify: bool = False, ) -> str: """Clone the Deformable-DETR repo, download weights, and export to ONNX. Args: model_key : Key from MODEL_CATALOG. output_dir : Directory to save the .onnx and .pth files. shape : Custom (height, width) or None for model default. opset : ONNX opset version. batch_size : Batch size embedded in the exported graph. verbose : Show additional progress messages. custom_weights: Path to a local .pth checkpoint; None = pretrained. force : Re-export even if .onnx already exists. force_reclone : Force re-clone of the source repo. Returns: Absolute path of the saved .onnx file. """ import torch # noqa: PLC0415 info = MODEL_CATALOG[model_key] export_shape = shape if shape is not None else info["shape"] h, w = export_shape os.makedirs(output_dir, exist_ok=True) shape_tag = f"_{h}x{w}" if shape is not None else "" dst_name = f"{model_key}{shape_tag}.onnx" dst_path = os.path.join(output_dir, dst_name) if not force and os.path.exists(dst_path): print(f"[SKIP] {dst_name} already exists. Use --force to re-export.\n") return dst_path print(f"[INFO] Variant : {model_key}") print(f"[INFO] Feature lvls : {info['num_feature_levels']}") print(f"[INFO] Box refine : {info['with_box_refine']}") print(f"[INFO] Two-stage : {info['two_stage']}") print(f"[INFO] DC5 dilation : {info['dilation']}") print(f"[INFO] Input shape : {h}×{w} (batch {batch_size})") print(f"[INFO] ONNX opset : {opset}") if custom_weights: print(f"[INFO] Weights : {custom_weights}") else: print(f"[INFO] Weights : COCO pretrained (Google Drive)") # ── Step 1: Clone repo ──────────────────────────────────────────────────── repo_root = setup_repo(force_reclone=force_reclone) # ── Step 2: Install Python fallback for deformable attention ────────────── _install_python_fallback(repo_root) # ── Step 3: Download or locate weights ──────────────────────────────────── if custom_weights: weights_path = custom_weights if not os.path.exists(weights_path): print(f"[ERROR] Custom weights not found: {weights_path}") sys.exit(1) else: weights_path = os.path.join(output_dir, f"{model_key}.pth") if not download_weights(info["gdrive_id"], weights_path, force=force): print(f"[ERROR] Failed to download weights for '{model_key}'.") print( "\n[TIP] The default export method (torch) downloads weights from\n" " Google Drive, which may be unreachable on corporate networks.\n" " Try the HuggingFace-based method instead — no Google Drive\n" " required, proxy-friendly:\n" f"\n" f" python prepare_model.py --method optimum --model {model_key}\n" ) sys.exit(1) # ── Step 4: Build model ─────────────────────────────────────────────────── print("[INFO] Building model …") if repo_root not in sys.path: sys.path.insert(0, repo_root) from models import build_model # noqa: PLC0415 args = argparse.Namespace( # Backbone backbone = "resnet50", dilation = info["dilation"], position_embedding = "sine", position_embedding_scale= 2 * math.pi, num_feature_levels = info["num_feature_levels"], # Transformer enc_layers = 6, dec_layers = 6, dim_feedforward = 1024, hidden_dim = 256, dropout = 0.1, nheads = 8, num_queries = 300, dec_n_points = 4, enc_n_points = 4, # Variant flags with_box_refine = info["with_box_refine"], two_stage = info["two_stage"], # Segmentation (not used for detection export) masks = False, frozen_weights = None, # Loss (needed by SetCriterion constructor, not used for inference) aux_loss = False, set_cost_class = 2.0, set_cost_bbox = 5.0, set_cost_giou = 2.0, mask_loss_coef = 1.0, dice_loss_coef = 1.0, cls_loss_coef = 2.0, bbox_loss_coef = 5.0, giou_loss_coef = 2.0, focal_alpha = 0.25, # Dataset (determines num_classes = 91 for coco) dataset_file = "coco", coco_path = "./data/coco", coco_panoptic_path = None, remove_difficult = False, # Device device = "cpu", ) model, _criterion, _postprocessors = build_model(args) model.eval() print("[INFO] Model built.\n") # ── Step 5: Load pretrained weights ─────────────────────────────────────── print(f"[INFO] Loading weights from: {weights_path}") checkpoint = torch.load(weights_path, map_location="cpu") state_dict = checkpoint.get("model", checkpoint) missing, unexpected = model.load_state_dict(state_dict, strict=False) unexpected = [k for k in unexpected if not k.endswith(("total_params", "total_ops"))] if missing: print(f"[WARN] Missing keys : {missing[:5]}{'…' if len(missing) > 5 else ''}") if unexpected: print(f"[WARN] Unexpected keys: {unexpected[:5]}{'…' if len(unexpected) > 5 else ''}") print("[INFO] Weights loaded.\n") # ── Step 6: Build ONNX wrapper ──────────────────────────────────────────── import torch # noqa: PLC0415 import torch.nn as nn # noqa: PLC0415 from util.misc import NestedTensor # noqa: PLC0415 class _Wrapper(nn.Module): def __init__(self): super().__init__() self.model = model self._NT = NestedTensor def forward(self, images: torch.Tensor): B, _, H, W = images.shape mask = torch.zeros((B, H, W), dtype=torch.bool, device=images.device) out = self.model(self._NT(images, mask)) return out["pred_boxes"], out["pred_logits"] wrapper = _Wrapper().eval() # ── Step 7: Export to ONNX ──────────────────────────────────────────────── print(f"[INFO] Exporting to ONNX (opset {opset}) …") dummy = torch.zeros(batch_size, 3, h, w) with torch.no_grad(): torch.onnx.export( wrapper, (dummy,), dst_path, input_names = ["images"], output_names = ["pred_boxes", "pred_logits"], opset_version = opset, do_constant_folding = True, ) # ── Optional ONNX validation ────────────────────────────────────────────── try: import onnx # noqa: PLC0415 onnx_model = onnx.load(dst_path) onnx.checker.check_model(onnx_model) print("[INFO] ONNX model validation passed.") except ImportError: pass except Exception as exc: print(f"[WARN] ONNX validation: {exc}") # ── Optional onnxsim simplification ────────────────────────────────────── if not skip_simplify: simplify_onnx(dst_path, force=force) # ── Fix float64 nodes for TIDL compatibility ────────────────────────────── fix_float64_nodes(dst_path) size_mb = os.path.getsize(dst_path) / (1024 * 1024) print(f"\n[SUCCESS] ONNX model saved to: {dst_path} ({size_mb:.1f} MB)\n") return dst_path # ───────────────────────────────────────────── # Optimum / HuggingFace export # ───────────────────────────────────────────── def export_model_optimum( model_key: str, output_dir: str, shape: tuple[int, int] | None, opset: int, batch_size: int, force: bool, skip_simplify: bool = False, ) -> str: """Download from HuggingFace and export Deformable-DETR to ONNX. Uses the HuggingFace transformers implementation of Deformable DETR, which is a pure-Python port of the original architecture. No Google Drive access, no CUDA compilation, and no repo cloning required. The transformers model is downloaded via HuggingFace Hub. Proxy settings are picked up automatically from the HTTPS_PROXY / https_proxy environment variables (TI corporate proxy is supported). Args: model_key : Key from MODEL_CATALOG. output_dir : Directory to save the .onnx file. shape : Custom (height, width) or None for model default. opset : ONNX opset version. batch_size : Batch size embedded in the exported graph. force : Re-export even if .onnx already exists. Returns: Absolute path of the saved .onnx file. """ import torch # noqa: PLC0415 import torch.nn as nn # noqa: PLC0415 from transformers import DeformableDetrForObjectDetection # noqa: PLC0415 info = MODEL_CATALOG[model_key] hf_model_id = info["hf_model_id"] export_shape = shape if shape is not None else info["shape"] h, w = export_shape os.makedirs(output_dir, exist_ok=True) shape_tag = f"_{h}x{w}" if shape is not None else "" dst_name = f"{model_key}{shape_tag}.onnx" dst_path = os.path.join(output_dir, dst_name) if not force and os.path.exists(dst_path): print(f"[SKIP] {dst_name} already exists. Use --force to re-export.\n") return dst_path print(f"[INFO] Method : optimum (HuggingFace transformers)") print(f"[INFO] HF model ID : {hf_model_id}") print(f"[INFO] Input shape : {h}×{w} (batch {batch_size})") print(f"[INFO] ONNX opset : {opset}") print() # ── Download / load from HuggingFace ───────────────────────────────────── print(f"[INFO] Loading model from HuggingFace …") print("[INFO] (First run downloads ~150–200 MB; cached at ~/.cache/huggingface/)") model = DeformableDetrForObjectDetection.from_pretrained(hf_model_id) model.eval() print("[INFO] Model ready.\n") # ── ONNX export wrapper ─────────────────────────────────────────────────── # DeformableDetrForObjectDetection.forward(pixel_values, pixel_mask=None) # outputs: DeformableDetrObjectDetectionOutput with .pred_boxes and .logits # We rename logits → pred_logits to match our postprocess configs. class _HFWrapper(nn.Module): def __init__(self): super().__init__() self.model = model def forward(self, images: torch.Tensor): out = self.model(pixel_values=images) return out.pred_boxes, out.logits wrapper = _HFWrapper().eval() dummy = torch.zeros(batch_size, 3, h, w) # ── Export ──────────────────────────────────────────────────────────────── print(f"[INFO] Exporting to ONNX (opset {opset}) …") with torch.no_grad(): torch.onnx.export( wrapper, (dummy,), dst_path, input_names = ["images"], output_names = ["pred_boxes", "pred_logits"], opset_version = opset, do_constant_folding= True, ) # ── Optional validation ─────────────────────────────────────────────────── try: import onnx # noqa: PLC0415 onnx_model = onnx.load(dst_path) onnx.checker.check_model(onnx_model) print("[INFO] ONNX model validation passed.") except ImportError: pass except Exception as exc: print(f"[WARN] ONNX validation: {exc}") # ── Optional onnxsim simplification ────────────────────────────────────── if not skip_simplify: simplify_onnx(dst_path, force=force) # ── Fix float64 nodes for TIDL compatibility ────────────────────────────── fix_float64_nodes(dst_path) size_mb = os.path.getsize(dst_path) / (1024 * 1024) print(f"\n[SUCCESS] ONNX model saved to: {dst_path} ({size_mb:.1f} MB)\n") return dst_path # ───────────────────────────────────────────── # CLI # ───────────────────────────────────────────── def build_parser() -> argparse.ArgumentParser: default_output = os.path.dirname(os.path.abspath(__file__)) parser = argparse.ArgumentParser( description=( "Export Deformable-DETR pretrained ONNX models.\n\n" "The Deformable-DETR source is cloned from GitHub on first use.\n" "Pretrained COCO weights are downloaded from Google Drive via gdown.\n" "Run --list-models to see all available variants." ), formatter_class=argparse.RawDescriptionHelpFormatter, epilog=( "Examples:\n" " %(prog)s\n" " %(prog)s --method optimum # proxy-friendly HF download\n" " %(prog)s --model deformable_detr_single_scale\n" " %(prog)s --model deformable_detr_single_scale --method optimum\n" " %(prog)s --model deformable_detr deformable_detr_two_stage\n" " %(prog)s --model deformable_detr --shape 640 640\n" " %(prog)s --model deformable_detr --weights /path/to/checkpoint.pth\n" " %(prog)s --model deformable_detr --opset 18 --output-dir ./exports\n" " %(prog)s --model deformable_detr --skip-simplify\n" " %(prog)s --model all\n" " %(prog)s --list-models" ), ) parser.add_argument( "--model", nargs="+", default=[DEFAULT_MODEL], choices=list(MODEL_CATALOG.keys()) + ["all"], metavar="VARIANT", help=( f"Model variant(s) to export. Default: {DEFAULT_MODEL}. " "Use 'all' to export every variant that has not yet been exported. " "Run --list-models to see all options." ), ) parser.add_argument( "--shape", nargs=2, type=int, default=None, metavar=("H", "W"), help=( "Custom input resolution (height width). " "Default: each model's native 800×800." ), ) parser.add_argument( "--opset", type=int, default=17, metavar="N", help="ONNX opset version. Default: 17.", ) parser.add_argument( "--batch-size", type=int, default=1, metavar="N", help="Batch size embedded in the exported ONNX graph. Default: 1.", ) parser.add_argument( "--weights", default=None, metavar="PATH", help=( "Path to a local .pth checkpoint (format: {'model': state_dict, ...}). " "When omitted the official COCO pretrained weights are downloaded " "automatically from Google Drive." ), ) parser.add_argument( "--output-dir", default=default_output, metavar="DIR", help=f"Directory where .onnx and .pth files will be saved. Default: {default_output}", ) parser.add_argument( "--force", action="store_true", default=False, help="Re-export and re-download even if output files already exist.", ) parser.add_argument( "--force-reclone", action="store_true", default=False, help=( "Force re-clone of the Deformable-DETR repository, " "removing the cached copy in ~/.cache/deformable_detr." ), ) parser.add_argument( "--skip-simplify", action="store_true", default=False, help=( "Skip the onnxsim simplification step. " "By default the exported ONNX is simplified in-place with " "onnx-simplifier (pip install onnx-simplifier). " "Use this flag to skip if onnxsim is unavailable or causing issues." ), ) parser.add_argument( "--quiet", action="store_true", default=False, help="Suppress verbose progress messages.", ) # ── Export method ───────────────────────────────────────────────────────── parser.add_argument( "--method", choices=["torch", "optimum"], default="torch", metavar="METHOD", help=( "Export method. " "'torch' (default): clones the official GitHub repo and downloads " "weights from Google Drive via gdown. " "'optimum': downloads from HuggingFace Hub using the transformers " "library — proxy-friendly, no CUDA ops, no Google Drive required." ), ) # ── Utility ─────────────────────────────────────────────────────────────── parser.add_argument( "--list-models", action="store_true", default=False, help="Print the model catalogue table and exit.", ) return parser # ───────────────────────────────────────────── # Entry point # ───────────────────────────────────────────── def main() -> None: parser = build_parser() args = parser.parse_args() if args.list_models: print_model_table() return if "all" in args.model: shape_tag = f"_{args.shape[0]}x{args.shape[1]}" if args.shape else "" output_dir = os.path.abspath(args.output_dir) pending = [ k for k in MODEL_CATALOG if not os.path.exists(os.path.join(output_dir, f"{k}{shape_tag}.onnx")) ] if not pending: print("[INFO] All models already exported. Use --force to re-export.") return skipped = [k for k in MODEL_CATALOG if k not in pending] if skipped: print("[INFO] Already exported (skipping):") for k in skipped: print(f" {k}") print("[INFO] Will export:") for k in pending: print(f" {k}") print() args.model = pending if args.weights and len(args.model) > 1: print( "[WARN] --weights applies the same checkpoint to every model in " "--model.\n This is unusual; pass a single --model variant " "when using custom weights." ) if args.weights and args.method == "optimum": print("[WARN] --weights is ignored with --method optimum. " "HuggingFace weights are always downloaded from the Hub.\n") # Install dependencies appropriate to the chosen method if args.method == "optimum": ensure_hf_dependencies() else: ensure_dependencies() shape = (args.shape[0], args.shape[1]) if args.shape else None output_dir = os.path.abspath(args.output_dir) exported: list[str] = [] failed: list[str] = [] for model_key in args.model: # if MODEL_CATALOG[model_key]["dilation"]: # print( # f"\n[WARN] '{model_key}' is a DC5 (dilation) variant and is " # "temporarily disabled because TIDL does not support dilated " # "convolution in ResNet. Skipping.\n" # ) # continue print(f"\n{'='*60}") print(f" Exporting: {model_key} [method={args.method}]") print(f"{'='*60}\n") try: if args.method == "optimum": out_path = export_model_optimum( model_key = model_key, output_dir = output_dir, shape = shape, opset = args.opset, batch_size = args.batch_size, force = args.force, skip_simplify = args.skip_simplify, ) else: out_path = export_model( model_key = model_key, output_dir = output_dir, shape = shape, opset = args.opset, batch_size = args.batch_size, verbose = not args.quiet, custom_weights= args.weights, force = args.force, force_reclone = args.force_reclone, skip_simplify = args.skip_simplify, ) exported.append(out_path) except SystemExit: raise except Exception as exc: print(f"[ERROR] Export failed for '{model_key}': {exc}") import traceback traceback.print_exc() failed.append(model_key) print("\n" + "=" * 60) print(" Export Summary") print("=" * 60) for path in exported: size_mb = os.path.getsize(path) / (1024 * 1024) print(f" ✔ {os.path.basename(path)} ({size_mb:.1f} MB)") print(f" {path}") if failed: for key in failed: print(f" ✘ {key} (FAILED)") print("=" * 60 + "\n") if failed and args.method == "torch": failed_str = " ".join(failed) print( "[TIP] The torch method failed (common causes: Google Drive blocked\n" " by a corporate proxy, or missing CUDA ops).\n" " Try the HuggingFace-based export instead — it downloads from\n" " HuggingFace Hub and requires no Google Drive access:\n" f"\n" f" python prepare_model.py --method optimum --model {failed_str}\n" ) elif failed and args.method == "optimum": failed_str = " ".join(failed) print( "[TIP] The optimum method failed.\n" " If HuggingFace Hub is accessible, check your transformers\n" " installation. You can also try the torch method with a\n" " manually downloaded checkpoint:\n" f"\n" f" python prepare_model.py --method torch --model {failed_str}\n" f" python prepare_model.py --method torch --model {failed_str} " f"--weights /path/to/checkpoint.pth\n" ) if failed: sys.exit(1) if __name__ == "__main__": main()