diff --git a/src/configs/config.py b/src/configs/config.py index a086a2c..bb466d7 100644 --- a/src/configs/config.py +++ b/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") diff --git a/src/dataloaders/build_dataloader.py b/src/dataloaders/build_dataloader.py index af229d3..37e4ac3 100644 --- a/src/dataloaders/build_dataloader.py +++ b/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 diff --git a/src/dataloaders/data_representation/base.py b/src/dataloaders/data_representation/base.py index da1b19d..ef4efe9 100644 --- a/src/dataloaders/data_representation/base.py +++ b/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() diff --git a/src/dataloaders/data_representation/rgb.py b/src/dataloaders/data_representation/rgb.py index 621cbbf..5eb5b12 100644 --- a/src/dataloaders/data_representation/rgb.py +++ b/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) diff --git a/src/dataloaders/data_representation/signal.py b/src/dataloaders/data_representation/signal.py index 6730544..6186c00 100644 --- a/src/dataloaders/data_representation/signal.py +++ b/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) diff --git a/src/dataloaders/data_representation/stacked_signal.py b/src/dataloaders/data_representation/stacked_signal.py index 9f078c3..56c92ba 100644 --- a/src/dataloaders/data_representation/stacked_signal.py +++ b/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) diff --git a/src/dataloaders/data_representation/symbolic.py b/src/dataloaders/data_representation/symbolic.py index 1e9b0cc..010cf39 100644 --- a/src/dataloaders/data_representation/symbolic.py +++ b/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 ### diff --git a/src/dataloaders/dataset_mixer.py b/src/dataloaders/dataset_mixer.py index 7860fa0..c488e47 100644 --- a/src/dataloaders/dataset_mixer.py +++ b/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): diff --git a/src/main_trainer.py b/src/main_trainer.py index 8b7636c..98bacac 100644 --- a/src/main_trainer.py +++ b/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) diff --git a/src/runners/validator.py b/src/runners/validator.py new file mode 100644 index 0000000..e96198e --- /dev/null +++ b/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} diff --git a/src/utils/gpu_manager.py b/src/utils/gpu_manager.py index 61ab8aa..396b7ac 100644 --- a/src/utils/gpu_manager.py +++ b/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