Loderunner cylex test eval

applications/harnesses/ch_DDP_loderunner_cylex/START_study.py

@@ -0,0 +1 @@
+../START_study.py

applications/harnesses/ch_DDP_loderunner_cylex/cp_files.txt

@@ -0,0 +1 @@
+train_LodeRunner_ddp_cylex.py

applications/harnesses/ch_DDP_loderunner_cylex/ddp_test.csv

@@ -0,0 +1,4 @@
+studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZE,NTRN_BATCH,NVAL_BATCH,ANCHOR_LR,NUM_CYCLES,MIN_FRACTION,TERMINAL_STEPS,WARMUP_STEPS,train_script
+# This is a longer epoch production run after tuning.
+# Single epoch with validation should be ~30 mins
+1,torch_se_gpu_120226,1,4,128,1,1,9,1,2,1,1000,500,5.0e-4,0.5,0.5,1000,500,train_LodeRunner_ddp_cylex.py

applications/harnesses/ch_DDP_loderunner_cylex/eval/avg_eval_csv.py

@@ -0,0 +1,40 @@
+"""Compute the statistics for MSE over evaluation CSV."""
+
+import argparse
+import pandas as pd
+
+
+def main():
+    """Compute and print."""
+    # Parser CLI CSV-filename
+    parser = argparse.ArgumentParser(
+        description="Compute summary statistics of evaluation CSV."
+        )
+    parser.add_argument(
+        '--csv',
+        type=str,
+        default='./testing_evaluation.csv',
+        help="Path to evaluation CSV."
+        )
+
+    args = parser.parse_args()
+
+    # Read the CSV file
+    df = pd.read_csv(args.csv, header=None)
+
+    # Extract the third column
+    col = df[2]
+
+    # Compute statistics of the third column
+    mean_MSE = col.mean()
+    p2_5, p50, p97_5 = col.quantile([0.025, 0.5, 0.975])
+
+    # Print results
+    print(f"Average MSE: {mean_MSE}")
+    print(f"2.5-percentile: {p2_5}")
+    print(f"50-percentile: {p50}")
+    print(f"97.5-percentile: {p97_5}")
+
+
+if __name__ == "__main__":
+    main()

applications/harnesses/ch_DDP_loderunner_cylex/eval/eval_LodeRunner_cylex.py

@@ -0,0 +1,235 @@
+"""Evaluate trained model on test set.
+
+Model is evaluated on `cycle_epoch` number of epochs with `test_batches` number of
+batches each of a set `batch_size`.
+
+"""
+
+import argparse
+import os
+import time
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from yoke.datasets.load_npz_dataset import TemporalDataSet
+from yoke.utils.checkpointing import load_model_and_optimizer
+from yoke.utils.training.epoch.loderunner import eval_loderunner_epoch
+from yoke.helpers import cli
+
+
+descr_str = (
+    "Single-GPU evaluation for a saved Yoke LodeRunner checkpoint on real dataset "
+    "batches. Mirrors train_LodeRunner_ddp.py structure but without DDP."
+)
+parser = argparse.ArgumentParser(
+    prog="LodeRunner Evaluation", description=descr_str, fromfile_prefix_chars="@"
+)
+
+# Reuse the same CLI arg groups as training so you can pass the same @argfiles.
+parser = cli.add_default_args(parser=parser)
+parser = cli.add_filepath_args(parser=parser)
+parser = cli.add_computing_args(parser=parser)
+parser = cli.add_training_args(parser=parser)
+
+# Keep the same default filelists as train_LodeRunner_ddp.py
+parser.set_defaults(
+    train_filelist="cx241203_prefixes_train_80pct_noBe_noVoid_truncated.txt",
+    validation_filelist="cx241203_prefixes_val_10pct_noBe_noVoid_truncated.txt",
+    test_filelist="cx241203_prefixes_test_10pct_noBe_noVoid_truncated.txt",
+)
+
+def main(args: argparse.Namespace) -> None:
+    """Main evaluation function."""
+    #############################################
+    # Process Inputs
+    #############################################
+    # Device (single GPU)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Paths
+    filelist_dir = args.FILELIST_DIR
+    filelist_path = os.path.join(filelist_dir, args.test_filelist)
+
+    # Dataloader params
+    batch_size = args.batch_size
+    num_workers = args.num_workers
+    test_batches = args.test_batches
+    batch_size = args.batch_size
+    cycle_epochs = args.cycle_epochs
+    test_rcrd_filename = args.test_rcrd_filename
+
+    model_args = {
+        "default_vars": [
+            "Rcoord",  # 4 kinematic variable fields
+            "Zcoord",
+            "Uvelocity",
+            "Wvelocity",
+            "density_Air",  # 39 thermodynamic variable fields
+            "energy_Air",
+            "pressure_Air",
+            "density_Al",
+            "energy_Al",
+            "pressure_Al",
+            "density_Be",
+            "energy_Be",
+            "pressure_Be",
+            "density_booster",
+            "energy_booster",
+            "pressure_booster",
+            "density_Cu",
+            "energy_Cu",
+            "pressure_Cu",
+            "density_U.DU",
+            "energy_U.DU",
+            "pressure_U.DU",
+            "density_maincharge",
+            "energy_maincharge",
+            "pressure_maincharge",
+            "density_N",
+            "energy_N",
+            "pressure_N",
+            "density_Sn",
+            "energy_Sn",
+            "pressure_Sn",
+            "density_Steel.alloySS304L",
+            "energy_Steel.alloySS304L",
+            "pressure_Steel.alloySS304L",
+            "density_Polymer.Sylgard",
+            "energy_Polymer.Sylgard",
+            "pressure_Polymer.Sylgard",
+            "density_Ta",
+            "energy_Ta",
+            "pressure_Ta",
+            "density_Void",
+            "energy_Void",
+            "pressure_Void",
+            "density_Water",
+            "energy_Water",
+            "pressure_Water",
+        ],
+        "image_size": (1120, 400),
+        "patch_size": (10, 5),
+        "embed_dim": 128,
+        "emb_factor": 2,
+        "num_heads": 8,
+        "block_structure": (1, 1, 9, 1),
+        "window_sizes": [(8, 8), (8, 8), (4, 4), (2, 2)],
+        "patch_merge_scales": [(2, 2), (2, 2), (2, 2)],
+    }
+
+    model = LodeRunner(**model_args)
+
+    #############################################
+    # Load Model Checkpoint
+    #############################################
+    available_models = {"LodeRunner": LodeRunner}
+
+    # NOTE: optimizer args are required by load_model_and_optimizer, even for eval.
+    model, _optimizer, starting_epoch = load_model_and_optimizer(
+        args.pretrained_model,
+        optimizer_class=torch.optim.AdamW,
+        optimizer_kwargs={
+            "lr": 1e-6,
+            "betas": (0.9, 0.999),
+            "eps": 1e-08,
+            "weight_decay": 0.01,
+        },
+        available_models=available_models,
+        device=device,
+    )
+    starting_epoch = 0
+    model.to(device)
+    model.eval()
+
+    # load_and_eval_YokePth.py prints these; keep similar behavior here
+    print(f"Loaded checkpoint: {args.pretrained_model}", flush=True)
+    print(f"Checkpoint starting_epoch: {starting_epoch}", flush=True)
+    if hasattr(model, "default_vars"):
+        print("Default LodeRunner fields:", model.default_vars, flush=True)
+    if hasattr(model, "image_size"):
+        print("LodeRunner image size:", model.image_size, flush=True)
+
+    #############################################
+    # Dataset / Dataloader (non-distributed)
+    #############################################
+    testing_dataset = TemporalDataSet(
+        args.NPZ_DIR,
+        args.CSV_FILEPATH,
+        file_prefix_list=filelist_path,
+        max_timeIDX_offset=2,
+        max_file_checks=10,
+        half_image=True,
+    )
+
+    from torch.utils.data.dataloader import default_collate
+
+    def collate_skip_none(batch):
+        batch = [b for b in batch if b is not None]
+        if len(batch) == 0:
+            return None
+        return default_collate(batch)
+
+    test_dataloader = DataLoader(
+        dataset=testing_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available(),
+        #drop_last=False,
+        collate_fn=collate_skip_none,
+        #prefetch_factor=2,
+    )
+
+    #############################################
+    # Loss + Evaluation Loop
+    #############################################
+    # Match train_LodeRunner_ddp.py: use per-element MSE so we can reduce ourselves
+    loss_fn = nn.MSELoss(reduction="none")
+
+
+    #############################################
+    # Testing Loop
+    #############################################
+    # Train Model
+    print("Testing Model . . .")
+    starting_epoch += 1
+    ending_epoch = starting_epoch + cycle_epochs
+
+    for epochIDX in range(starting_epoch, ending_epoch):
+        # Time each epoch and print to stdout
+        startTime = time.time()
+
+        # Testing epoch
+        # for cylex channel_map changes per sample. So pass None here
+        # & calculate channel_map later in the datastep function.
+        eval_loderunner_epoch(
+            testing_data=test_dataloader,
+            num_test_batches=test_batches,
+            model=model,
+            dataset='cylex',
+            channel_map=None,
+            loss_fn=loss_fn,
+            epochIDX=epochIDX,
+            test_rcrd_filename=test_rcrd_filename,
+            device=device,
+        )
+
+        # Time each epoch and print to stdout
+        endTime = time.time()
+
+        epoch_time = (endTime - startTime) / 60
+
+        # Print Summary Results
+        print(f"Completed epoch {epochIDX}...", flush=True)
+        print(f"Epoch time (minutes): {epoch_time:.2f}", flush=True)
+
+
+if __name__ == "__main__":
+    """Parse arguments and run main evaluation function."""
+
+    args = parser.parse_args()
+
+    main(args)

applications/harnesses/ch_DDP_loderunner_cylex/eval/eval_START.input

@@ -0,0 +1,22 @@
+--pretrain_checkpoint
+/net/sescratch1/exempt/artimis/soumide/projects/yoke_runs/cylex_full_run_lr5e-4_w_val/runs/study_001/study001_modelState_epoch0100.pth
+--FILELIST_DIR
+/usr/projects/artimis/mpmm/hickmank/github_yoke/applications/filelists/
+--NPZ_DIR
+/net/sescratch1/exempt/artimis/data/cx241203/
+--CSV_FILEPATH
+/net/sescratch1/exempt/artimis/mpmm/design_cx241203_MASTER.csv
+--test_filelist
+cx241203_prefixes_test_10pct_noBe_noVoid.txt
+--test_rcrd_filename
+./testing_evaluation.csv
+--batch_size
+1
+--num_workers
+2
+--total_epochs
+10
+--cycle_epochs
+10
+--test_batches
+1000

applications/harnesses/ch_DDP_loderunner_cylex/eval/eval_START.slurm

@@ -0,0 +1,79 @@
+#!/bin/bash
+
+# This is a setup for GPU training on Selene. Find out how much
+# memory per node, number of CPUs/node.
+
+# NOTE: Number of CPUs per GPU must be an even number since there are
+# 2 threads per core. If an odd number is requested the next higher
+# even number gets used. 
+
+# The following are one set of SBATCH options for the Selene GPU
+# partition. There are optional other constraints.
+
+#SBATCH --job-name=cylfull_eval
+#SBATCH --account=y26_artimis-fmod_g
+#SBATCH --partition=standard
+#SBATCH --time=8:00:00
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=4
+#SBATCH --gpus-per-node=4
+#SBATCH --cpus-per-task=8
+#SBATCH --mem-per-gpu=120G
+#SBATCH --output=eval_test.out
+#SBATCH --error=eval_test.err
+#SBATCH -vvv
+
+# Set the master node's address and port
+MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
+MASTER_PORT=$(shuf -i 1024-65535 -n 1)  # Choose a random port
+export MASTER_ADDR
+export MASTER_PORT
+
+# Enable shell debugging: Debugging selene slurm
+set -xv
+
+# Check available GPUs
+sinfo  -o "%P %.24G %N"
+#srun -vv --cpu-bind=verbose /usr/bin/echo $CUDA_AVAILABLE_DEVICES
+nvidia-smi
+
+# Specify NCCL communication
+export NCCL_SOCKET_IFNAME=ib0  # Check possible interfaces with `ip link show`
+
+# for multi-node training
+export NCCL_IB_HCA=mlx5_0
+export NCCL_IB_GID_INDEX=3  # sometimes 0 or 3 depending on subnet manager config
+
+# Debugging distributed data parallel
+# export NCCL_DEBUG=INFO
+# export NCCL_DEBUG_SUBSYS=INIT
+
+# Debugging selene slurm
+export SLURM_CPU_BIND=verbose
+
+# Load correct conda environment
+module load anaconda/3.12
+source activate
+conda activate torch_se_gpu_120226
+
+# Set number of threads per GPU
+export OMP_NUM_THREADS=10
+
+# Get start time
+export date00=`date`
+
+# Start the Code
+# Explicitly set TCP environment for the following...
+#srun -vv --cpu-bind=verbose python eval_LodeRunner.py @eval_START.input
+python eval_LodeRunner.py @eval_START.input
+
+# Start the Code
+#python eval_LodeRunner.py @eval_START.input
+
+# Get end time and print to stdout
+export date01=`date`
+
+echo "===================TIME STARTED==================="
+echo $date00
+echo "===================TIME FINISHED==================="
+echo $date01