Feat/val split

src/configs/config.py

@@ -66,7 +66,7 @@ def get_args(mode: Mode) -> argparse.Namespace:
         parser.add_argument("--epochs", type=int, default=1, help="Number of epochs")
         parser.add_argument("--weight_decay", type=float, default=1e-2, help="Weight decay")
         parser.add_argument("--patience", type=int, default=5, help="Patience for early stopping")
-        parser.add_argument("--patience_delta", type=float, default=0.1, help="Delta for early stopping")
+        parser.add_argument("--patience_delta", type=float, default=0.0001, help="Delta for early stopping")
         parser.add_argument("--early_stopping", action="store_true", default=False, help="Enable early stopping")
         parser.add_argument("--beta1", type=float, default=0.9, help="Beta1 for optimizer")
         parser.add_argument("--beta2", type=float, default=0.99, help="Beta2 for optimizer")
@@ -82,6 +82,8 @@ def get_args(mode: Mode) -> argparse.Namespace:
         parser.add_argument("--grad_clip", type=float, default=0.0, help="Max gradient norm for clipping (0 to disable)")
         parser.add_argument("--scale_wd", type=str, default="none", choices=["none", "inv_sqrt", "inv_linear"])
         parser.add_argument("--resume_ckpt", type=str, default=None, help="Full training resume: restores model, optimizer, and LR schedule state")
+        parser.add_argument("--val_split", type=float, default=None,
+                            help="Hold out examples for validation (<1 = fraction of train, >=1 = absolute count). Default: none.")
 
         # RL (train_phase=rl) — agnostic policy-gradient pipeline
         parser.add_argument("--rl_algo", type=str, default="sapo", choices=["sapo"], help="RL policy-loss algorithm")

src/dataloaders/build_dataloader.py

@@ -17,21 +17,23 @@ def __init__(
     ):
         self.args = args
         self.dataset_mixer = DatasetMixer(self.args)
+        self.val_dataloader = None
 
     def build_dataloader(
         self,
     ):
-        torch_dataset = self.dataset_mixer.build_torch_dataset()
-        torch_data_loader = self.build_torch_dataloader(torch_dataset)
-        return torch_data_loader
+        train_dataset, val_dataset = self.dataset_mixer.build_torch_dataset()
+        if val_dataset is not None:
+            self.val_dataloader = self.build_torch_dataloader(val_dataset, is_val=True)
+        return self.build_torch_dataloader(train_dataset)
 
-    def build_torch_dataloader(self, torch_dataset):
-        sampler = self.get_torch_dataloader_sampler(torch_dataset)
+    def build_torch_dataloader(self, torch_dataset, is_val=False):
+        sampler = self.get_torch_dataloader_sampler(torch_dataset, shuffle=not is_val)
         if "train" in self.args.mode:
             torch_data_loader = DataLoader(
                 torch_dataset,
                 batch_size=self.args.batch_size,
-                shuffle=(sampler is None),
+                shuffle=(sampler is None and not is_val),
                 num_workers=self.args.num_workers,
                 sampler=sampler,
                 pin_memory=torch.cuda.is_available(),
@@ -52,10 +54,11 @@ def build_torch_dataloader(self, torch_dataset):
     def get_torch_dataloader_sampler(
         self,
         torch_dataset,
+        shuffle=True,
     ):
         if self.args.distributed:
             sampler = DistributedSampler(torch_dataset, num_replicas=get_world_size(),
-                                         rank=get_rank(), seed=self.args.seed, shuffle=True)
+                                         rank=get_rank(), seed=self.args.seed, shuffle=shuffle)
         else:
             sampler = None
         return sampler

src/dataloaders/data_representation/base.py

@@ -23,6 +23,7 @@ class Base(Dataset):
     def __init__(self, data, args):
         self.data = data
         self.args = args
+        self.is_train = True
         self.fm = DirFileManager()
         if self.args.llm:
             self.chat_template = self.make_chat_template()

src/dataloaders/data_representation/rgb.py

@@ -17,7 +17,7 @@ def __init__(self, data, llm_tokenizer_components,
         self.llm_tokenizer = llm_tokenizer_components["llm_tokenizer"] if llm_tokenizer_components else None
         self.encoder_tokenizer = encoder_tokenizer_components["encoder_tokenizer"]
         self.viz = VizManager()
-        if self.args.augment_rgb:
+        if self.is_train and self.args.augment_rgb:
             self.aug = T.Compose([T.RandomApply([T.ColorJitter(brightness=0.2)], p=0.5),
                                 T.RandomApply([T.RandomRotation(5)], p=0.5),
                                 T.RandomApply([T.GaussianBlur(kernel_size=5, sigma=(0.0, 1.5))], p=0.5),
@@ -34,7 +34,7 @@ def __getitem__(self, index):
         else:
             ecg_np_file = self.fm.open_npy(instance["ecg_path"])
             ecg_signal = ecg_np_file["ecg"][self.args.leads]
-            if self.args.augment_ecg:
+            if self.is_train and self.args.augment_ecg:
                 ecg_signal = self.augment_ecg(ecg_signal)
 
         if self.args.dev and is_main():
@@ -102,7 +102,7 @@ def prepare_eval_inference_set(
     ### SIGNAL TO IMAGE FUNCTIONS ###
     def signal_to_image(self, ecg_signal: np.array):
         image = self.viz.get_plot_as_image(ecg_signal, 250, lead_names = self.lead_names)  # 250 Hz
-        if self.args.augment_rgb and random.random() < 0.6:
+        if self.is_train and self.args.augment_rgb and random.random() < 0.6:
             return self.augment_image(image)
         return Image.fromarray(image)
 

src/dataloaders/data_representation/signal.py

@@ -19,7 +19,7 @@ def __getitem__(self, index):
         else:
             ecg_np_file = self.fm.open_npy(instance["ecg_path"])
             ecg_signal = ecg_np_file["ecg"][self.args.leads]
-            if self.args.augment_ecg:
+            if self.is_train and self.args.augment_ecg:
                 ecg_signal = self.augment_ecg(ecg_signal)
 
         ecg_signal, _ = self.normalize(ecg_signal)

src/dataloaders/data_representation/stacked_signal.py

@@ -21,7 +21,7 @@ def __getitem__(self, index):
         else:
             ecg_np_file = self.fm.open_npy(instance["ecg_path"])
             ecg_signal = ecg_np_file["ecg"][self.args.leads]
-            if self.args.augment_ecg:
+            if self.is_train and self.args.augment_ecg:
                 ecg_signal = self.augment_ecg(ecg_signal)
         ecg_stacked_signal = self.signal_to_stacked_signal(ecg_signal)
 

src/dataloaders/data_representation/symbolic.py

@@ -21,7 +21,7 @@ def __getitem__(self, index):
         else:
             ecg_np_file = self.fm.open_npy(instance["ecg_path"])
             ecg_signal = ecg_np_file["ecg"][self.args.leads]
-            if self.args.augment_ecg:
+            if self.is_train and self.args.augment_ecg:
                 ecg_signal = self.augment_ecg(ecg_signal)
 
         ### PREPARE ECG INPUT ###

src/dataloaders/dataset_mixer.py

@@ -1,5 +1,6 @@
 from datasets import load_dataset
 import json
+import random
 from transformers import AutoTokenizer, AutoProcessor
 
 from utils.dir_file_manager import DirFileManager
@@ -27,9 +28,32 @@ def build_torch_dataset(self, ):
             print(f"Using {self.args.data_representation} representation")
         encoder_tokenizer_components = self.build_encoder_tokenizer()
         llm_tokenizer_components = self.build_llm_tokenizer()
-        torch_dataset = self.build_data_representation(data, llm_tokenizer_components,
+        train_data, val_data = self.split_train_val(data)
+        train_dataset = self.build_data_representation(train_data, llm_tokenizer_components,
                                                        encoder_tokenizer_components)
-        return torch_dataset
+        val_dataset = None
+        if val_data is not None:
+            val_dataset = self.build_data_representation(val_data, llm_tokenizer_components,
+                                                         encoder_tokenizer_components)
+            val_dataset.is_train = False
+        return train_dataset, val_dataset
+
+    def split_train_val(self, data):
+        val_split = getattr(self.args, "val_split", None)
+        if not val_split or "train" not in self.args.mode or getattr(self.args, "train_phase", "sft") == "rl":
+            return data, None
+        n_total = len(data)
+        n_val = int(n_total * val_split) if val_split < 1 else int(val_split)
+        n_val = max(0, min(n_val, n_total))
+        if n_val == 0:
+            return data, None
+        indices = list(range(n_total))
+        random.Random(self.args.seed).shuffle(indices)
+        val_data = [data[i] for i in indices[:n_val]]
+        train_data = [data[i] for i in indices[n_val:]]
+        if is_main():
+            print(f"Validation split: {len(train_data)} train / {len(val_data)} val (val_split={val_split})")
+        return train_data, val_data
 
     def build_data_representation(self, data, llm_tokenizer_components,
                                   encoder_tokenizer_components):

src/main_trainer.py

@@ -9,6 +9,7 @@
 from elms.build_elm import BuildELM
 
 from runners.trainer import run_train
+from runners.validator import run_validation
 from runners.rl_trainer import run_rl_train
 
 from utils.checkpoint_manager import CheckpointManager
@@ -51,6 +52,7 @@ def main():
         set_seed(args.seed)
         build_dataloader = BuildDataLoader(args)
         dataloader = build_dataloader.build_dataloader()
+        val_dataloader = build_dataloader.val_dataloader
         args.max_steps = math.ceil(len(dataloader) / args.grad_accum_steps) * args.epochs
         build_elm = BuildELM(args)
         elm_components = build_elm.build_elm(dataloader.dataset.llm_tokenizer)
@@ -69,11 +71,13 @@ def main():
         runner = run_rl_train if getattr(args, "train_phase", "sft") == "rl" else run_train
         for epoch in range(start_epoch, args.epochs):
             train_result = runner(elm, optimizer, dataloader, epoch, args, checkpoint_manager)
+            val_result = run_validation(elm, val_dataloader, epoch, args) if val_dataloader is not None else None
+            monitor_loss = val_result["average_loss"] if val_result is not None else train_result["average_loss"]
             should_stop = False
             if checkpoint_manager and is_main():
-                if checkpoint_manager.save_epoch(train_result["average_loss"]):
+                if checkpoint_manager.save_epoch(monitor_loss):
                     checkpoint_manager.save_checkpoint(elm, optimizer, epoch, -1, is_best=True, prefix="epoch_")
-                if args.early_stopping and checkpoint_manager.stop_early():
+                if args.early_stopping and val_result is not None and checkpoint_manager.stop_early():
                     print(f"Early stopping at epoch {epoch}")
                     should_stop = True
             should_stop = broadcast_value(should_stop, src=0)

src/runners/validator.py

@@ -0,0 +1,49 @@
+import torch
+from tqdm import tqdm
+import wandb
+
+from utils.gpu_manager import is_main, train_dev_break, batch_to_device, all_reduce_sum
+
+
+@torch.no_grad()
+def run_validation(
+    nn,
+    dataloader,
+    epoch,
+    args,
+):
+    if getattr(args, "distributed", False) and hasattr(getattr(dataloader, "sampler", None), "set_epoch"):
+        dataloader.sampler.set_epoch(epoch)
+
+    show_progress = is_main()
+    total_loss = 0.0
+    total_steps = 0
+    progress = tqdm(
+        dataloader,
+        desc=f"Validating LLM: {args.llm} ENCODER: {args.encoder};Epoch: {epoch}",
+        disable=not show_progress,
+        leave=False,
+    )
+
+    device = next(nn.parameters()).device
+
+    nn.eval()
+    for step, batch in enumerate(progress):
+        batch = {k: batch_to_device(v, device) for k, v in batch.items()}
+
+        out = nn(**batch)
+        raw_loss = out.loss
+
+        total_loss += raw_loss.item()
+        total_steps += 1
+
+        if train_dev_break(getattr(args, "dev", False), batch, raw_loss.item()):
+            break
+    nn.train()
+
+    total_loss = all_reduce_sum(total_loss)
+    total_steps = all_reduce_sum(total_steps)
+    average_loss = total_loss / total_steps if total_steps > 0 else float("inf")
+    if getattr(args, "wandb", False) and is_main():
+        wandb.log({"val/loss": average_loss, "epoch": epoch})
+    return {"average_loss": average_loss, "total_steps": total_steps}

src/utils/gpu_manager.py

@@ -55,6 +55,15 @@ def broadcast_value(val, src: int = 0):
     return obj[0]
 
 
+def all_reduce_sum(value: float) -> float:
+    if not (dist.is_available() and dist.is_initialized()):
+        return value
+    device = torch.device(f"cuda:{get_local_rank()}") if torch.cuda.is_available() else torch.device("cpu")
+    t = torch.tensor([value], dtype=torch.float64, device=device)
+    dist.all_reduce(t, op=dist.ReduceOp.SUM)
+    return t.item()
+
+
 def train_dev_break(enabled: bool, batch: dict, loss_value: float) -> bool:
     if not enabled:
         return False