raft testing

src/orisys/tracker.py

@@ -2,7 +2,8 @@ import cv2
 import numpy as np
 import torch
 import torchvision.transforms as T
-from torchvision.models.optical_flow import raft_large
+from torchvision.models.optical_flow import raft_large, raft_small
+import time
 
 class Tracker:
     def __init__(self,
@@ -54,6 +55,7 @@ class Tracker:
         self.__upper_threshold = config_tracker["upper_threshold"]  # error greater than threshold will add new reference frame
         ### RAFT
         self.__raft_device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"RAFT device: {self.__raft_device}")
         self.__raft_model = self.setup_simple_raft()
         self.__raft_ref_frame = None
         self.__raft_ref_tensor = None
@@ -250,7 +252,7 @@ class Tracker:
         model = model.eval()
         
         # Optimize matmul precision for faster inference (TF32 on Ampere+ GPUs)
-        torch.set_float32_matmul_precision('medium')
+        # torch.set_float32_matmul_precision('medium')
         
         # Enable CUDA-specific optimizations
         if self.__raft_device == 'cuda':
@@ -263,7 +265,7 @@ class Tracker:
             
         print(f"RAFT model loaded on {self.__raft_device}")
         return model
-
+    
     def _preprocess_frame_to_tensor(self, frame, resize=None):
         """Efficiently convert numpy frame to RAFT-ready tensor"""
         # Ensure contiguous memory for faster tensor conversion
@@ -296,7 +298,6 @@ class Tracker:
         # Add batch dimension
         tensor = tensor.unsqueeze(0)
         
-        # Keep in float32 - autocast will handle mixed precision during inference
         return tensor
     
     def track_raft(self, frame):
@@ -319,16 +320,22 @@ class Tracker:
         prev_tensor = self.__raft_ref_tensor
         curr_tensor = self._preprocess_frame_to_tensor(frame)
         
-        # OPTIMIZATION: Remove CUDA stream since immediate sync eliminates async benefit
-        # The bottleneck is inference, not data transfer (already optimized with pin_memory)
-        with torch.no_grad():
+        # Synchronize before timing to ensure preprocessing is complete
+        if self.__raft_device == 'cuda':
+            torch.cuda.synchronize()
+        
+        with torch.inference_mode():
             if self.__raft_device == 'cuda':
                 with torch.amp.autocast('cuda'):
                     flow = self.__raft_model(prev_tensor, curr_tensor)
             else:
                 flow = self.__raft_model(prev_tensor, curr_tensor)
             predicted_flows = flow[-1]
-
+        
+        # Synchronize after inference to get accurate timing
+        if self.__raft_device == 'cuda':
+            torch.cuda.synchronize()
+        
         # Optimize tensor->numpy conversion: contiguous() ensures optimal memory layout
         flow_np = predicted_flows[0].contiguous().cpu().numpy()