""" Inference Script for Solar Flare Prediction Models ================================================= Runs a trained ViTLocal checkpoint over a folder of AIA data. Computes soft X-ray (SXR) predictions, and saves attention weights, flux contributions, and final outputs. The workflow includes: - Loading configuration parameters from a YAML file. - Resolving dynamic variables in the config. - Loading the model checkpoint and preparing it for inference. - Performing batched evaluation over AIA/GOES datasets. - Saving predicted fluxes, ground truth, and visualization-ready artifacts. """ import argparse import re import sys import gc import pandas as pd import torch import numpy as np from torch.utils.data import DataLoader from pathlib import Path import os import yaml import torch.nn.functional as F from torch.nn import HuberLoss from tqdm import tqdm # Add project root to Python path PROJECT_ROOT = Path(__file__).parent.parent.absolute() sys.path.insert(0, str(PROJECT_ROOT)) from forecasting.dataset import AIAGOESDataset from forecasting.model import ViTLocal device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') def evaluate_model_on_dataset(model, dataset, batch_size=16, times=None, config_data=None, save_weights=True, input_size=512, patch_size=16, save_flux=False): """ Run batched inference on the dataset and yield predictions, attention maps, and flux data. Memory optimization: Processes attention weights immediately and moves to CPU to reduce GPU memory usage. Parameters ---------- model : torch.nn.Module Loaded solar flare prediction model. dataset : torch.utils.data.Dataset Dataset containing AIA images and corresponding SXR values. batch_size : int, default=16 Number of samples per batch. times : list, optional List of timestamps corresponding to each sample. config_data : dict, optional YAML configuration dictionary. save_weights : bool, default=True Whether to save attention weights for visualization. input_size : int, default=512 Input image resolution. patch_size : int, default=16 Patch size for ViT-based models. save_flux : bool, default=False Whether to save flux contributions. Yields ------ tuple (predictions, ground_truth, attention_map, flux_map, global_index) """ model.eval() data_device = next(model.parameters()).device # Optimize DataLoader settings num_workers = config_data.get('num_workers', 4) if config_data else 4 pin_memory = config_data.get('pin_memory', True) if config_data else True n_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1 use_multi_gpu = str(config_data.get('multi_gpu', False) if config_data else False).lower() == 'true' # Unwrapped model used as fallback for batches smaller than n_gpus base_model = model.module if isinstance(model, torch.nn.DataParallel) else model loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=False, multiprocessing_context='spawn' if num_workers > 0 else None) # Load SXR normalization only if path is provided and not empty sxr_norm = None if (config_data and 'data' in config_data and 'sxr_norm_path' in config_data['data'] and config_data['data']['sxr_norm_path'] and config_data['data']['sxr_norm_path'].strip()): try: sxr_norm = np.load(config_data['data']['sxr_norm_path']) except FileNotFoundError: sxr_norm = None # All models are ViTLocal with localized attention (no CLS token) grid_h, grid_w = input_size // patch_size, input_size // patch_size use_amp = (data_device.type == 'cuda' and config_data.get('use_amp', False) if config_data else False) try: with torch.no_grad(): for batch_idx, batch in enumerate(loader): aia_imgs = batch[0] sxr = batch[1] aia_imgs = aia_imgs.to(data_device, non_blocking=True) # Call model — optionally use AMP (set use_amp: true in config). # FP16 on V100 can spike peak memory due to FP32 fallbacks in attention. # CRITICAL: ViTLocal defaults to return_attention=True, which uses massive memory # Only compute attention weights if we're saving them (save_weights=True) # Fall back to single GPU for batches smaller than n_gpus to avoid # DataParallel crashing when some replicas receive empty inputs. active_model = (base_model if isinstance(model, torch.nn.DataParallel) and aia_imgs.shape[0] < n_gpus else model) with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp): if save_weights: pred = active_model(aia_imgs, return_attention=True) else: pred = active_model(aia_imgs, return_attention=False) # Extract outputs (ViTLocal returns tuple of (predictions, attention_weights, flux_contributions) when return_attention=True) if isinstance(pred, tuple) and len(pred) >= 3: predictions = pred[0] weights = pred[1] if save_weights else None flux_contributions = pred[2] if save_flux else None elif isinstance(pred, tuple) and len(pred) == 2: # When return_attention=False, returns (predictions, flux_contributions) predictions = pred[0] weights = None flux_contributions = pred[1] if save_flux else None else: predictions = pred weights = None flux_contributions = None current_batch_size = predictions.shape[0] # Process each sample in the batch to reduce memory footprint for i in range(current_batch_size): global_idx = batch_idx * batch_size + i # Process attention weights immediately and move to CPU to free GPU memory weight_data = None if save_weights and weights is not None: try: # Extract attention for this sample only (localized model - no CLS token) last_layer_attention = weights[-1][i].detach() # Detach to allow garbage collection if last_layer_attention is not None: avg_attention = last_layer_attention.mean(dim=0) if not torch.isnan(avg_attention).any(): # Move to CPU immediately to free GPU memory # Localized attention: average over heads, then over patches patch_attention = avg_attention.mean(dim=0).cpu() attention_map = patch_attention.reshape(grid_h, grid_w) weight_data = attention_map.numpy() # Clear GPU tensors del last_layer_attention, avg_attention, patch_attention # Save weight if needed if config_data and 'weight_path' in config_data and global_idx < len(times): weight_dir = Path(config_data['weight_path']) weight_dir.mkdir(parents=True, exist_ok=True) weight_file = os.path.join(config_data['weight_path'], f"{times[global_idx]}.npy") np.save(weight_file, weight_data) except Exception as e: weight_data = np.zeros((grid_h, grid_w)) # Process flux contributions flux_data = None if save_flux and flux_contributions is not None: try: flux_contrib = flux_contributions[i].detach().cpu() flux_contrib_map = flux_contrib.reshape(grid_h, grid_w) flux_data = flux_contrib_map.numpy() # Save flux if needed if config_data and 'flux_path' in config_data and global_idx < len(times): flux_dir = Path(config_data['flux_path']) flux_dir.mkdir(parents=True, exist_ok=True) flux_file = os.path.join(config_data['flux_path'], f"{times[global_idx]}.npy") np.save(flux_file, flux_data) del flux_contrib, flux_contrib_map except Exception: flux_data = None # Yield result immediately - don't accumulate in memory yield (predictions[i].cpu().numpy(), sxr[i].cpu().numpy(), weight_data, flux_data, global_idx) # Clear all batch-level tensors from GPU memory immediately # CRITICAL: Delete weights first as they're the biggest memory consumer if weights is not None: # Delete each layer's attention weights explicitly to free memory for layer_idx, layer_weights in enumerate(weights): if isinstance(layer_weights, torch.Tensor): del layer_weights elif isinstance(layer_weights, (list, tuple)): for w in layer_weights: if isinstance(w, torch.Tensor): del w del weights # Delete the list itself weights = None # Ensure reference is cleared del predictions, aia_imgs if flux_contributions is not None: del flux_contributions flux_contributions = None # Only clear cache occasionally — doing it every batch stalls the GPU pipeline #if batch_idx % 2 == 0: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() finally: # Explicitly shut down DataLoader workers so they don't linger into the # next condition and exhaust shared memory / semaphore limits. del loader gc.collect() def load_model_from_config(config_data): """ Load the model from checkpoint based on configuration data. Parameters ---------- config_data : dict Configuration dictionary from YAML file. Returns ------- torch.nn.Module Loaded model ready for inference. """ checkpoint_path = config_data['data']['checkpoint_path'] print("Loading ViTLocal model...") # Use GPU(s) if available if torch.cuda.is_available(): load_device = torch.device('cuda:0') print(f"Using GPU(s) for inference") else: load_device = torch.device('cpu') print("Using CPU for inference") if ".ckpt" in checkpoint_path: # Lightning checkpoint format model = ViTLocal.load_from_checkpoint(checkpoint_path, map_location=load_device, weights_only=False) else: state = torch.load(checkpoint_path, map_location=load_device, weights_only=False) model = state['model'] model = model.to(load_device) model.eval() raw_multi_gpu = config_data.get('multi_gpu', False) if config_data else False use_multi_gpu = (str(raw_multi_gpu).lower() == 'true' and torch.cuda.is_available() and torch.cuda.device_count() > 1) if use_multi_gpu: n_gpus = torch.cuda.device_count() model = torch.nn.DataParallel(model) print(f"Using DataParallel across {n_gpus} GPUs — ensure batch_size >= {n_gpus}") return model def main(): """ Main function to execute solar flare model inference pipeline. Steps ----- 1. Parse YAML configuration and resolve ${variable} placeholders. 2. Load pretrained model and dataset. 3. Run batched inference and optionally save attention/flux maps. 4. Save predictions and ground truth results to CSV. """ def resolve_config_variables(config_dict): """Recursively resolve ${variable} references within config.""" variables = {} for key, value in config_dict.items(): if isinstance(value, str) and not value.startswith('${'): variables[key] = value def substitute_value(value, variables): if isinstance(value, str): pattern = r'\$\{([^}]+)\}' for match in re.finditer(pattern, value): var_name = match.group(1) if var_name in variables: value = value.replace(f'${{{var_name}}}', variables[var_name]) return value def recursive_substitute(obj, variables): if isinstance(obj, dict): return {k: recursive_substitute(v, variables) for k, v in obj.items()} elif isinstance(obj, list): return [recursive_substitute(item, variables) for item in obj] else: return substitute_value(obj, variables) return recursive_substitute(config_dict, variables) parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='inference_config.yaml', required=True, help='Path to the inference configuration YAML file.') args = parser.parse_args() with open(args.config, 'r') as stream: config_data = yaml.load(stream, Loader=yaml.SafeLoader) config_data: dict = resolve_config_variables(config_data) model_params = config_data.get('model_params', {}) input_size = model_params.get('input_size', 512) patch_size = model_params.get('patch_size', 16) batch_size = model_params.get('batch_size', 10) no_weights = model_params.get('no_weights', False) no_flux = model_params.get('no_flux', False) print(f"Using parameters from config:") print(f" Input size: {input_size}\n Patch size: {patch_size}\n Batch size: {batch_size}\n Skip weights: {no_weights}\n Skip flux: {no_flux}") model = load_model_from_config(config_data) save_weights = not no_weights if no_weights: print("Skipping attention weight saving (no_weights=true).") print(" Note: This saves ~3GB per batch by not computing attention weights.") else: print("Will save attention weights during inference.") print("\n Memory note:") print(" - Attention weights from all layers use significant GPU memory") print(" - For ViT with 8 layers, 8 heads, 4096 patches: ~3GB+ per batch with attention!") print(" - If you get OOM errors, set no_weights=true to skip attention saving\n") save_flux = config_data and 'flux_path' in config_data and not no_flux if no_flux: print("Skipping flux contribution saving (no_flux=true).") elif config_data and 'flux_path' in config_data: print("Will save flux contributions during inference.") else: print("No flux path specified.") torch.backends.cudnn.benchmark = True print("Loading dataset...") # Check if running in prediction-only mode prediction_only = config_data.get('prediction_only', 'false').lower() == 'true' dataset = AIAGOESDataset(aia_dir=config_data['data']['aia_dir'], sxr_dir=config_data['data'].get('sxr_dir') if not prediction_only else None, wavelengths=config_data['wavelengths'], only_prediction=prediction_only) times = dataset.samples # Load SXR normalization only if not in prediction-only mode if prediction_only: print("Running in prediction-only mode - skipping SXR normalization loading") sxr_norm = None else: try: sxr_norm = np.load(config_data['data']['sxr_norm_path']) print("SXR normalization loaded successfully") except FileNotFoundError: print(f"Warning: SXR normalization file not found: {config_data['data']['sxr_norm_path']}") print("Continuing without normalization...") sxr_norm = None timestamp, predictions, ground = [], [], [] total_samples = len(times) print(f"Processing {total_samples} samples with batch size {batch_size}...") print("Running inference...") # Create progress bar pbar = tqdm( evaluate_model_on_dataset( model, dataset, batch_size, times, config_data, save_weights, input_size, patch_size, save_flux ), total=total_samples, desc="Inference", unit="sample", ncols=100 ) for prediction, sxr, weight, flux_data, idx in pbar: # ViTLocal models already return unnormalized predictions, so no need to unnormalize pred = prediction predictions.append(pred.item() if hasattr(pred, 'item') else float(pred)) # Only collect ground truth if not in prediction-only mode if not prediction_only: ground.append(sxr.item() if hasattr(sxr, 'item') else float(sxr)) else: ground.append(0.0) # Dummy value for prediction-only mode timestamp.append(str(times[idx])) # Update progress bar pbar.set_postfix({'sample': idx + 1, 'total': total_samples}) output_df = pd.DataFrame({'timestamp': timestamp, 'predictions': predictions, 'groundtruth': ground}) output_dir = Path(config_data['output_path']).parent output_dir.mkdir(parents=True, exist_ok=True) output_df.to_csv(config_data['output_path'], index=False) if prediction_only: print(f"Predictions saved to {config_data['output_path']} (prediction-only mode)") else: print(f"Predictions saved to {config_data['output_path']}") if __name__ == '__main__': main()