Upload feature_extractor.py

src/modules/feature_extractor.py

@@ -34,6 +34,9 @@ class FeatureExtractor:
         # Check if ONNX model exists
         if os.path.exists(onnx_path):
             print(f"Loading existing ONNX model from {onnx_path}")
+            # Get model size in MB
+            model_size = os.path.getsize(onnx_path) / (1024 * 1024)
+            print(f"Model size: {model_size:.2f} MB")
             self.onnx_session = onnxruntime.InferenceSession(onnx_path)
         else:
             print(
@@ -42,8 +45,8 @@ class FeatureExtractor:
             # Initialize PyTorch model
             self.model, _ = self.init_model(base_model)
             self.model.eval()
-            self.device = torch.device("cpu")
-            self.model.to(self.device)
+            # self.device = torch.device("cpu")
+            # self.model.to(self.device).half()
 
             # Create directory if it doesn't exist
             os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
@@ -53,7 +56,10 @@ class FeatureExtractor:
 
             # Load the newly created ONNX model
             self.onnx_session = onnxruntime.InferenceSession(onnx_path)
+            # Get model size in MB
+            model_size = os.path.getsize(onnx_path) / (1024 * 1024)
             print(f"Successfully created and loaded ONNX model from {onnx_path}")
+            print(f"Model size: {model_size:.2f} MB")
 
     def init_model(self, base_model):
         """Initialize the model for feature extraction
@@ -92,10 +98,12 @@ class FeatureExtractor:
         # add batch dimension
         x = x.unsqueeze(0)
 
-        # Convert to numpy for ONNX Runtime
+        # Convert to float16 and then to numpy for ONNX Runtime
+        x = x.half()
         x_numpy = x.numpy()
-        # Run inference with ONNX Runtime
-        print("Running inference with ONNX Runtime")
+        # 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)")
         output = self.onnx_session.run(
             None,
             {'input': x_numpy}
@@ -118,6 +126,10 @@ class FeatureExtractor:
         dummy_input = torch.randn(1, 3, 224, 224, device=self.device)
 
         # Export the model
+        # Convert model to float16
+        self.model.half()
+        dummy_input = dummy_input.half()
+        
         torch.onnx.export(
             self.model,
             dummy_input,