refactor model configuration and callbacks; update data paths and loss functions

flaring/MEGS_AI_baseline/callback.py

@@ -28,7 +28,7 @@ class ImagePredictionLogger_SXR(Callback):
     def __init__(self, data_samples, sxr_norm):
         super().__init__()
         self.data_samples = data_samples
-        self.val_aia = data_samples[0][0]
+        self.val_aia = data_samples[0]
         self.val_sxr = data_samples[1]
         self.sxr_norm = sxr_norm
 
@@ -178,7 +178,7 @@ class AttentionMapCallback(Callback):
         fig, axes = plt.subplots(1, 3, figsize=(15, 5))
 
         # Plot 1: Original image
-        axes[0].imshow(img_np[:, :, :3])  # only first 3 channels if more than 3
+        axes[0].imshow((img_np[:, :,0]+1)/2)
         axes[0].set_title(f'Original Image (Epoch {epoch})')
         axes[0].axis('off')
 
@@ -189,7 +189,7 @@ class AttentionMapCallback(Callback):
         plt.colorbar(im, ax=axes[1])
 
         # Plot 3: Overlay attention on image
-        axes[2].imshow(img_np[:, :, :3])
+        axes[2].imshow((img_np[:, :,0]+1)/2)
 
         # Overlay attention as colored patches
         max_attention = attention_map.max().numpy()

flaring/MEGS_AI_baseline/config.yaml

@@ -1,7 +1,7 @@
 
 # Base directories - change these to switch datasets
-base_data_dir: "/mnt/data/ML-Ready/flares_event_dir"  # Change this line for different datasets
-base_checkpoint_dir: "/mnt/data/ML-Ready/flares_event_dir"    # Change this line for different datasets
+base_data_dir: "/mnt/data/ML-Ready/mixed_data"  # Change this line for different datasets
+base_checkpoint_dir: "/mnt/data/ML-Ready/mixed_data"    # Change this line for different datasets
 
 # Model configuration
 selected_model: "ViT"  # Options: "cnn", "vit",
@@ -20,20 +20,35 @@ model:
   epochs:
     100
   batch_size:
-    64
+    16
 
 vit:
     embed_dim: 512
     num_channels: 6  # AIA has 6 channels
     num_classes: 1  # Regression task, predicting SXR flux
     patch_size: 16
-    num_patches: 262144
+    num_patches: 1024
     hidden_dim: 512
     num_heads: 4
-    num_layers: 6
-    dropout: 0.1
+    num_layers: 4
+    dropout: 0.25
     lr: .0001
 
+
+#vit:
+#    embed_dim: 512
+#    num_channels: 6  # AIA has 6 channels
+#    num_classes: 1  # Regression task, predicting SXR flux
+#    patch_size: 8
+#    num_patches: 4096
+#    hidden_dim: 512
+#    num_heads: 2
+#    num_layers: 3
+#    dropout: 0.25
+#    lr: .0001
+
+
+
 # Data paths (automatically constructed from base directories)
 data:
   aia_dir:

flaring/MEGS_AI_baseline/models/base_model.py

@@ -16,7 +16,7 @@ class BaseModel(LightningModule):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
         scheduler = {
             'scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3),
-            'monitor': 'valid_loss',  # name of the metric to monitor
+            'monitor': 'val_loss',  # name of the metric to monitor
             'interval': 'epoch',
         }
         return {'optimizer': optimizer, 'lr_scheduler': scheduler}

flaring/MEGS_AI_baseline/models/vision_transformer_custom.py

@@ -19,20 +19,27 @@ class ViT(pl.LightningModule):
         filtered_kwargs.pop('lr', None)
         self.model = VisionTransformer(**filtered_kwargs)
 
-    def forward(self, x, return_attention=False):
+    def forward(self, x, return_attention=True):
         return self.model(x, return_attention=return_attention)
 
     def configure_optimizers(self):
-        optimizer = optim.AdamW(self.parameters(), lr=self.lr)
-        lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1)
-        return [optimizer], [lr_scheduler]
+        # optimizer = optim.AdamW(self.parameters(), lr=self.lr)
+        # lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1)
+        # return [optimizer], [lr_scheduler]
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
+        scheduler = {
+            'scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3),
+            'monitor': 'val_loss',  # name of the metric to monitor
+            'interval': 'epoch',
+        }
+        return {'optimizer': optimizer, 'lr_scheduler': scheduler}
 
     def _calculate_loss(self, batch, mode="train"):
         imgs, sxr = batch
         preds = self.model(imgs)
 
         # Change loss function for regression
-        loss = F.huber_loss(torch.squeeze(preds), sxr)  # or F.l1_loss() or F.huber_loss()
+        loss = F.mse_loss(torch.squeeze(preds), sxr)  # or F.l1_loss() or F.huber_loss()
 
         # Change accuracy to a regression metric
         mae = F.l1_loss(torch.squeeze(preds), sxr)  # Mean Absolute Error

flaring/MEGS_AI_baseline/train.py

@@ -203,7 +203,7 @@ trainer = Trainer(
     accelerator="gpu" if torch.cuda.is_available() else "cpu",
     devices=1,
     max_epochs=config_data['model']['epochs'],
-    callbacks=[sxr_plot_callback, attention, pth_callback],
+    callbacks=[attention, pth_callback],
     logger=wandb_logger,
     log_every_n_steps=10
 )