HarshithReddy01/srmamamba-liver-segmentation
1
1import os2import time3import torch4from monai.inferers import SlidingWindowInferer5from config import BUILD_SRMAMAMBA_AVAILABLE, build_SRMAMamba, SRMA_MAMBA_DIR6 7MODEL_T1 = None8MODEL_T2 = None9DEVICE = torch.device('cpu')10WINDOW_INFER = None11 12def clear_gpu_memory():13 global MODEL_T1, MODEL_T2, WINDOW_INFER14 if torch.cuda.is_available():15 if MODEL_T1 is not None:16 del MODEL_T117 MODEL_T1 = None18 if MODEL_T2 is not None:19 del MODEL_T220 MODEL_T2 = None21 if WINDOW_INFER is not None:22 del WINDOW_INFER23 WINDOW_INFER = None24 torch.cuda.empty_cache()25 torch.cuda.synchronize()26 print(" → GPU memory cleared (models unloaded)")27 28def load_model(modality='T1'):29 global MODEL_T1, MODEL_T2, DEVICE, WINDOW_INFER, BUILD_SRMAMAMBA_AVAILABLE30 31 if torch.cuda.is_available():32 torch.cuda.empty_cache()33 torch.cuda.synchronize()34 35 if not BUILD_SRMAMAMBA_AVAILABLE or build_SRMAMamba is None:36 error_msg = "Model builder (build_SRMAMamba) is not available. Please check the logs for import errors."37 print(f"✗ {error_msg}")38 raise ImportError(error_msg)39 40 print(f"Loading {modality} model...")41 42 if torch.cuda.is_available():43 try:44 max_retries = 345 retry_delay = 246 47 for attempt in range(max_retries):48 try:49 torch.cuda.empty_cache()50 test_tensor = torch.zeros(1).cuda()51 del test_tensor52 torch.cuda.synchronize()53 DEVICE = torch.device('cuda')54 print(f"✓ Using device: {DEVICE}")55 break56 except RuntimeError as e:57 if "CUDA" in str(e) and attempt < max_retries - 1:58 print(f"⚠ GPU wake-up attempt {attempt + 1}/{max_retries}: {e}")59 print(f"⚠ Waiting {retry_delay}s for GPU to wake up...")60 time.sleep(retry_delay)61 retry_delay *= 262 else:63 raise64 except Exception as e:65 print(f"⚠ CUDA available but failed to initialize: {e}. Falling back to CPU.")66 DEVICE = torch.device('cpu')67 else:68 DEVICE = torch.device('cpu')69 print(f"ℹ CUDA not available. Using device: {DEVICE}")70 71 if DEVICE.type == 'cuda':72 torch.cuda.empty_cache()73 torch.cuda.synchronize()74 allocated = torch.cuda.memory_allocated(0) / (1024**3)75 reserved = torch.cuda.memory_reserved(0) / (1024**3)76 total = torch.cuda.get_device_properties(0).total_memory / (1024**3)77 free_memory_gb = total - allocated78 print(f" → GPU memory: {allocated:.2f} GB allocated, {reserved:.2f} GB reserved, {free_memory_gb:.2f} GB free (total: {total:.2f} GB)")79 80 if free_memory_gb < 1.0:81 print(f" ⚠ CRITICAL: Very low free memory ({free_memory_gb:.2f} GB). Using ultra-minimal settings.")82 size = [192, 192, 32]83 batch_size = 184 overlap = 0.2585 elif free_memory_gb < 2.0:86 print(f" ⚠ WARNING: Very low free memory ({free_memory_gb:.2f} GB). Using minimal settings.")87 size = [192, 192, 32]88 batch_size = 189 overlap = 0.2590 elif free_memory_gb < 5.0:91 size = [224, 224, 48]92 batch_size = 193 overlap = 0.294 elif free_memory_gb > 40:95 print(f" Very high VRAM GPU detected ({free_memory_gb:.2f} GB free). Using optimal settings for maximum speed.")96 size = [256, 256, 80]97 batch_size = 298 overlap = 0.199 elif free_memory_gb > 30:100 print(f" High VRAM GPU detected ({free_memory_gb:.2f} GB free). Using optimal settings for speed.")101 size = [256, 256, 64]102 batch_size = 2103 overlap = 0.1104 elif free_memory_gb > 25:105 print(f" ✓ Large VRAM GPU detected ({free_memory_gb:.2f} GB free). Using optimal settings.")106 size = [256, 256, 64]107 batch_size = 1108 overlap = 0.1109 elif free_memory_gb > 20:110 size = [256, 256, 64]111 batch_size = 1112 overlap = 0.1113 elif free_memory_gb > 15:114 size = [256, 256, 64]115 batch_size = 1116 overlap = 0.1117 elif free_memory_gb > 10:118 size = [224, 224, 64]119 batch_size = 1120 overlap = 0.1121 elif free_memory_gb > 8:122 size = [224, 224, 48]123 batch_size = 1124 overlap = 0.2125 else:126 size = [192, 192, 48]127 batch_size = 1128 overlap = 0.2129 else:130 size = [224, 224, 64]131 batch_size = 1132 overlap = 0.15133 134 print(f" → Sliding window config: roi_size={size}, sw_batch_size={batch_size}, overlap={overlap}")135 136 print("Building model architecture...")137 if SRMA_MAMBA_DIR:138 original_cwd = os.getcwd()139 try:140 os.chdir(SRMA_MAMBA_DIR)141 print(f"Changed working directory to: {SRMA_MAMBA_DIR}")142 model = build_SRMAMamba()143 print("✓ Model architecture built")144 finally:145 os.chdir(original_cwd)146 else:147 model = build_SRMAMamba()148 print("✓ Model architecture built")149 150 model = model.to(DEVICE)151 print(f"✓ Model moved to {DEVICE}")152 153 checkpoint_path = f"checkpoint_{modality}.pth"154 possible_paths = [155 checkpoint_path,156 os.path.join(os.path.dirname(__file__), checkpoint_path),157 f"../../Chkpoints/checkpoint_{modality}.pth",158 f"Chkpoints/checkpoint_{modality}.pth",159 f"../Chkpoints/checkpoint_{modality}.pth",160 f"Model/Chkpoints/checkpoint_{modality}.pth",161 os.path.join(os.path.dirname(__file__), f"Chkpoints/checkpoint_{modality}.pth"),162 ]163 164 found = False165 for path in possible_paths:166 abs_path = os.path.abspath(path)167 if os.path.exists(path) or os.path.exists(abs_path):168 checkpoint_path = path if os.path.exists(path) else abs_path169 found = True170 print(f"✓ Found checkpoint at: {checkpoint_path}")171 break172 173 if not found:174 try:175 from huggingface_hub import hf_hub_download176 repo_id = os.environ.get("HF_MODEL_REPO", "HarshithReddy01/srmamamba-liver-segmentation")177 print(f"Attempting to download checkpoint from Hugging Face: {repo_id}")178 checkpoint_path = hf_hub_download(179 repo_id=repo_id,180 filename=f"checkpoint_{modality}.pth",181 cache_dir="."182 )183 found = True184 print(f"✓ Downloaded checkpoint to: {checkpoint_path}")185 except Exception as e:186 error_msg = f"Checkpoint not found. Searched: {possible_paths}. Hugging Face download failed: {str(e)}"187 print(f"✗ {error_msg}")188 raise FileNotFoundError(error_msg)189 190 print(f"Loading checkpoint weights from: {checkpoint_path}")191 try:192 checkpoint = torch.load(checkpoint_path, map_location=DEVICE)193 if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:194 ckpt_sd = checkpoint['state_dict']195 else:196 ckpt_sd = checkpoint197 model_keys = set(model.state_dict().keys())198 ckpt_keys = set(ckpt_sd.keys())199 missing = model_keys - ckpt_keys200 unexpected = ckpt_keys - model_keys201 if missing:202 print(f" ⚠ WARNING: Checkpoint missing keys for model: {len(missing)} keys")203 if unexpected:204 print(f" ⚠ WARNING: Checkpoint has unexpected keys: {len(unexpected)} keys")205 if not missing and not unexpected:206 print(" ✓ State dict keys match (no missing/unexpected keys)")207 if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:208 model.load_state_dict(ckpt_sd)209 else:210 model.load_state_dict(ckpt_sd)211 print("✓ Checkpoint loaded successfully")212 except Exception as e:213 print(f"✗ Failed to load checkpoint: {e}")214 raise215 216 model.eval()217 print("✓ Model set to evaluation mode")218 219 if DEVICE.type == 'cuda':220 import config221 from packaging import version222 223 torch_version = version.parse(torch.__version__)224 if torch_version >= version.parse("2.9.0"):225 torch.backends.cuda.matmul.fp32_precision = 'tf32'226 torch.backends.cudnn.conv.fp32_precision = 'tf32'227 tf32_matmul = torch.backends.cuda.matmul.fp32_precision228 tf32_conv = torch.backends.cudnn.conv.fp32_precision229 else:230 torch.backends.cuda.matmul.allow_tf32 = True231 torch.backends.cudnn.allow_tf32 = True232 tf32_matmul = 'tf32' if torch.backends.cuda.matmul.allow_tf32 else 'ieee'233 tf32_conv = 'tf32' if torch.backends.cudnn.allow_tf32 else 'ieee'234 torch.backends.cudnn.benchmark = True235 print(f"TF32 enabled: matmul={tf32_matmul}, conv={tf32_conv}")236 print("cuDNN benchmarking enabled")237 238 if config.ENABLE_TORCH_COMPILE:239 try:240 compile_mode = os.environ.get('TORCH_COMPILE_MODE', 'reduce-overhead')241 if compile_mode == 'max-autotune':242 print(f" → Compiling with max-autotune (may take 2-5 min on first run)...")243 model = torch.compile(model, mode='max-autotune', fullgraph=False)244 print(f"✓ Model compiled with torch.compile (mode=max-autotune, fullgraph=False)")245 elif compile_mode == 'default':246 print(f" → Compiling with default mode (may take 1-3 min on first run)...")247 model = torch.compile(model, fullgraph=False)248 print(f"✓ Model compiled with torch.compile (mode=default, fullgraph=False)")249 else:250 print(f" → Compiling with reduce-overhead (faster first run, ~30-60s)...")251 model = torch.compile(model, mode='reduce-overhead', fullgraph=False)252 print(f"✓ Model compiled with torch.compile (mode=reduce-overhead, fullgraph=False)")253 except Exception as e:254 print(f" ⚠ torch.compile failed: {e}. Continuing without compilation.")255 else:256 print(" ℹ torch.compile disabled (set ENABLE_TORCH_COMPILE=true to enable)")257 258 torch.cuda.empty_cache()259 torch.cuda.synchronize()260 allocated_after_load = torch.cuda.memory_allocated(0) / (1024**3)261 free_after_load = (torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated(0)) / (1024**3)262 print(f" → GPU memory after model load: {allocated_after_load:.2f} GB allocated, {free_after_load:.2f} GB free")263 264 if free_after_load < 1.0:265 print(f" ⚠ CRITICAL: Only {free_after_load:.2f} GB free after model load. Using ultra-minimal settings.")266 size = [192, 192, 32]267 batch_size = 1268 overlap = 0.25269 elif free_after_load < 2.0:270 print(f" ⚠ WARNING: Low free memory ({free_after_load:.2f} GB) after model load. Adjusting to minimal settings.")271 size = [192, 192, 32]272 batch_size = 1273 overlap = 0.25274 elif free_after_load > 40:275 print(f" Excellent free memory ({free_after_load:.2f} GB) after model load. Using optimal settings for maximum speed.")276 size = [256, 256, 80]277 batch_size = 2278 overlap = 0.1279 elif free_after_load > 30:280 print(f" Excellent free memory ({free_after_load:.2f} GB) after model load. Using optimal settings for speed.")281 size = [256, 256, 64]282 batch_size = 2283 overlap = 0.1284 elif free_after_load > 25:285 print(f" ✓ Good free memory ({free_after_load:.2f} GB) after model load. Using optimal settings.")286 size = [256, 256, 64]287 batch_size = 1288 overlap = 0.1289 elif free_after_load > 20:290 print(f" ✓ Good free memory ({free_after_load:.2f} GB) after model load. Using optimal settings.")291 size = [256, 256, 64]292 batch_size = 1293 overlap = 0.1294 elif free_after_load > 15:295 size = [256, 256, 64]296 batch_size = 1297 overlap = 0.1298 elif free_after_load < 5.0 and (size[0] > 224 or batch_size > 1):299 print(f" ⚠ WARNING: Limited free memory ({free_after_load:.2f} GB). Reducing window size and batch size.")300 size = [224, 224, 48]301 batch_size = 1302 overlap = 0.1303 304 aggregation_device = 'cuda'305 if free_after_load < 2.0:306 aggregation_device = 'cpu'307 print(f" → Very low VRAM ({free_after_load:.2f} GB), using CPU aggregation to prevent OOM")308 else:309 print(f" → Using GPU aggregation for maximum speed (VRAM: {free_after_load:.2f} GB free)")310 311 WINDOW_INFER = SlidingWindowInferer(312 roi_size=size, 313 sw_batch_size=batch_size, 314 overlap=overlap,315 sw_device='cuda',316 device=aggregation_device317 )318 print(f"✓ Sliding window inferer created (GPU compute, {aggregation_device.upper()} aggregation)")319 320 if DEVICE.type == 'cuda':321 if config.ENABLE_TORCH_COMPILE:322 print(" Running warm-up inference to trigger compilation and kernel autotuning...")323 print(" This may take 30-60s (reduce-overhead) or 2-5min (max-autotune) on first run...")324 else:325 print(" Running warm-up inference to trigger kernel autotuning...")326 try:327 dummy_input = torch.randn(1, 1, size[0], size[1], size[2], device=DEVICE, dtype=torch.float32)328 dummy_input = dummy_input.contiguous(memory_format=torch.channels_last_3d)329 warmup_start = time.time()330 with torch.no_grad():331 from torch.amp import autocast332 with autocast(device_type='cuda'):333 _ = model(dummy_input)334 torch.cuda.synchronize()335 warmup_time = time.time() - warmup_start336 del dummy_input, _337 torch.cuda.empty_cache()338 if config.ENABLE_TORCH_COMPILE:339 print(f" Warm-up completed in {warmup_time:.1f}s (compilation + kernel autotuning)")340 else:341 print(f" Warm-up completed in {warmup_time:.1f}s (kernels autotuned)")342 except RuntimeError as e:343 if "out of memory" in str(e):344 print(f" Warm-up OOM (non-critical): {e}")345 print(f" Will use progressive fallback during inference")346 else:347 print(f" Warm-up failed (non-critical): {e}")348 except Exception as e:349 print(f" Warm-up failed (non-critical): {e}")350 351 if modality == 'T1':352 MODEL_T1 = model353 else:354 MODEL_T2 = model355 356 print(f"✓ {modality} model loaded and ready")357 return model358 359