From d750a039fef20c6f9a48856842a8365fd73cd23e Mon Sep 17 00:00:00 2001 From: bluepulsar Date: Sat, 31 Jan 2026 16:10:34 +0800 Subject: [PATCH] pin memory --- src/orisys/tracker.py | 42 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 34 insertions(+), 8 deletions(-) diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index d31e117..5028119 100644 --- a/src/orisys/tracker.py +++ b/src/orisys/tracker.py @@ -243,21 +243,45 @@ class Tracker: def setup_simple_raft(self): """ - Setup a simpler version using available models + Setup a simpler version using available models with performance optimizations """ # Load model and move to GPU model = raft_large(pretrained=True, progress=False).to(self.__raft_device) model = model.eval() + + # Optimize matmul precision for faster inference (TF32 on Ampere+ GPUs) + torch.set_float32_matmul_precision('medium') + + # Enable CUDA-specific optimizations + if self.__raft_device == 'cuda': + # cuDNN optimizations + torch.backends.cudnn.benchmark = True # Auto-select fastest algorithms + torch.backends.cudnn.deterministic = False # Allow non-deterministic for speed + # TensorFloat-32 for Ampere+ GPUs (faster with minimal accuracy loss) + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True 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""" - # Upload to device as uint8 first to save bandwidth - # permute returns a view, but from_numpy usually shares memory. - # Moving to GPU as uint8 is fast. - tensor = torch.from_numpy(frame).permute(2, 0, 1).to(self.__raft_device) + # Ensure contiguous memory for faster tensor conversion + if not frame.flags['C_CONTIGUOUS']: + frame = np.ascontiguousarray(frame) + + # Upload to device as uint8 first to save bandwidth (4x less data than float32) + tensor = torch.from_numpy(frame).permute(2, 0, 1) + + # Move to device with non_blocking for async transfer (if CUDA) + # Pin memory for faster CPU-GPU transfer (similar to DataLoader pin_memory=True) + if self.__raft_device == 'cuda': + # Pin memory before transferring to GPU (must be on CPU first) + if not tensor.is_pinned(): + tensor = tensor.pin_memory() + tensor = tensor.to(self.__raft_device, non_blocking=True) + else: + tensor = tensor.to(self.__raft_device) # Convert to float and normalize on device tensor = tensor.float() / 255.0 @@ -274,7 +298,7 @@ class Tracker: # Keep in float32 - autocast will handle mixed precision during inference return tensor - + def track_raft(self, frame): # No need to copy if we just read # frame is expected to be (H, W, 3) or (H, W) @@ -295,7 +319,8 @@ class Tracker: prev_tensor = self.__raft_ref_tensor curr_tensor = self._preprocess_frame_to_tensor(frame) - # Use autocast for automatic mixed precision (FP16 where supported, FP32 where needed) + # 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(): if self.__raft_device == 'cuda': with torch.amp.autocast('cuda'): @@ -304,7 +329,8 @@ class Tracker: flow = self.__raft_model(prev_tensor, curr_tensor) predicted_flows = flow[-1] - flow_np = predicted_flows[0].float().cpu().numpy() # Convert back to float32 for numpy + # Optimize tensor->numpy conversion: contiguous() ensures optimal memory layout + flow_np = predicted_flows[0].contiguous().cpu().numpy() # Transpose to (H, W, 2) format expected by draw_arrows flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2)