#!/usr/bin/env python3 """ Export DINOv2 classification models (backbone + linear head) from PyTorch Hub to ONNX with fixed input shapes for TI EdgeAI hardware deployment. Supported models (backbone + linear classification head, 1000-class ImageNet): dinov2_vits14_lc - ViT-S/14 distilled, 21M params, 81.1% top-1 dinov2_vitb14_lc - ViT-B/14 distilled, 86M params, 84.5% top-1 dinov2_vitl14_lc - ViT-L/14 distilled, 307M params, 86.3% top-1 dinov2_vitg14_lc - ViT-g/14, 1100M params, 86.5% top-1 dinov2_vits14_reg_lc - ViT-S/14 + registers, 21M params, 80.9% top-1 dinov2_vitb14_reg_lc - ViT-B/14 + registers, 86M params, 84.6% top-1 dinov2_vitl14_reg_lc - ViT-L/14 + registers, 307M params, 86.7% top-1 dinov2_vitg14_reg_lc - ViT-g/14 + registers, 1100M params, 87.1% top-1 Usage: python prepare_model.py --model dinov2_vits14_lc python prepare_model.py --model dinov2_vitb14_lc --no-simplifier python prepare_model.py --model dinov2_vitl14_lc --skip-export """ import sys import subprocess import tempfile from pathlib import Path SUPPORTED_MODELS = { 'dinov2_vits14_lc': {'params': '21M', 'accuracy_top1': 81.1, 'gflops': 4.6, 'edge_suitable': True}, 'dinov2_vitb14_lc': {'params': '86M', 'accuracy_top1': 84.5, 'gflops': 17.6, 'edge_suitable': True}, 'dinov2_vitl14_lc': {'params': '307M', 'accuracy_top1': 86.3, 'gflops': 61.6, 'edge_suitable': False}, 'dinov2_vitg14_lc': {'params': '1100M', 'accuracy_top1': 86.5, 'gflops': 314.0, 'edge_suitable': False}, 'dinov2_vits14_reg_lc': {'params': '21M', 'accuracy_top1': 80.9, 'gflops': 4.6, 'edge_suitable': True}, 'dinov2_vitb14_reg_lc': {'params': '86M', 'accuracy_top1': 84.6, 'gflops': 17.6, 'edge_suitable': True}, 'dinov2_vitl14_reg_lc': {'params': '307M', 'accuracy_top1': 86.7, 'gflops': 61.6, 'edge_suitable': False}, 'dinov2_vitg14_reg_lc': {'params': '1100M', 'accuracy_top1': 87.1, 'gflops': 314.0, 'edge_suitable': False}, } # Models exported with --model all (edge-suitable only; ViT-g excluded due to ~16GB RAM requirement) EDGE_MODELS = [name for name, info in SUPPORTED_MODELS.items() if info['edge_suitable']] def _ensure_dependencies(): required = { 'onnx': 'onnx', 'onnxsim': 'onnx-simplifier', 'torch': 'torch', } for module, package in required.items(): try: __import__(module) except ImportError: print(f"Installing missing dependency: {package}") subprocess.check_call([sys.executable, '-m', 'pip', 'install', package]) _ensure_dependencies() import torch import onnx from onnx import shape_inference import argparse def _consolidate_external_data(src_onnx_path, dst_onnx_path): """ Consolidate scattered per-tensor external data files into a single .data file by streaming — never loads all tensor weights into RAM at once. Used for models > 2 GB where onnx.load() would cause an OOM kill. Produces two files: dst_onnx_path (proto) and dst_onnx_path.name + '.data'. """ from onnx import TensorProto, AttributeProto src_dir = src_onnx_path.parent data_filename = dst_onnx_path.name + '.data' data_path = dst_onnx_path.parent / data_filename # Load proto structure only — tensor data stays on disk model = onnx.load(str(src_onnx_path), load_external_data=False) def _external_tensors(model_proto): for t in model_proto.graph.initializer: if t.data_location == TensorProto.EXTERNAL: yield t for node in model_proto.graph.node: for attr in node.attribute: if attr.type == AttributeProto.TENSOR and attr.t.data_location == TensorProto.EXTERNAL: yield attr.t elif attr.type == AttributeProto.TENSORS: for t in attr.tensors: if t.data_location == TensorProto.EXTERNAL: yield t offset = 0 with open(str(data_path), 'wb') as out_f: for tensor in _external_tensors(model): info = {e.key: e.value for e in tensor.external_data} src_file = src_dir / info['location'] t_offset = int(info.get('offset', 0)) t_length = int(info['length']) with open(str(src_file), 'rb') as src_f: src_f.seek(t_offset) out_f.write(src_f.read(t_length)) del tensor.external_data[:] for k, v in [('location', data_filename), ('offset', str(offset)), ('length', str(t_length))]: e = tensor.external_data.add() e.key = k e.value = v offset += t_length # Save only the proto — tensor data lives in data_path onnx.save(model, str(dst_onnx_path)) print(f"✓ Proto: {dst_onnx_path.name}") print(f"✓ Tensor data: {data_filename} ({data_path.stat().st_size / 1024**3:.2f} GB)") def export_to_onnx(model_name, output_path, height=224, width=224): """ Load DINOv2 classifier from PyTorch Hub and export to ONNX. torch.onnx.export writes per-tensor external data files alongside its output for models exceeding the 2 GB protobuf limit. The export runs inside a TemporaryDirectory so those files never appear in the destination directory. Small models (S/B/L, < 2 GB): external data is merged inline → single ONNX file. Large models (g, > 2 GB): _consolidate_external_data streams tensors into a single .data file without loading all weights into RAM. """ info = SUPPORTED_MODELS[model_name] print(f"\nLoading model from PyTorch Hub:") print("=" * 80) print(f"Model: {model_name}") print(f"Params: {info['params']}") print(f"Top-1: {info['accuracy_top1']}%") print(f"GFLOPs: {info['gflops']}") print() try: model = torch.hub.load('facebookresearch/dinov2', model_name, pretrained=True) except Exception as e: print(f"✗ Failed to load model: {e}") print(" Ensure you have an internet connection and PyTorch installed.") return False model.eval() # aten::_upsample_bicubic2d_aa (antialias=True) has no opset-17 ONNX handler. # Register models set interpolate_antialias=True on the backbone for positional # embedding interpolation. Patching it to False uses standard bicubic, which # is fully supported and has negligible quality difference at fixed 224x224. for module in model.modules(): if hasattr(module, 'interpolate_antialias'): module.interpolate_antialias = False dummy_input = torch.randn(1, 3, height, width) print(f"\nExporting to ONNX (opset 17):") print("-" * 80) print(f"Input shape: [1, 3, {height}, {width}]") print(f"Output path: {output_path}") try: with tempfile.TemporaryDirectory() as tmpdir: tmp_onnx = Path(tmpdir) / f"{model_name}.onnx" torch.onnx.export( model, dummy_input, str(tmp_onnx), export_params=True, opset_version=17, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes={ 'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}, }, ) # Probe proto (no data loaded) to detect external data format probe = onnx.load(str(tmp_onnx), load_external_data=False) has_external = any( t.data_location == onnx.TensorProto.EXTERNAL for t in probe.graph.initializer ) if has_external: # Model > 2 GB: stream-consolidate without loading all weights into RAM print("\nModel exceeds 2 GB — streaming consolidation (no full RAM load)...") _consolidate_external_data(tmp_onnx, output_path) else: # Model fits inline: load and save as single self-contained ONNX print("\nMerging external data into single ONNX file...") exported = onnx.load(str(tmp_onnx)) onnx.save(exported, str(output_path)) # tmpdir and all scattered tensor files are deleted here except Exception as e: print(f"✗ ONNX export failed: {e}") return False if output_path.exists(): file_size = output_path.stat().st_size print(f"\n✓ Export completed successfully!") print(f"✓ File size: {file_size:,} bytes ({file_size / 1024 / 1024:.2f} MB)") return True print("✗ Export failed: output file not created") return False def fix_model_shape(model_path, output_path, batch_size=1, channels=3, height=224, width=224, use_simplifier=True): """ Convert dynamic ONNX model input shape to fixed shape in all layers. Uses ONNX shape inference to propagate fixed shapes through all intermediate layers. Optionally runs onnxsim for additional simplification and optimization. """ print(f"\nFixing Model Shapes:") print("=" * 80) print(f"Input model: {model_path}") print(f"Output model: {output_path}") print(f"\nLoading model...") model = onnx.load(str(model_path)) graph = model.graph input_tensor = None for inp in graph.input: if any(init.name == inp.name for init in graph.initializer): continue input_tensor = inp break if input_tensor is None: print("✗ Error: No input tensor found!") return False print(f"\nOriginal Shape:") print("-" * 80) print(f"Input name: {input_tensor.name}") original_shape = [] for dim in input_tensor.type.tensor_type.shape.dim: if dim.dim_value: original_shape.append(str(dim.dim_value)) elif dim.dim_param: original_shape.append(f"'{dim.dim_param}'") else: original_shape.append("?") print(f"Shape: [{', '.join(original_shape)}]") print(f"\nSetting Fixed Shape:") print("-" * 80) new_shape = [batch_size, channels, height, width] print(f"New shape: {new_shape}") print(f"Format: [batch_size, channels, height, width]") input_tensor.type.tensor_type.shape.ClearField('dim') for dim_value in new_shape: dim = input_tensor.type.tensor_type.shape.dim.add() dim.dim_value = dim_value print(f"\nRunning Shape Inference:") print("-" * 80) try: model = shape_inference.infer_shapes(model) value_info_count = len(model.graph.value_info) print(f"✓ Propagated shapes through {value_info_count} intermediate tensors") except Exception as e: print(f"⚠ Warning: Shape inference issue: {e}") print(" Continuing with partial inference...") print(f"\nValidating Model:") print("-" * 80) try: onnx.checker.check_model(model) print("✓ Model validation passed") except Exception as e: print(f"✗ Model validation failed: {e}") return False if use_simplifier: print(f"\nRunning ONNX Simplifier:") print("-" * 80) try: import onnxsim model_simplified, check = onnxsim.simplify( model, check_n=3, perform_optimization=True, skip_fuse_bn=False, overwrite_input_shapes={input_tensor.name: new_shape}, ) if check: print("✓ Model simplified and optimized") model = model_simplified original_nodes = len(graph.node) simplified_nodes = len(model.graph.node) if simplified_nodes < original_nodes: print(f"✓ Reduced nodes: {original_nodes} → {simplified_nodes}") else: print("⚠ Simplification validation failed, using non-simplified version") except ImportError: print("⚠ onnx-simplifier not installed, skipping") except Exception as e: print(f"⚠ Simplification failed: {e}") print(" Continuing with non-simplified model") print(f"\nSaving Fixed Model:") print("-" * 80) onnx.save(model, str(output_path)) output_size = output_path.stat().st_size print(f"✓ Saved to: {output_path}") print(f"✓ File size: {output_size:,} bytes ({output_size / 1024 / 1024:.2f} MB)") print(f"\nFinal Verification:") print("-" * 80) try: verified_model = onnx.load(str(output_path)) onnx.checker.check_model(verified_model) verified_graph = verified_model.graph for inp in verified_graph.input: if any(init.name == inp.name for init in verified_graph.initializer): continue shape = [dim.dim_value for dim in inp.type.tensor_type.shape.dim] all_fixed = all(isinstance(s, int) and s > 0 for s in shape) if all_fixed: print(f"✓ Input '{inp.name}': {shape}") else: print(f"⚠ Input '{inp.name}' has dynamic dimensions") if verified_graph.value_info: fixed_count = sum( 1 for vi in verified_graph.value_info if all(dim.dim_value > 0 for dim in vi.type.tensor_type.shape.dim) ) total_count = len(verified_graph.value_info) print(f"✓ Fixed shapes: {fixed_count}/{total_count} intermediate tensors") print(f"\n✨ Success! Fixed model ready for deployment") return True except Exception as e: print(f"✗ Final verification failed: {e}") return False def _prepare_single_model(model_name, args, script_dir): """Export and fix shapes for one model. Returns True on success.""" onnx_filename = f"{model_name}.onnx" final_output = script_dir / onnx_filename info = SUPPORTED_MODELS[model_name] if not info['edge_suitable']: print(f"\n⚠ Warning: {model_name} has {info['params']} parameters (~" f"{int(info['gflops'])} GFLOPs).") print(f" This model requires ~16 GB+ RAM and is not suitable for edge deployment.") print(f" Shape inference and onnxsim will be skipped to reduce memory usage.") print() print(f"\nDINOv2 Model Preparation") print("=" * 80) print(f"Model: {model_name}") print(f"Output: {onnx_filename}") print(f"Input shape: [{args.batch_size}, {args.channels}, {args.height}, {args.width}]") if args.skip_export: if not final_output.exists(): print(f"\n✗ Error: ONNX file not found: {final_output}") print(" Run without --skip-export to export it first.") return False print(f"\nUsing existing ONNX file: {final_output.name}") else: if final_output.exists() and not args.force_export: print(f"\nONNX file already exists: {final_output}") file_size = final_output.stat().st_size print(f"File size: {file_size:,} bytes ({file_size / 1024 / 1024:.2f} MB)") print("Use --force-export to re-export.") return True success = export_to_onnx(model_name, final_output, args.height, args.width) if not success: return False # Disable onnxsim for non-edge models (>2 GB) to avoid OOM during shape fixing use_simplifier = not args.no_simplifier and info['edge_suitable'] success = fix_model_shape( final_output, final_output, batch_size=args.batch_size, channels=args.channels, height=args.height, width=args.width, use_simplifier=use_simplifier, ) if success: print("\n" + "=" * 80) print("COMPLETE!") print("=" * 80) print(f"Model: {model_name}") print(f"Output: {final_output.name}") print(f"Location: {script_dir}") print(f"Input shape: [{args.batch_size}, {args.channels}, {args.height}, {args.width}]") print(f"Config: {model_name}_config.yaml") else: print("\n✗ Shape fixing failed") return success def main(): parser = argparse.ArgumentParser( description=( 'Export DINOv2 classification models (backbone + linear head) from PyTorch Hub\n' 'to ONNX format with fixed static shapes for TI EdgeAI hardware deployment.\n' '\n' 'Processing pipeline:\n' ' 1. Load pretrained model from torch.hub (facebookresearch/dinov2)\n' ' 2. Export to ONNX (opset 17) with dynamic batch axis\n' ' 3. Fix dynamic input shapes to static [batch, channels, height, width]\n' ' 4. Run ONNX shape inference to propagate fixed shapes through all layers\n' ' 5. Run ONNX simplification via onnxsim (use --no-simplifier to skip)\n' ' 6. Validate the final model with onnx.checker' ), formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Export default model (ViT-S/14) %(prog)s # Export a specific variant %(prog)s --model dinov2_vitb14_lc # Export all supported models in sequence %(prog)s --model all # Export all models, skip onnxsim %(prog)s --model all --no-simplifier # Export with registers variant %(prog)s --model dinov2_vits14_reg_lc # Skip onnxsim (faster export, larger model file, shape inference still runs) %(prog)s --model dinov2_vitl14_lc --no-simplifier # Re-run shape inference + onnxsim on an already-exported ONNX file %(prog)s --model dinov2_vits14_lc --skip-export # Force re-export even if ONNX file exists %(prog)s --model dinov2_vits14_lc --force-export Available models: all Export all models listed below in sequence dinov2_vits14_lc ViT-S/14 distilled, 21M params, 81.1%% top-1 (recommended) dinov2_vitb14_lc ViT-B/14 distilled, 86M params, 84.5%% top-1 dinov2_vitl14_lc ViT-L/14 distilled, 307M params, 86.3%% top-1 dinov2_vitg14_lc ViT-g/14, 1100M params, 86.5%% top-1 (not for edge) dinov2_vits14_reg_lc ViT-S/14 + registers, 21M params, 80.9%% top-1 (recommended) dinov2_vitb14_reg_lc ViT-B/14 + registers, 86M params, 84.6%% top-1 dinov2_vitl14_reg_lc ViT-L/14 + registers, 307M params, 86.7%% top-1 dinov2_vitg14_reg_lc ViT-g/14 + registers, 1100M params, 87.1%% top-1 (not for edge) """ ) parser.add_argument( '--model', type=str, default='dinov2_vits14_lc', choices=list(SUPPORTED_MODELS.keys()) + ['all'], help=( 'Model variant to export, or "all" to export every supported model ' '(default: dinov2_vits14_lc)' ), ) parser.add_argument('--batch-size', type=int, default=1, help='Fixed batch size (default: 1)') parser.add_argument('--channels', type=int, default=3, help='Number of channels (default: 3)') parser.add_argument('--height', type=int, default=224, help='Image height (default: 224)') parser.add_argument('--width', type=int, default=224, help='Image width (default: 224)') parser.add_argument('--force-export', action='store_true', help='Force re-export even if ONNX file already exists') parser.add_argument('--skip-export', action='store_true', help='Skip export, only re-run shape inference + onnxsim on existing ONNX') parser.add_argument('--no-simplifier', action='store_true', help='Skip onnx-simplifier (onnxsim) step; shape inference still runs') args = parser.parse_args() script_dir = Path(__file__).parent if args.model == 'all': models = EDGE_MODELS skipped = [m for m in SUPPORTED_MODELS if m not in EDGE_MODELS] print(f"Exporting {len(models)} edge-suitable DINOv2 models...") if skipped: print(f"Skipping (not for edge, ~16 GB RAM required): {', '.join(skipped)}") results = {} for model_name in models: results[model_name] = _prepare_single_model(model_name, args, script_dir) print("\n" + "=" * 80) print("ALL MODELS SUMMARY") print("=" * 80) succeeded = [m for m, ok in results.items() if ok] failed = [m for m, ok in results.items() if not ok] for m in succeeded: print(f" ✓ {m}") for m in failed: print(f" ✗ {m}") print(f"\n{len(succeeded)}/{len(models)} models completed successfully.") sys.exit(0 if not failed else 1) else: ok = _prepare_single_model(args.model, args, script_dir) sys.exit(0 if ok else 1) if __name__ == '__main__': main()