update model configuration: adjust batch size, patch size, dropout rate, and checkpoint paths

.idea/workspace.xml

@@ -4,8 +4,14 @@
     <option name="autoReloadType" value="SELECTIVE" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="e883d810-3b4e-470d-b72e-9ed3eee941a0" name="Changes" comment="fixed some comments">
+    <list default="true" id="e883d810-3b4e-470d-b72e-9ed3eee941a0" name="Changes" comment="weight decay">
+      <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/flaring/forecasting/inference/inference_config.yaml" beforeDir="false" afterPath="$PROJECT_DIR$/flaring/forecasting/inference/inference_config.yaml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/flaring/forecasting/inference/inference_on_patch.py" beforeDir="false" afterPath="$PROJECT_DIR$/flaring/forecasting/inference/inference_on_patch.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/flaring/forecasting/inference/inference_on_patch_config.yaml" beforeDir="false" afterPath="$PROJECT_DIR$/flaring/forecasting/inference/inference_on_patch_config.yaml" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/flaring/forecasting/models/vit_patch_model.py" beforeDir="false" afterPath="$PROJECT_DIR$/flaring/forecasting/models/vit_patch_model.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/flaring/forecasting/training/config.yaml" beforeDir="false" afterPath="$PROJECT_DIR$/flaring/forecasting/training/config.yaml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/flaring/forecasting/training/train.py" beforeDir="false" afterPath="$PROJECT_DIR$/flaring/forecasting/training/train.py" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -143,7 +149,9 @@
       <workItem from="1755801805292" duration="58749000" />
       <workItem from="1756171392559" duration="40222000" />
       <workItem from="1756308588067" duration="10459000" />
-      <workItem from="1756386184391" duration="57683000" />
+      <workItem from="1756386184391" duration="64689000" />
+      <workItem from="1756818485332" duration="17000" />
+      <workItem from="1756821188374" duration="7721000" />
     </task>
     <task id="LOCAL-00004" summary="update train code to save model as .pth, creating plotting jupyter notebook, updated inference code to accept checkpoints">
       <option name="closed" value="true" />
@@ -489,7 +497,15 @@
       <option name="project" value="LOCAL" />
       <updated>1756678669786</updated>
     </task>
-    <option name="localTasksCounter" value="47" />
+    <task id="LOCAL-00047" summary="weight decay">
+      <option name="closed" value="true" />
+      <created>1756678864389</created>
+      <option name="number" value="00047" />
+      <option name="presentableId" value="LOCAL-00047" />
+      <option name="project" value="LOCAL" />
+      <updated>1756678864389</updated>
+    </task>
+    <option name="localTasksCounter" value="48" />
     <servers />
   </component>
   <component name="TypeScriptGeneratedFilesManager">
@@ -507,7 +523,6 @@
     </option>
   </component>
   <component name="VcsManagerConfiguration">
-    <MESSAGE value="add visualization notebook and refactor model configuration; update batch size, patch size, and dropout rate; implement dynamic loss for SXR regression" />
     <MESSAGE value="update model configuration and training scripts; adjust learning rate, epochs, and data paths; refactor trainer initialization and logging" />
     <MESSAGE value="update model configuration and training scripts; change selected model to 'ViT Custom', increase epochs to 250, and adjust output paths; refactor model loading in inference to support hybrid models and attention weight saving" />
     <MESSAGE value="added evaluation script" />
@@ -532,7 +547,8 @@
     <MESSAGE value="implement wavelength dropout in training, just as a test... it's commented out" />
     <MESSAGE value="added support for a patch based model" />
     <MESSAGE value="fixed some comments" />
-    <option name="LAST_COMMIT_MESSAGE" value="fixed some comments" />
+    <MESSAGE value="weight decay" />
+    <option name="LAST_COMMIT_MESSAGE" value="weight decay" />
   </component>
   <component name="XDebuggerManager">
     <breakpoint-manager>

flaring/forecasting/inference/inference_config.yaml

@@ -1,7 +1,7 @@
 # Base directories - change these to switch datasets
 base_data_dir: "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/"  # Change this line for different datasets
-output_path: "${base_data_dir}/output//mnt/data/ML-Ready-mixed/ML-Ready-mixed/output/NEURIPS.csv"
-weight_path: "${base_data_dir}/NEURIPS_Weight/"
+output_path: "${base_data_dir}/output//mnt/data/ML-Ready-mixed/ML-Ready-mixed/output/patch.csv"
+weight_path: "${base_data_dir}/patch_Weight/"
 mc:
   active: "false"
   runs: 5
@@ -18,7 +18,7 @@ SolO_data:
     sxr_dir: "${base_data_dir}/SXR"
     sxr_norm_path: "${base_data_dir}/SolO/SXR/normalized_sxr.npy"
 model: "cnn"  # Options: "cnn", "vit", "ViT Custom"
-wavelengths: [94,131, 171, 193, 211, 304]  # AIA wavelengths in Angstroms
+wavelengths: [94, 131, 171, 193, 211, 304]  # AIA wavelengths in Angstroms
 data:
   aia_dir:
     "${base_data_dir}/AIA"
@@ -27,6 +27,6 @@ data:
   sxr_norm_path:
     "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/SXR/normalized_sxr.npy"
   checkpoint_path:
-    "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/new-checkpoint/megsai-added-data-final-20250807_065956.pth"
+    "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/new-checkpoint/vit-16-higher-weight-lower-decay-epoch=119-val_total_loss=0.0419.ckpt"
 
 

flaring/forecasting/inference/inference_on_patch.py

@@ -237,7 +237,7 @@ def main():
     parser.add_argument('-config', type=str, default='config.yaml', required=True, help='Path to config YAML.')
     parser.add_argument('-input_size', type=int, default=512, help='Input size for the model')
     parser.add_argument('-patch_size', type=int, default=16, help='Patch size for the model')
-    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for inference')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size for inference')
     parser.add_argument('--no_weights', action='store_true', help='Skip saving attention weights to speed up')
     args = parser.parse_args()
 

flaring/forecasting/inference/inference_on_patch_config.yaml

@@ -27,5 +27,5 @@ data:
   sxr_norm_path:
     "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/SXR/normalized_sxr.npy"
   checkpoint_path:
-    "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/new-checkpoint/vit-16-patch-final-final-20250831_113151.pth"
+    "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/new-checkpoint/vit-16-higher-weight-lower-decay-epoch=119-val_total_loss=0.0419.ckpt"
 

flaring/forecasting/models/vit_patch_model.py

@@ -289,7 +289,7 @@ def img_to_patch(x, patch_size, flatten_channels=True):
 
 
 class SXRRegressionDynamicLoss:
-    def __init__(self, window_size=1000):
+    def __init__(self, window_size=1500):
         self.c_threshold = 1e-6
         self.m_threshold = 1e-5
         self.x_threshold = 1e-4
@@ -302,14 +302,14 @@ class SXRRegressionDynamicLoss:
 
         self.base_weights = {
             'quiet': 1.0,
-            'c_class': 5.0,
+            'c_class': 8.0,
             'm_class': 10.0,
             'x_class': 20.0
         }
 
     def calculate_loss(self, preds_squeezed, sxr, sxr_un, preds_squeezed_un):
         base_loss = F.huber_loss(preds_squeezed, sxr, delta=1.0, reduction='none')
-        weights = self._get_adaptive_weights(sxr_un, preds_squeezed_un, base_loss)
+        weights = self._get_adaptive_weights(sxr_un)
         self._update_tracking(sxr_un, preds_squeezed_un, base_loss)
         weighted_loss = base_loss * weights
         loss = weighted_loss.mean()

flaring/forecasting/training/config.yaml

@@ -5,9 +5,9 @@ base_checkpoint_dir: "/mnt/data/ML-Ready-mixed/ML-Ready-mixed"    # Change this
 wavelengths: [94, 131, 171, 193, 211, 304]  # AIA wavelengths in Angstroms
 # Model configuration
 selected_model: "ViT Patch"  # Options: "cnn", "vit",
-batch_size:    80
+batch_size:    4
 epochs:        500
-oversample: true
+oversample: false
 balance_strategy: "upsample_minority"
 
 megsai:
@@ -21,12 +21,12 @@ vit_custom:
     embed_dim: 512
     num_channels: 6
     num_classes: 1
-    patch_size: 16
-    num_patches: 1024
+    patch_size: 8
+    num_patches: 4096
     hidden_dim: 512
     num_heads: 8
     num_layers: 6
-    dropout: 0.2
+    dropout: 0.1
     lr: 0.0001
 
 # Data paths (automatically constructed from base directories)
@@ -48,5 +48,5 @@ wandb:
     - aia
     - sxr
     - regression
-  wb_name: vit-16-patch-final-oversample-higher-weights
+  wb_name: vit-8-patch
   notes: Regression from AIA images (6 channels) to GOES SXR flux

flaring/forecasting/training/train.py

@@ -166,7 +166,7 @@ checkpoint_callback = ModelCheckpoint(
     dirpath=config_data['data']['checkpoints_dir'],
     monitor='val_total_loss',
     mode='min',
-    save_top_k=4,
+    save_top_k=10,
     filename=f"{config_data['wandb']['wb_name']}-{{epoch:02d}}-{{val_total_loss:.4f}}"
 )