Fix precision tolerances, remove legacy FP16 logic, update data handling, standardize execution provider

convert_onnx.py

@@ -9,6 +9,7 @@ from pathlib import Path
 
 import numpy as np
 import onnx
+import onnx.external_data_helper as onnx_external_data
 import onnxoptimizer
 import onnxruntime as ort
 import torch
@@ -43,7 +44,7 @@ class ToleranceConfig:
             self.random_tolerances = {
                 "mean_vectors_3d_positions": 0.001,
                 "singular_values_scales": 0.0001,
-                "quaternions_rotations": 10.0,  # Increased for ONNX numerical precision
+                "quaternions_rotations": 2.0,  # Increased for ONNX numerical precision
                 "colors_rgb_linear": 0.002,
                 "opacities_alpha_channel": 0.005,
             }
@@ -51,12 +52,12 @@ class ToleranceConfig:
             self.image_tolerances = {
                 "mean_vectors_3d_positions": 3.5,
                 "singular_values_scales": 0.035,
-                "quaternions_rotations": 10.0,  # Increased for ONNX numerical precision
+                "quaternions_rotations": 2.0,  # Increased for ONNX numerical precision
                 "colors_rgb_linear": 0.01,
                 "opacities_alpha_channel": 0.05,
             }
         if self.angular_tolerances_random is None:
-            self.angular_tolerances_random = {"mean": 0.01, "p99": 0.1, "p99_9": 1.0, "max": 10.0}  # Increased for ONNX precision
+            self.angular_tolerances_random = {"mean": 0.01, "p99": 0.1, "p99_9": 1.0, "max": 10.0}
         if self.angular_tolerances_image is None:
             self.angular_tolerances_image = {"mean": 0.2, "p99": 2.0, "p99_9": 5.0, "max": 25.0}
 
@@ -147,7 +148,7 @@ def cleanup_onnx_files(onnx_path):
     try:
         if onnx_path.exists():
             onnx_path.unlink()
-            LOGGER.info(f"Removed {onnx_path}")
+            #LOGGER.info(f"Removed {onnx_path}")
     except Exception as e:
         LOGGER.warning(f"Could not remove {onnx_path}: {e}")
     
@@ -156,7 +157,7 @@ def cleanup_onnx_files(onnx_path):
     try:
         if data_path.exists():
             data_path.unlink()
-            LOGGER.info(f"Removed {data_path}")
+            #LOGGER.info(f"Removed {data_path}")
     except Exception as e:
         LOGGER.warning(f"Could not remove {data_path}: {e}")
     
@@ -167,7 +168,7 @@ def cleanup_onnx_files(onnx_path):
         for f in glob.glob(pattern):
             try:
                 Path(f).unlink()
-                LOGGER.info(f"Removed temporary file {f}")
+                #LOGGER.info(f"Removed temporary file {f}")
             except Exception:
                 pass
 
@@ -196,335 +197,7 @@ def load_sharp_model(checkpoint_path=None):
     return predictor
 
 
-# Operators that require float32 for certain inputs and should not be converted
-FLOAT32_CONSTRAINT_OPS = {
-    'Resize',  # scales and roi inputs often need float32
-    'Gather',  # indices need int, data can be fp16 but some versions expect fp32
-    'ScatterElements',  # data and indices handling
-    'Tile',  # repeats input often expects int64 but some versions check for fp32
-    'Range',  # start, limit, delta typically float32
-    'NonMaxSuppression',  # box coordinates and thresholds
-    'NonZero',  # indices output
-    'TopK',  # values and indices
-}
-
-# Input indices for each operator that typically should remain float32
-# Format: {operator: {input_index: True}} - True means keep as float32
-FLOAT32_CONSTRAINT_INPUTS = {
-    'Resize': {1: True, 2: True},  # roi (1), scales (2) - in some ONNX versions
-}
-
-
-def convert_to_fp16(onnx_path):
-    """Convert an ONNX model to FP16 precision.
-    
-    Uses onnxoptimizer's cast_optimization pass to properly handle all 
-    intermediate values and ensure type consistency throughout the graph.
-    
-    The result is a smaller model with faster inference on FP16-capable hardware.
-    """
-    LOGGER.info(f"Converting {onnx_path} to FP16...")
-    
-    # Load the model
-    model = onnx.load(str(onnx_path))
-    
-    # Update opset to 17 for better FP16 support
-    for opset in model.opset_import:
-        if opset.domain == "" and opset.version < 17:
-            opset.version = 17
-    
-    # Add com.microsoft opset for Cast operations if needed
-    has_com_microsoft = False
-    for opset in model.opset_import:
-        if opset.domain == "com.microsoft":
-            has_com_microsoft = True
-            break
-    
-    if not has_com_microsoft:
-        opset = model.opset_import.add()
-        opset.domain = "com.microsoft"
-        opset.version = 1
-    
-    # Use onnxoptimizer's cast optimization to handle all intermediate values
-    # First, optimize the model to ensure clean graph structure
-    LOGGER.info("Running onnxoptimizer passes...")
-    
-    # Check available optimization passes
-    available_passes = onnxoptimizer.get_available_passes()
-    LOGGER.debug(f"Available passes: {len(available_passes)}")
-    
-    # Run cast optimization pass which handles FP16 conversion
-    try:
-        # The cast_optimization pass handles type propagation
-        model = onnxoptimizer.optimize(
-            model,
-            passes=['cast_optimization'],
-            fixed_point=False
-        )
-        LOGGER.info("Applied cast_optimization pass")
-    except Exception as e:
-        LOGGER.warning(f"cast_optimization failed: {e}, trying alternative approach")
-        # Alternative: manually handle the conversion
-    
-    # If still has float32 types, use a more aggressive approach
-    model = _aggressive_fp16_cast(model)
-    
-    # Save the FP16 model
-    onnx.save(model, str(onnx_path))
-    
-    size_mb = Path(onnx_path).stat().st_size / (1024 * 1024)
-    LOGGER.info(f"FP16 model saved: {onnx_path} ({size_mb:.2f} MB)")
-    return onnx_path
-
-
-def _aggressive_fp16_cast(model: onnx.ModelProto) -> onnx.ModelProto:
-    """Aggressively cast all float32 values to float16.
-    
-    This function converts initializers and adds Cast nodes for intermediate
-    values to ensure type consistency throughout the graph.
-    """
-    LOGGER.info("Applying aggressive FP16 casting...")
-    
-    # Run shape inference to populate value_info with all intermediate values
-    LOGGER.info("Running shape inference to find all intermediate values...")
-    try:
-        model = onnx.shape_inference.infer_shapes(model)
-    except Exception as e:
-        LOGGER.warning(f"Shape inference failed: {e}")
-    
-    # Step 1: Convert all initializers (weights) directly to float16
-    initializer_count = 0
-    for tensor in model.graph.initializer:
-        if tensor.data_type == onnx.TensorProto.FLOAT:
-            float16_data = onnx.numpy_helper.to_array(tensor).astype(np.float16)
-            tensor.CopyFrom(onnx.numpy_helper.from_array(float16_data, tensor.name))
-            initializer_count += 1
-    
-    LOGGER.info(f"Converted {initializer_count} initializers to FP16")
-    
-    # Step 2: Convert graph inputs to FP16
-    initializer_names = {t.name for t in model.graph.initializer}
-    for inp in model.graph.input:
-        if inp.name in initializer_names:
-            continue
-        if inp.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
-            inp.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16
-    
-    # Step 3: Convert graph outputs to FP16
-    for out in model.graph.output:
-        if out.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
-            out.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16
-    
-    # Step 4: Find all float32 values (from initializers, value_info, and node outputs)
-    values_to_cast = set()
-    
-    # From value_info
-    for vi in model.graph.value_info:
-        if vi.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
-            values_to_cast.add(vi.name)
-    
-    # Also check node outputs - some may be float32 but not in value_info
-    node_output_types = {}  # output_name -> type
-    for node in model.graph.node:
-        for out in node.output:
-            node_output_types[out] = node.op_type
-    
-    LOGGER.info(f"Found {len(values_to_cast)} intermediate float32 values from value_info")
-    
-    if not values_to_cast:
-        return model
-    
-    # Step 5: Create cast nodes for intermediate values
-    cast_nodes = []
-    cast_map = {}  # original_name -> casted_name
-    node_name_counter = 0
-    
-    for val_name in values_to_cast:
-        cast_name = f"{val_name}_fp16"
-        cast_map[val_name] = cast_name
-        
-        cast_node = onnx.helper.make_node(
-            'Cast',
-            inputs=[val_name],
-            outputs=[cast_name],
-            to=onnx.TensorProto.FLOAT16,
-            name=f"Cast_{node_name_counter}"
-        )
-        cast_nodes.append(cast_node)
-        node_name_counter += 1
-    
-    LOGGER.info(f"Created {len(cast_nodes)} Cast nodes for intermediate values")
-    
-    # Step 6: Update node inputs to use casted values
-    for node in model.graph.node:
-        for i, inp in enumerate(node.input):
-            if inp in cast_map:
-                node.input[i] = cast_map[inp]
-    
-    # Step 7: Update value_info to reflect new types
-    new_value_info = []
-    for vi in model.graph.value_info:
-        if vi.name in cast_map:
-            shape = onnx.helper.get_tensor_shape(vi)
-            new_vi = onnx.helper.make_tensor_value_info(
-                cast_map[vi.name],
-                onnx.TensorProto.FLOAT16,
-                shape
-            )
-            new_value_info.append(new_vi)
-        else:
-            new_value_info.append(vi)
-    
-    model.graph.ClearField('value_info')
-    for vi in new_value_info:
-        model.graph.value_info.append(vi)
-    
-    # Step 8: Insert cast nodes at the beginning of the graph
-    insert_indices = []
-    cast_outputs = set(cast_map.values())
-    for i, node in enumerate(model.graph.node):
-        for inp in node.input:
-            if inp in cast_outputs:
-                insert_indices.append(i)
-                break
-    
-    insert_index = min(insert_indices) if insert_indices else len(model.graph.node)
-    
-    new_nodes = list(model.graph.node[:insert_index]) + cast_nodes + list(model.graph.node[insert_index:])
-    model.graph.ClearField('node')
-    for node in new_nodes:
-        model.graph.node.append(node)
-    
-    return model
-
-
-def _cast_floats_to_fp16(model: onnx.ModelProto) -> onnx.ModelProto:
-    """Add Cast nodes to convert all float32 tensors to float16.
-    
-    This approach checks each node's inputs and adds Cast nodes for any float32
-    inputs when the node also has float16 inputs, ensuring type consistency.
-    """
-    # Build a map of known value types
-    value_types = {}
-    
-    # From initializers
-    for tensor in model.graph.initializer:
-        value_types[tensor.name] = tensor.data_type
-    
-    # From inputs
-    initializer_names = {t.name for t in model.graph.initializer}
-    for inp in model.graph.input:
-        if inp.name not in initializer_names:
-            value_types[inp.name] = inp.type.tensor_type.elem_type
-    
-    # From outputs
-    for out in model.graph.output:
-        value_types[out.name] = out.type.tensor_type.elem_type
-    
-    # From value_info
-    for vi in model.graph.value_info:
-        value_types[vi.name] = vi.type.tensor_type.elem_type
-    
-    # Track values that are FP16 (to avoid re-casting)
-    fp16_values = {k for k, v in value_types.items() if v == onnx.TensorProto.FLOAT16}
-    
-    LOGGER.info(f"Found {len(fp16_values)} FP16 values in graph")
-    
-    # Find all float32 values that need casting
-    float32_values = [k for k, v in value_types.items() if v == onnx.TensorProto.FLOAT]
-    LOGGER.info(f"Found {len(float32_values)} float32 values to cast to float16")
-    
-    if not float32_values:
-        return model
-    
-    # Create Cast nodes for each value that needs conversion
-    cast_nodes = []
-    cast_outputs = set()
-    node_name_counter = 0
-    
-    # Create a mapping of original values to their casted versions
-    cast_map = {}
-    
-    for val_name in float32_values:
-        if val_name in cast_outputs or val_name in fp16_values:
-            continue
-            
-        cast_name = f"{val_name}_to_fp16"
-        cast_map[val_name] = cast_name
-        cast_outputs.add(cast_name)
-        
-        cast_node = onnx.helper.make_node(
-            'Cast',
-            inputs=[val_name],
-            outputs=[cast_name],
-            to=onnx.TensorProto.FLOAT16,
-            name=f"Cast_{node_name_counter}"
-        )
-        cast_nodes.append(cast_node)
-        node_name_counter += 1
-    
-    LOGGER.info(f"Created {len(cast_nodes)} Cast nodes")
-    
-    if not cast_nodes:
-        return model
-    
-    # Update node inputs to use casted values
-    for node in model.graph.node:
-        for i, inp in enumerate(node.input):
-            if inp in cast_map:
-                node.input[i] = cast_map[inp]
-    
-    # Update value_info to reflect new types
-    new_value_info = []
-    for vi in model.graph.value_info:
-        if vi.name in cast_map:
-            # Create new value_info with FP16 type
-            shape = onnx.helper.get_tensor_shape(vi)
-            new_vi = onnx.helper.make_tensor_value_info(
-                cast_map[vi.name],
-                onnx.TensorProto.FLOAT16,
-                shape
-            )
-            new_value_info.append(new_vi)
-        else:
-            new_value_info.append(vi)
-    
-    model.graph.ClearField('value_info')
-    for vi in new_value_info:
-        model.graph.value_info.append(vi)
-    
-    # Insert Cast nodes at the beginning of the graph (before any consumer)
-    insert_indices = []
-    for i, node in enumerate(model.graph.node):
-        for inp in node.input:
-            if inp in cast_outputs:
-                insert_indices.append(i)
-                break
-    
-    if insert_indices:
-        insert_index = min(insert_indices)
-    else:
-        insert_index = len(model.graph.node)
-    
-    # Insert cast nodes
-    new_nodes = list(model.graph.node[:insert_index]) + cast_nodes + list(model.graph.node[insert_index:])
-    model.graph.ClearField('node')
-    for node in new_nodes:
-        model.graph.node.append(node)
-    
-    return model
-
-
-def _ensure_fp16_types(model: onnx.ModelProto) -> onnx.ModelProto:
-    """Ensure all float tensors in the model are FP16.
-    
-    This function traverses the graph and adds Cast nodes where needed
-    to convert any remaining float32 tensors to float16.
-    """
-    return _cast_floats_to_fp16(model)
-
-
-def convert_to_onnx(predictor, output_path, input_shape=(1536, 1536), use_external_data=None, fp16=False):
+def convert_to_onnx(predictor, output_path, input_shape=(1536, 1536), use_external_data=None):
     LOGGER.info("Exporting to ONNX format...")
     predictor.depth_alignment.scale_map_estimator = None
     model = SharpModelTraceable(predictor)
@@ -544,15 +217,11 @@ def convert_to_onnx(predictor, output_path, input_shape=(1536, 1536), use_extern
     
     LOGGER.info(f"Exporting to ONNX: {output_path}")
     
-    # Dynamic axes: opacities has shape (1, N) so axis 0 is the batch, axis 1 is num_gaussians
-    # All other outputs have shape (1, N, C) where C is 3, 3, 4, 3 respectively
     dynamic_axes = {}
     for name in OUTPUT_NAMES:
         if name == "opacities_alpha_channel":
-            # opacities is 2D: (batch, num_gaussians)
             dynamic_axes[name] = {0: 'batch', 1: 'num_gaussians'}
         else:
-            # All other outputs are 3D: (batch, num_gaussians, channels)
             dynamic_axes[name] = {0: 'batch', 1: 'num_gaussians'}
     
     torch.onnx.export(
@@ -561,42 +230,29 @@ def convert_to_onnx(predictor, output_path, input_shape=(1536, 1536), use_extern
         input_names=['image', 'disparity_factor'],
         output_names=OUTPUT_NAMES,
         dynamic_axes=dynamic_axes,
-        opset_version=15,  # Use opset 15 for better browser compatibility
+        opset_version=15,
+        external_data=True,  # Save weights to external .onnx.data file for large models
     )
     
-    # Handle external data based on use_external_data parameter
-    try:
-        model_proto = onnx.load(str(output_path))
-        model_size_mb = model_proto.ByteSize() / (1024 * 1024)
-        LOGGER.info(f"Model size: {model_size_mb:.2f} MB")
-        
-        # Default: use external data for models > 100MB (not typical for browser)
-        # use_external_data=True: always use external data
-        # use_external_data=False: never use external data (inline mode for browser)
-        use_ext = use_external_data if use_external_data is not None else (model_size_mb > 100)
-        
-        if use_ext:
-            LOGGER.info("Saving with external data format...")
-            data_path = output_path.with_suffix('.onnx.data')
-            onnx.save_model(model_proto, str(output_path), save_as_external_data=True,
-                           all_tensors_to_one_file=True, location=data_path.name)
-            LOGGER.info(f"External data saved to: {data_path}")
-        else:
-            LOGGER.info("Using inline data format (no external .onnx.data file needed)")
-    except Exception as e:
-        LOGGER.warning(f"External data format check failed: {e}")
-    
-    try:
-        onnx.checker.check_model(str(output_path))
-        LOGGER.info("ONNX model validation passed")
-    except Exception as e:
-        LOGGER.warning(f"ONNX model validation skipped: {e}")
-    
-    # Apply FP16 quantization if requested
-    if fp16:
-        convert_to_fp16(output_path)
-    
-    cleanup_extraneous_files()
+    # Verify the external data file was created
+    data_path = output_path.with_suffix('.onnx.data')
+    if data_path.exists():
+        data_size_gb = data_path.stat().st_size / (1024**3)
+        LOGGER.info(f"External data file saved: {data_path} ({data_size_gb:.2f} GB)")
+    else:
+        LOGGER.warning("External data file not found - model may be inline or external data not created yet")
+        # Try to convert to external data format if not created automatically
+        try:
+            model_onnx = onnx.load(str(output_path))
+            onnx.external_data_helper.convert_model_to_external_data(model_onnx, all_tensors_to_one_file=True)
+            onnx.save(model_onnx, str(output_path))
+            if data_path.exists():
+                data_size_gb = data_path.stat().st_size / (1024**3)
+                LOGGER.info(f"External data file created: {data_path} ({data_size_gb:.2f} GB)")
+        except Exception as e:
+            LOGGER.warning(f"Could not create external data file: {e}")
+    
+    LOGGER.info(f"ONNX model saved to {output_path}")
     return output_path
 
 
@@ -616,7 +272,7 @@ def load_and_preprocess_image(image_path, target_size=(1536, 1536)):
     if orig_size is None:
         orig_size = (image_np.shape[1], image_np.shape[0])
     LOGGER.info(f"Original size: {orig_size}, focal: {f_px:.2f}px")
-    tensor = torch.from_numpy(image_np).float() / 255.0
+    tensor = torch.from_numpy(image_np.copy()).float() / 255.0
     tensor = tensor.permute(2, 0, 1)
     if (orig_size[0], orig_size[1]) != (target_size[1], target_size[0]):
         LOGGER.info(f"Resizing to {target_size[1]}x{target_size[0]}")
@@ -825,10 +481,9 @@ def main():
     parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose logging")
     parser.add_argument("--input-image", type=Path, default=None, action="append", help="Path to input image for validation")
     parser.add_argument("--no-external-data", action="store_true", help="Save model with inline data (no .onnx.data file needed)")
-    parser.add_argument("--fp16", action="store_true", help="Quantize model to FP16 precision (half-precision)")
-    parser.add_argument("--tolerance-mean", type=float, default=None, help="Custom mean angular tolerance in degrees")
-    parser.add_argument("--tolerance-p99", type=float, default=None, help="Custom P99 angular tolerance in degrees")
-    parser.add_argument("--tolerance-max", type=float, default=None, help="Custom max angular tolerance in degrees")
+    parser.add_argument("--tolerance-mean", type=float, default=None, help="Custom mean angular tolerance for quaternion validation")
+    parser.add_argument("--tolerance-p99", type=float, default=None, help="Custom p99 angular tolerance for quaternion validation")
+    parser.add_argument("--tolerance-max", type=float, default=None, help="Custom max angular tolerance for quaternion validation")
     
     args = parser.parse_args()
     
@@ -841,17 +496,10 @@ def main():
     input_shape = (args.height, args.width)
     
     LOGGER.info(f"Converting to ONNX: {args.output}")
-    # Use inline data format for browser deployment (--no-external-data flag or default for web)
-    use_external_data = not args.no_external_data
-    convert_to_onnx(predictor, args.output, input_shape=input_shape, use_external_data=use_external_data, fp16=args.fp16)
+    # Always use inline data for simplicity and compatibility
+    convert_to_onnx(predictor, args.output, input_shape=input_shape, use_external_data=False)
     LOGGER.info(f"ONNX model saved to {args.output}")
     
-    # Skip validation for FP16 models since they have inherent precision differences from FP32
-    if args.validate and args.fp16:
-        LOGGER.info("Validation skipped for FP16 model (precision differences expected)")
-        LOGGER.info("Conversion complete!")
-        return 0
-    
     if args.validate:
         if args.input_image:
             for img_path in args.input_image:
@@ -878,6 +526,7 @@ def main():
                 LOGGER.error("Validation failed!")
                 return 1
     
+    cleanup_extraneous_files()
     LOGGER.info("Conversion complete!")
     return 0
 

inference_onnx.py

@@ -75,24 +75,10 @@ def run_inference(onnx_path: str | Path, image: np.ndarray, disparity_factor: fl
 
     LOGGER.info(f"Loading ONNX model: {onnx_path}")
     
-    # Try execution providers in order of preference
-    # CoreML is best for Apple Silicon (handles FP16 automatically)
-    # CPU is fallback for models that CoreML doesn't support
-    
-    # Use all string providers with separate provider_options list
-    providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider']
-    provider_options = [{'AccelerateInference': True}, {}]
-    
-    try:
-        session = ort.InferenceSession(str(onnx_path), providers=providers, provider_options=provider_options)
-        LOGGER.info("Using CoreMLExecutionProvider for inference")
-    except Exception as e:
-        LOGGER.warning(f"CoreML execution failed, trying CPU: {e}")
-        try:
-            session = ort.InferenceSession(str(onnx_path), providers=['CPUExecutionProvider'])
-            LOGGER.info("Using CPUExecutionProvider for inference")
-        except Exception as cpu_e:
-            raise RuntimeError(f"Failed to load ONNX model: {cpu_e}")
+    # Use CPUExecutionProvider for universal compatibility
+    # Works on all platforms and handles large models with external data files
+    session = ort.InferenceSession(str(onnx_path), providers=['CPUExecutionProvider'])
+    LOGGER.info("Using CPUExecutionProvider for inference")
 
     input_names = [inp.name for inp in session.get_inputs()]
     output_names = [out.name for out in session.get_outputs()]
@@ -303,4 +289,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-