Update src/modules/feature_extractor.py

src/modules/feature_extractor.py

@@ -4,7 +4,8 @@ import os
 import torch
 import onnx
 import onnxruntime
-
+from onnxconverter_common import float16
+import numpy as np
 from src.modules.config_extractor import MODEL_CONFIG
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
@@ -99,7 +100,7 @@ class FeatureExtractor:
         x = x.unsqueeze(0)
 
         # Convert to numpy for ONNX Runtime
-        x_numpy = x.numpy()
+        x_numpy = x.numpy().astype(np.float16)
         # Get model size in MB
         model_size = os.path.getsize(self.onnx_path) / (1024 * 1024)
         print(f"Running inference with ONNX model (size: {model_size:.2f} MB)")
@@ -124,7 +125,7 @@ class FeatureExtractor:
         # Create a dummy input tensor
         dummy_input = torch.randn(1, 3, 224, 224, device=self.device)
 
-        # Export the model
+        # Export the model to ONNX (FP32 first)
         torch.onnx.export(
             self.model,
             dummy_input,
@@ -140,7 +141,14 @@ class FeatureExtractor:
             }
         )
 
-        # Verify the exported model
+        # Load and verify the exported model
+        print("Converting model to float16...")
         onnx_model = onnx.load(save_path)
         onnx.checker.check_model(onnx_model)
-        print(f"ONNX model saved to {save_path}")
+        
+        # Convert to float16
+        model_fp16 = float16.convert_float_to_float16(onnx_model)
+        
+        # Save the float16 model
+        onnx.save(model_fp16, save_path)
+        print(f"Float16 ONNX model saved to {save_path}")