refactor attention visualization; update callback to log attention maps and adjust model configuration

flaring/MEGS_AI_baseline/callback.py

@@ -11,6 +11,9 @@ import numpy as np
 from pytorch_lightning.callbacks import Callback
 from PIL import Image
 import matplotlib.patches as patches
+import matplotlib.cm as cm
+import matplotlib.colors as mcolors
+from scipy.ndimage import zoom
 
 # Custom Callback
 sdoaia94 = matplotlib.colormaps['sdoaia94']
@@ -148,29 +151,28 @@ class AttentionMapCallback(Callback):
             patch_size: Size of patches
         """
         # Convert image to numpy for plotting
+        # Convert image to numpy and transpose
         img_np = image.cpu().numpy()
-        # Transpose from [C, H, W] to [H, W, C]
-
-        # Normalize image for display
-        #img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min())
+        if len(img_np.shape) == 3 and img_np.shape[0] in [1, 3]:  # Check if channels first
+            img_np = np.transpose(img_np, (1, 2, 0))
 
 
-        # Get attention from the last layer (or you can average across layers)
+        # Get attention from the last layer
         last_layer_attention = attention_weights[-1]  # [B, num_heads, seq_len, seq_len]
 
         # Extract attention for this sample
         sample_attention = last_layer_attention[sample_idx]  # [num_heads, seq_len, seq_len]
 
-        # Average across heads (or you can visualize individual heads)
+        # Average across heads
         avg_attention = sample_attention.mean(dim=0)  # [seq_len, seq_len]
 
         # Get attention from CLS token to patches (exclude CLS->CLS)
         cls_attention = avg_attention[0, 1:].cpu()  # [num_patches]
 
-        # Calculate grid size
-        H, W = img_np.shape[:2]
+        # Calculate grid size - NOW USING CORRECT DIMENSIONS
+        H, W = img_np.shape[:2]  # Now this is correct after transpose
         grid_h, grid_w = H // patch_size, W // patch_size
-        #print(grid_h, grid_w)
+
         # Reshape attention to spatial grid
         attention_map = cls_attention.reshape(grid_h, grid_w)
 
@@ -178,46 +180,52 @@ class AttentionMapCallback(Callback):
         fig, axes = plt.subplots(1, 3, figsize=(15, 5))
 
         # Plot 1: Original image
-        axes[0].imshow((img_np[:, :,0]+1)/2)
+        # if img_np.shape[2] == 1:  # Grayscale
+        #     img_display = (img_np[:, :, 0] + 1) / 2
+        #     axes[0].imshow(img_display, cmap='gray')
+        # elif img_np.shape[2] == 3:  # RGB
+        #     # Normalize RGB image properly
+        #     img_display = (img_np + 1) / 2  # Assuming images are in [-1, 1] range
+        #     img_display = np.clip(img_display, 0, 1)  # Ensure valid range
+        #     axes[0].imshow(img_display)
+        # else:  # Multi-channel (6 channels in your case)
+        #     # Option 1: Display first channel as grayscale
+        #     img_display = (img_np[:, :, 0] + 1) / 2
+        #     axes[0].imshow(img_display, cmap='gray')
+
+            # Option 2: Create RGB composite from 3 channels (uncomment if preferred)
+        rgb_channels = [0, 2, 4]  # Select which channels to use for R, G, B
+        img_display = np.stack([(img_np[:, :, i] + 1) / 2 for i in rgb_channels], axis=2)
+        img_display = np.clip(img_display, 0, 1)
+        axes[0].imshow(img_display)
         axes[0].set_title(f'Original Image (Epoch {epoch})')
         axes[0].axis('off')
 
         # Plot 2: Attention heatmap
-        im = axes[1].imshow(attention_map.numpy(), cmap='hot', interpolation='nearest')
+        attention_np = np.log1p(attention_map.numpy())
+        # Resize attention map to match image size
+        attention_resized = zoom(attention_np, (H / grid_h, W / grid_w), order=1)
+
+        # Create colormap for attention - FIX: Use the scalar values, not RGB
+        im = axes[1].imshow(attention_resized, cmap='hot')
         axes[1].set_title(f'Attention Map (Sample {sample_idx})')
         axes[1].axis('off')
+        # FIXED: Create colorbar from the scalar image, not RGB
         plt.colorbar(im, ax=axes[1])
 
         # Plot 3: Overlay attention on image
-        axes[2].imshow((img_np[:, :,0]+1)/2)
-
-        # Overlay attention as colored patches
-        max_attention = attention_map.max().numpy()
-        for i in range(grid_h):
-            for j in range(grid_w):
-                attention_val = attention_map[i, j].item()
-                # Create a colored rectangle with alpha based on attention
-                rect = patches.Rectangle(
-                    (j * patch_size, i * patch_size),
-                    patch_size, patch_size,
-                    linewidth=0,
-                    facecolor='red',
-                    alpha=(attention_val/max_attention) * .9
-                )
-                axes[2].add_patch(rect)
+        #img_display_overlay = (img_np[:, :, 0] + 1) / 2
+        axes[2].imshow(img_display)
 
-        axes[2].set_title(f'Attention Overlay (Sample {sample_idx})')
+        # Overlay attention with proper alpha blending
+        axes[2].imshow(attention_resized, cmap='hot', alpha=0.5)
+        axes[2].set_title(f'Log-Scaled Attention Overlay (Sample {sample_idx})')
         axes[2].axis('off')
 
         plt.tight_layout()
-        return fig
 
-        # Save the plot
-        # import os
-        # os.makedirs(self.save_dir, exist_ok=True)
-        # plt.savefig(f'{self.save_dir}/attention_epoch_{epoch}_sample_{sample_idx}.png',
-        #             dpi=150, bbox_inches='tight')
-        # plt.close()
+        plt.tight_layout()
+        return fig
 
 
 class MultiHeadAttentionCallback(AttentionMapCallback):

flaring/MEGS_AI_baseline/config.yaml

@@ -5,8 +5,9 @@ base_checkpoint_dir: "/mnt/data/ML-Ready/mixed_data"    # Change this line for d
 
 # Model configuration
 selected_model: "ViT"  # Options: "cnn", "vit",
-
-model:
+batch_size:    64
+epochs:        100
+megsai:
   architecture:
     "cnn"
   seed:
@@ -17,20 +18,16 @@ model:
     "resnet"
   cnn_dp:
     0.5
-  epochs:
-    100
-  batch_size:
-    16
 
 vit:
-    embed_dim: 512
+    embed_dim: 256
     num_channels: 6  # AIA has 6 channels
     num_classes: 1  # Regression task, predicting SXR flux
-    patch_size: 16
-    num_patches: 1024
-    hidden_dim: 512
-    num_heads: 4
-    num_layers: 4
+    patch_size: 32
+    num_patches: 256
+    hidden_dim: 256
+    num_heads: 1
+    num_layers: 1
     dropout: 0.25
     lr: .0001
 
@@ -68,5 +65,5 @@ wandb:
     - aia
     - sxr
     - regression
-  wb_name: flaring-vit-lr-scheduler
+  wb_name: mixed-vit-lr-scheduler
   notes: Regression from AIA images (6 channels) to GOES SXR flux

flaring/MEGS_AI_baseline/train.py

@@ -79,8 +79,8 @@ sxr_norm = np.load(config_data['data']['sxr_norm_path'])
 
 n = 0
 
-torch.manual_seed(config_data['model']['seed'])
-np.random.seed(config_data['model']['seed'])
+torch.manual_seed(config_data['megsai']['seed'])
+np.random.seed(config_data['megsai']['seed'])
 
 # DataModule
 data_loader = AIA_GOESDataModule(
@@ -90,7 +90,7 @@ data_loader = AIA_GOESDataModule(
     sxr_train_dir=config_data['data']['sxr_dir']+"/train",
     sxr_val_dir=config_data['data']['sxr_dir']+"/val",
     sxr_test_dir=config_data['data']['sxr_dir']+"/test",
-    batch_size=config_data['model']['batch_size'],
+    batch_size=config_data['batch_size'],
     num_workers=os.cpu_count(),
     sxr_norm=sxr_norm,
 )
@@ -105,7 +105,7 @@ wandb_logger = WandbLogger(
     tags=config_data['wandb']['tags'],
     name=config_data['wandb']['wb_name'],
     notes=config_data['wandb']['notes'],
-    config=config_data['model']
+    config=config_data['megsai']
 )
 
 # Logging callback
@@ -155,7 +155,7 @@ class PTHCheckpointCallback(Callback):
 # Checkpoint callback
 checkpoint_callback = ModelCheckpoint(
     dirpath=config_data['data']['checkpoints_dir'],
-    monitor='valid_loss',
+    monitor='val_loss',
     mode='min',
     save_top_k=1,
     filename=f"{config_data['wandb']['wb_name']}-{{epoch:02d}}-{{valid_loss:.4f}}.pth"
@@ -187,12 +187,6 @@ elif config_data['selected_model'] == 'hybrid':
         lr=config_data['model']['lr'],
     )
 elif config_data['selected_model'] == 'ViT':
-    print("Using ViT")
-#     model = ViT(embed_dim=config_data['vit']['embed_dim'], hidden_dim=config_data['vit']['hidden_dim'],
-#                 num_channels=config_data['vit']['num_channels'],num_heads=config_data['vit']['num_heads'],
-#                 num_layers=config_data['vit']['num_layers'], num_classes=config_data['vit']['num_classes'],
-#                 patch_size=config_data['vit']['patch_size'], num_patches=config_data['vit']['num_patches'],
-#                 dropout=config_data['vit']['dropout'], lr=config_data['vit']['lr'])
     model = ViT(model_kwargs=config_data['vit'])
 else:
     raise NotImplementedError(f"Architecture {config_data['selected_model']} not supported.")
@@ -202,8 +196,8 @@ trainer = Trainer(
     default_root_dir=config_data['data']['checkpoints_dir'],
     accelerator="gpu" if torch.cuda.is_available() else "cpu",
     devices=1,
-    max_epochs=config_data['model']['epochs'],
-    callbacks=[attention, pth_callback],
+    max_epochs=config_data['epochs'],
+    callbacks=[attention, pth_callback,checkpoint_callback],
     logger=wandb_logger,
     log_every_n_steps=10
 )