Simplify training

Author: deinal

README.md

@@ -46,23 +46,19 @@ python download_data.py -d era5 -s 1987-01-01 -e 2024-05-31
 
 Mediterranean reanalysis
 ```
-python prepare_states.py -d data/mediterranean/raw/reanalysis -o data/mediterranean/samples/train -n 6 -p rea_data -s 1987-01-01 -e 2021-11-30
-python prepare_states.py -d data/mediterranean/raw/reanalysis -o data/mediterranean/samples/val -n 6 -p rea_data -s 2021-11-01 -e 2021-12-31
-python prepare_states.py -d data/mediterranean/raw/reanalysis -o data/mediterranean/samples/test -n 16 -p rea_data -s 2022-01-01 -e 2022-07-31
+python prepare_states.py -d data/mediterranean/raw/reanalysis -o data/mediterranean/samples/train -n 6 -p rea_data -s 1987-01-01 -e 2021-12-31
 ```
 
 Mediterranean analysis
 ```
-python prepare_states.py -d data/mediterranean/raw/analysis -o data/mediterranean/samples/train -n 6 -p ana_data -s 2021-11-01 -e 2024-03-31
-python prepare_states.py -d data/mediterranean/raw/analysis -o data/mediterranean/samples/val -n 6 -p ana_data -s 2024-04-01 -e 2024-05-31
+python prepare_states.py -d data/mediterranean/raw/analysis -o data/mediterranean/samples/train -n 6 -p ana_data -s 2022-01-01 -e 2024-04-30
+python prepare_states.py -d data/mediterranean/raw/analysis -o data/mediterranean/samples/val -n 6 -p ana_data -s 2024-05-01 -e 2024-06-30
 ```
 
 ERA5
 ```
-python prepare_states.py -d data/mediterranean/raw/era5 -o data/mediterranean/samples/train -n 6 -p forcing -s 1987-01-01 -e 2024-03-31
-python prepare_states.py -d data/mediterranean/raw/era5 -o data/mediterranean/samples/val -n 6 -p forcing -s 2021-11-01 -e 2021-12-31
-python prepare_states.py -d data/mediterranean/raw/era5 -o data/mediterranean/samples/val -n 6 -p forcing -s 2024-04-01 -e 2024-05-31
-python prepare_states.py -d data/mediterranean/raw/era5 -o data/mediterranean/samples/test -n 16 -p forcing -s 2022-01-01 -e 2022-07-31
+python prepare_states.py -d data/mediterranean/raw/era5 -o data/mediterranean/samples/train -n 6 -p forcing -s 1987-01-01 -e 2024-04-30
+python prepare_states.py -d data/mediterranean/raw/era5 -o data/mediterranean/samples/val -n 6 -p forcing -s 2024-05-01 -e 2024-06-30
 ```
 
 Forecast data
@@ -111,13 +107,13 @@ wandb off
 SeaCast was trained on 4 nodes with 8 GPUs each:
 ```
 python train_model.py \
-  --data_subset reanalysis \
   --epochs 200 \
   --n_workers 4 \
   --batch_size 1 \
   --step_length 1 \
   --ar_steps 4 \
   --lr 0.001 \
+  --optimizer momo_adam \
   --scheduler cosine \
   --finetune_start 0.6 \
   --model hi_lam \
@@ -126,16 +122,15 @@ python train_model.py \
   --hidden_dim 128 \
   --n_nodes 4
 ```
-For finetuing update arguments `--data_subset analysis`, `--epochs 50` and `--lr 0.0001`.
 
 For a full list of possible training options, check `python train_model.py --help`.
 
 ## Evaluation
 
-SeaCast was evaluated on 1 GPU using `--eval test`, and by choosing the correct data subset + loading the appropriate model:
+SeaCast was evaluated on 1 GPU using `--eval test`:
 ```
 python train_model.py \
-  --data_subset analysis \
+  --forcing_prefix aifs_forcing \
   --n_workers 4 \
   --batch_size 1 \
   --step_length 1 \
@@ -188,7 +183,8 @@ data
 │       ├── parameter_std.pt                - Std.-dev. of state parameters (create_parameter_weights.py)
 │       ├── diff_mean.pt                    - Means of one-step differences (create_parameter_weights.py)
 │       ├── diff_std.pt                     - Std.-dev. of one-step differences (create_parameter_weights.py)
-│       ├── forcing_stats.pt                - Mean and std.-dev. of forcing (create_parameter_weights.py)
+│       ├── forcing_mean.pt                 - Means of atmospheric forcing (create_parameter_weights.py)
+│       ├── forcing_std.pt                  - Std.-dev. of atmospheric forcing (create_parameter_weights.py)
 │       └── parameter_weights.npy           - Loss weights for different state parameters (create_parameter_weights.py)
 ├── baltic
 ├── ...

create_parameter_weights.py

@@ -90,23 +90,29 @@ def main():
         )  # (N_batch, d_features,)
 
         # Atmospheric forcing at 1st windowed position
-        forcing_batch = forcing_batch[:, :, :, :6]
-        forcing_means.append(torch.mean(forcing_batch))  # (,)
-        forcing_squares.append(torch.mean(forcing_batch**2))  # (,)
+        forcing_batch = forcing_batch[
+            :, :, :, :4
+        ]  # (N_batch, N_t-2, N_grid, d_atm)
+        forcing_means.append(
+            torch.mean(forcing_batch, dim=(1, 2))
+        )  # (N_batch, d_atm)
+        forcing_squares.append(
+            torch.mean(forcing_batch**2, dim=(1, 2))
+        )  # (N_batch, d_atm)
 
     mean = torch.mean(torch.cat(means, dim=0), dim=0)  # (d_features)
     second_moment = torch.mean(torch.cat(squares, dim=0), dim=0)
     std = torch.sqrt(second_moment - mean**2)  # (d_features)
 
-    forcing_mean = torch.mean(torch.stack(forcing_means))  # (,)
-    forcing_second_moment = torch.mean(torch.stack(forcing_squares))  # (,)
-    forcing_std = torch.sqrt(forcing_second_moment - forcing_mean**2)  # (,)
-    forcing_stats = torch.stack((forcing_mean, forcing_std))
+    forcing_mean = torch.mean(torch.cat(forcing_means, dim=0), dim=0)  # (d_atm)
+    forcing_second_moment = torch.mean(torch.cat(forcing_squares, dim=0), dim=0)
+    forcing_std = torch.sqrt(forcing_second_moment - forcing_mean**2)  # (d_atm)
 
     print("Saving mean, std.-dev...")
     torch.save(mean, os.path.join(static_dir_path, "parameter_mean.pt"))
     torch.save(std, os.path.join(static_dir_path, "parameter_std.pt"))
-    torch.save(forcing_stats, os.path.join(static_dir_path, "forcing_stats.pt"))
+    torch.save(forcing_mean, os.path.join(static_dir_path, "forcing_mean.pt"))
+    torch.save(forcing_std, os.path.join(static_dir_path, "forcing_std.pt"))
 
     # Compute mean and std.-dev. of one-step differences across the dataset
     print("Computing mean and std.-dev. for one-step differences...")