From e5ecb714ba2d8c3107a673b193b548f9b5c7712e Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 20 Jun 2026 11:26:03 +0800 Subject: [PATCH] Refactor loss computation and model input handling for improved clarity and efficiency --- losses.py | 55 ++++++------- models.py | 75 ++++++++++-------- train_all_future.py | 39 +++++++-- train_next_step.py | 187 ++++++++++++++++++++++++++++---------------- 4 files changed, 220 insertions(+), 136 deletions(-) diff --git a/losses.py b/losses.py index 8daee71..51883a4 100644 --- a/losses.py +++ b/losses.py @@ -99,29 +99,33 @@ class Delphi2MLoss(nn.Module): } return total_loss + logits_valid = logits[valid_mask] + target_events_valid = target_events[valid_mask] + target_times_valid = target_times[valid_mask] + current_times_valid = current_times[valid_mask] + logits_safe = torch.nan_to_num( - logits, + logits_valid, nan=0.0, posinf=self.max_exp_input, neginf=-self.max_exp_input, ) loss_ce = F.cross_entropy( - logits_safe.reshape(-1, vocab_size)[valid_mask.reshape(-1)], - target_events.reshape(-1)[valid_mask.reshape(-1)], + logits_safe, + target_events_valid, reduction="mean", ) logits_for_lse = logits_safe if self.exclude_ignored_from_intensity: ignore_mask = _make_ignore_mask(vocab_size, self.ignored_tokens, logits.device) - logits_for_lse = logits_safe.clone() - logits_for_lse[:, :, ignore_mask] = float("-inf") + logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf")) log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1) log_lambda_total = -torch.log(torch.exp(-log_lambda_total) + self.t_min) - dt = torch.clamp(target_times - current_times, min=self.t_min) + dt = torch.clamp(target_times_valid - current_times_valid, min=self.t_min) log_dt_inv = -torch.log(dt + self.t_min) loss_dt = -( log_lambda_total @@ -129,7 +133,7 @@ class Delphi2MLoss(nn.Module): torch.clamp(log_lambda_total - log_dt_inv, max=self.max_exp_input) ) ) - loss_dt = loss_dt[valid_mask].mean() + loss_dt = loss_dt.mean() total_loss = self.ce_weight * loss_ce + self.time_weight * loss_dt if return_components: @@ -181,17 +185,9 @@ class UniqueTimeSetExponentialLoss(nn.Module): if readout_mask.shape != (bsz, seq_len): raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}") - logits_safe = torch.nan_to_num( - logits, - nan=0.0, - posinf=self.max_exp_input, - neginf=-self.max_exp_input, - ) ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device) - target_valid = target_multi_hot.float().clone() - target_valid[:, :, ignore_mask] = 0.0 - num_targets = target_valid.sum(dim=-1) + num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1) valid_mask = readout_mask.bool() & (num_targets > 0) if not valid_mask.any(): @@ -204,34 +200,35 @@ class UniqueTimeSetExponentialLoss(nn.Module): } return total_loss - log_lambda_targets = logits_safe * target_valid - log_lambda_targets = torch.where( - target_valid.bool(), - log_lambda_targets, - torch.zeros_like(log_lambda_targets), + logits_safe = torch.nan_to_num( + logits[valid_mask], + nan=0.0, + posinf=self.max_exp_input, + neginf=-self.max_exp_input, ) + target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype) + target_valid[:, ignore_mask] = 0.0 - observed_term = log_lambda_targets.sum(dim=-1) - penalty_scale = num_targets + observed_term = (logits_safe * target_valid).sum(dim=-1) + penalty_scale = target_valid.sum(dim=-1) logits_for_lse = logits_safe if self.exclude_ignored_from_intensity: - logits_for_lse = logits_safe.clone() - logits_for_lse[:, :, ignore_mask] = float("-inf") + logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf")) - dt_clamped = torch.clamp(target_dt_unique, min=self.t_min) + dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min) log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1) log_penalty = log_lambda_total + dt_clamped.log() penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input)) observed_loss = -observed_term penalty_loss = penalty_scale * penalty - total_loss = (observed_loss + penalty_loss)[valid_mask].mean() + total_loss = (observed_loss + penalty_loss).mean() if return_components: return total_loss, { - "observed": observed_loss[valid_mask].mean().detach(), - "penalty": penalty_loss[valid_mask].mean().detach(), + "observed": observed_loss.mean().detach(), + "penalty": penalty_loss.mean().detach(), "total": total_loss.detach(), } return total_loss diff --git a/models.py b/models.py index 6d409da..d8b6e4e 100644 --- a/models.py +++ b/models.py @@ -260,12 +260,27 @@ class DeepHealth(nn.Module): other_time: torch.Tensor, other_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size, _n_other, n_embd = h_other.shape - group_counts = [ - int(torch.unique(other_time[b, other_mask[b]], sorted=True).numel()) - for b in range(batch_size) - ] - max_groups = max(group_counts, default=0) + batch_size, n_other, n_embd = h_other.shape + if n_other == 0: + empty_h = h_other.new_zeros(batch_size, 0, n_embd) + empty_t = other_time.new_zeros(batch_size, 0) + empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device) + return empty_h, empty_t, empty_m + + masked_time = other_time.masked_fill(~other_mask, float("inf")) + _sorted_time_with_pad, order = masked_time.sort(dim=1) + sorted_time = other_time.gather(1, order) + sorted_mask = other_mask.gather(1, order) + sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd)) + + group_start = torch.zeros_like(sorted_mask) + group_start[:, 0] = sorted_mask[:, 0] + group_start[:, 1:] = sorted_mask[:, 1:] & ( + sorted_time[:, 1:] != sorted_time[:, :-1] + ) + group_id = group_start.long().cumsum(dim=1) - 1 + max_groups = int(group_start.sum(dim=1).max().item()) + pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd) pooled_time = other_time.new_zeros(batch_size, max_groups) pooled_mask = torch.zeros( @@ -277,35 +292,29 @@ class DeepHealth(nn.Module): if max_groups == 0: return pooled_h, pooled_time, pooled_mask - for b in range(batch_size): - valid_time = other_time[b, other_mask[b]] - if valid_time.numel() == 0: - continue - valid_h = h_other[b, other_mask[b]] - unique_time, inverse = torch.unique( - valid_time, - sorted=True, - return_inverse=True, + safe_group_id = group_id.clamp_min(0) + pooled_h.scatter_add_( + 1, + safe_group_id.unsqueeze(-1).expand_as(sorted_h), + sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype), + ) + if self.extra_pool_reduce == "mean": + counts = h_other.new_zeros(batch_size, max_groups, 1) + counts.scatter_add_( + 1, + safe_group_id.unsqueeze(-1), + sorted_mask.unsqueeze(-1).to(h_other.dtype), ) - n_groups = unique_time.numel() - group_h = valid_h.new_zeros(n_groups, n_embd) - group_h.scatter_add_( - 0, - inverse[:, None].expand(-1, n_embd), - valid_h, - ) - if self.extra_pool_reduce == "mean": - counts = valid_h.new_zeros(n_groups, 1) - counts.scatter_add_( - 0, - inverse[:, None], - torch.ones_like(valid_h[:, :1]), - ) - group_h = group_h / counts.clamp_min(1.0) + pooled_h = pooled_h / counts.clamp_min(1.0) - pooled_h[b, :n_groups] = group_h - pooled_time[b, :n_groups] = unique_time - pooled_mask[b, :n_groups] = True + pooled_time.scatter_add_( + 1, + safe_group_id, + sorted_time * group_start.to(sorted_time.dtype), + ) + group_count = group_start.sum(dim=1) + arange_groups = torch.arange(max_groups, device=h_other.device) + pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1) return pooled_h, pooled_time, pooled_mask def _forward_shared( diff --git a/train_all_future.py b/train_all_future.py index 167b964..281d69f 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -44,6 +44,19 @@ from train_util import ( ) +MODEL_INPUT_KEYS = ( + "event_seq", + "time_seq", + "sex", + "padding_mask", + "t_query", + "other_type", + "other_value", + "other_value_kind", + "other_time", +) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train DeepHealth with all-future supervision") @@ -87,6 +100,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--min_lr_ratio", type=float, default=0.1) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", type=str, default="cuda") + parser.add_argument("--progress_interval", type=int, default=20) args = parser.parse_args() if args.min_history_events < 1: @@ -169,7 +183,14 @@ def compute_all_future_loss( batch: Dict[str, torch.Tensor], device: torch.device, ) -> torch.Tensor: - batch = move_batch_to_device(batch, device) + required_keys = set(MODEL_INPUT_KEYS) + required_keys.update(("future_targets", "exposure")) + if args.dist_mode in {"weibull", "mixed"}: + required_keys.add("future_dt") + batch = move_batch_to_device( + {key: batch[key] for key in required_keys}, + device, + ) hidden = model( event_seq=batch["event_seq"], @@ -224,10 +245,11 @@ def run_epoch( is_train: bool, ) -> float: model.train(is_train) - total = 0.0 + total = torch.zeros((), device=device) n_batches = 0 skipped = 0 desc = "train" if is_train else "val" + progress_interval = max(1, int(args.progress_interval)) progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) for batch_idx, batch in enumerate(progress): @@ -242,10 +264,15 @@ def run_epoch( clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() - total += float(loss.detach().cpu()) + total = total + loss.detach() n_batches += 1 - avg = total / max(1, n_batches) - progress.set_postfix(loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{avg:.4f}", skipped=skipped) + if (batch_idx + 1) % progress_interval == 0: + avg = total / max(1, n_batches) + progress.set_postfix( + loss=f"{float(loss.detach().cpu()):.4f}", + avg=f"{float(avg.detach().cpu()):.4f}", + skipped=skipped, + ) except RuntimeError as exc: if "Loss is not finite" not in str(exc): raise @@ -254,7 +281,7 @@ def run_epoch( if skipped: logger.info(f"Skipped {skipped} batches due to non-finite loss") - return total / max(1, n_batches) if n_batches else float("inf") + return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf") def build_metadata( diff --git a/train_next_step.py b/train_next_step.py index c57cae4..8965bbd 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -42,6 +42,18 @@ from train_util import ( ) +MODEL_INPUT_KEYS = ( + "event_seq", + "time_seq", + "sex", + "padding_mask", + "other_type", + "other_value", + "other_value_kind", + "other_time", +) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train DeepHealth with next-token/point supervision") @@ -94,6 +106,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--min_lr_ratio", type=float, default=0.1) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", type=str, default="cuda") + parser.add_argument("--progress_interval", type=int, default=20) args = parser.parse_args() use_eid_split = all( @@ -181,92 +194,111 @@ def build_next_step_loss(args: argparse.Namespace): def build_augmented_next_step_targets( - batch: Dict[str, torch.Tensor], + batch_cpu: Dict[str, torch.Tensor], model_out: DeepHealthOutput, + include_uts_targets: bool, ) -> Dict[str, torch.Tensor]: hidden_len = model_out.hidden.size(1) event_len = int(model_out.event_len) extra_len = hidden_len - event_len - if extra_len <= 0: - return { - "target_event_seq": batch["target_event_seq"], - "target_time_seq": batch["target_time_seq"], - "readout_mask": batch["readout_mask"], - "target_dt_unique": batch["target_dt_unique"], - "target_multi_hot": batch["target_multi_hot"], - } - device = model_out.hidden.device - bsz, _seq_len, vocab_size = batch["target_multi_hot"].shape - extra_mask = model_out.padding_mask[:, event_len:] - extra_time = model_out.time_seq[:, event_len:] + non_blocking = device.type == "cuda" + if extra_len <= 0: + targets = { + "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking), + "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking), + "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking), + } + if include_uts_targets: + targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to( + device, non_blocking=non_blocking + ) + targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to( + device, non_blocking=non_blocking + ) + return targets + + bsz = batch_cpu["target_event_seq"].size(0) + vocab_size = ( + batch_cpu["target_multi_hot"].size(2) + if include_uts_targets + else None + ) + other_valid = batch_cpu["other_type"] > 0 + extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len) + extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool) + for b in range(bsz): + unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True) + n_time = min(int(unique_time.numel()), extra_len) + if n_time > 0: + extra_time[b, :n_time] = unique_time[:n_time] + extra_mask[b, :n_time] = True target_event_seq = torch.cat( [ - batch["target_event_seq"], + batch_cpu["target_event_seq"], torch.full( (bsz, extra_len), PAD_IDX, - dtype=batch["target_event_seq"].dtype, - device=device, + dtype=batch_cpu["target_event_seq"].dtype, ), ], dim=1, ) target_time_seq = torch.cat( [ - batch["target_time_seq"], + batch_cpu["target_time_seq"], torch.zeros( bsz, extra_len, - dtype=batch["target_time_seq"].dtype, - device=device, - ), - ], - dim=1, - ) - readout_mask = torch.cat([batch["readout_mask"], extra_mask], dim=1) - target_dt_unique = torch.cat( - [ - batch["target_dt_unique"], - torch.zeros( - bsz, - extra_len, - dtype=batch["target_dt_unique"].dtype, - device=device, - ), - ], - dim=1, - ) - target_multi_hot = torch.cat( - [ - batch["target_multi_hot"], - torch.zeros( - bsz, - extra_len, - vocab_size, - dtype=batch["target_multi_hot"].dtype, - device=device, + dtype=batch_cpu["target_time_seq"].dtype, ), ], dim=1, ) + readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1) + target_dt_unique = None + target_multi_hot = None + if include_uts_targets: + target_dt_unique = torch.cat( + [ + batch_cpu["target_dt_unique"], + torch.zeros( + bsz, + extra_len, + dtype=batch_cpu["target_dt_unique"].dtype, + ), + ], + dim=1, + ) + target_multi_hot = torch.cat( + [ + batch_cpu["target_multi_hot"], + torch.zeros( + bsz, + extra_len, + vocab_size, + dtype=batch_cpu["target_multi_hot"].dtype, + ), + ], + dim=1, + ) for b in range(bsz): - valid_event = batch["padding_mask"][b].bool() + valid_event = batch_cpu["padding_mask"][b].bool() if not valid_event.any(): continue n_event = int(valid_event.sum().item()) events = torch.cat( [ - batch["event_seq"][b, :n_event], - batch["target_event_seq"][b, n_event - 1:n_event], + batch_cpu["event_seq"][b, :n_event], + batch_cpu["target_event_seq"][b, n_event - 1:n_event], ] ) times = torch.cat( [ - batch["time_seq"][b, :n_event], - batch["target_time_seq"][b, n_event - 1:n_event], + batch_cpu["time_seq"][b, :n_event], + batch_cpu["target_time_seq"][b, n_event - 1:n_event], ] ) valid_full = events > PAD_IDX @@ -291,6 +323,9 @@ def build_augmented_next_step_targets( target_event_seq[b, pos] = next_event target_time_seq[b, pos] = next_time + if not include_uts_targets: + continue + same_next_time = times == next_time next_events = events[same_next_time] valid_next_events = next_events[ @@ -302,13 +337,15 @@ def build_augmented_next_step_targets( target_multi_hot[b, pos, valid_next_events] = True target_dt_unique[b, pos] = next_time - t - return { - "target_event_seq": target_event_seq, - "target_time_seq": target_time_seq, - "readout_mask": readout_mask, - "target_dt_unique": target_dt_unique, - "target_multi_hot": target_multi_hot, + targets = { + "target_event_seq": target_event_seq.to(device, non_blocking=non_blocking), + "target_time_seq": target_time_seq.to(device, non_blocking=non_blocking), + "readout_mask": readout_mask.to(device, non_blocking=non_blocking), } + if include_uts_targets: + targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking) + targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking) + return targets def compute_next_step_loss( @@ -319,7 +356,11 @@ def compute_next_step_loss( batch: Dict[str, torch.Tensor], device: torch.device, ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: - batch = move_batch_to_device(batch, device) + batch_cpu = batch + batch = move_batch_to_device( + {key: batch_cpu[key] for key in MODEL_INPUT_KEYS}, + device, + ) model_out = model( event_seq=batch["event_seq"], time_seq=batch["time_seq"], @@ -334,7 +375,11 @@ def compute_next_step_loss( ) if not isinstance(model_out, DeepHealthOutput): raise TypeError("DeepHealth return_output=True must return DeepHealthOutput") - targets = build_augmented_next_step_targets(batch=batch, model_out=model_out) + targets = build_augmented_next_step_targets( + batch_cpu=batch_cpu, + model_out=model_out, + include_uts_targets=args.target_mode == "uts", + ) readout_out = readout( hidden=model_out.hidden, time_seq=model_out.time_seq, @@ -380,11 +425,12 @@ def run_epoch( ) -> float: model.train(is_train) readout.train(is_train) - total = 0.0 + total = torch.zeros((), device=device) n_batches = 0 skipped = 0 - parts_sum: Dict[str, float] = {} + parts_sum: Dict[str, torch.Tensor] = {} desc = "train" if is_train else "val" + progress_interval = max(1, int(args.progress_interval)) progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) for batch_idx, batch in enumerate(progress): @@ -399,15 +445,20 @@ def run_epoch( clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() - total += float(loss.detach().cpu()) + total = total + loss.detach() n_batches += 1 for name, value in parts.items(): - parts_sum[name] = parts_sum.get(name, 0.0) + float(value.detach().cpu()) - avg = total / max(1, n_batches) - postfix = {"loss": f"{float(loss.detach().cpu()):.4f}", "avg": f"{avg:.4f}", "skipped": skipped} - for name, value in parts_sum.items(): - postfix[name] = f"{value / max(1, n_batches):.4f}" - progress.set_postfix(postfix) + parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach() + if (batch_idx + 1) % progress_interval == 0: + avg = total / max(1, n_batches) + postfix = { + "loss": f"{float(loss.detach().cpu()):.4f}", + "avg": f"{float(avg.detach().cpu()):.4f}", + "skipped": skipped, + } + for name, value in parts_sum.items(): + postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}" + progress.set_postfix(postfix) except RuntimeError as exc: if "Loss is not finite" not in str(exc): raise @@ -416,7 +467,7 @@ def run_epoch( if skipped: logger.info(f"Skipped {skipped} batches due to non-finite loss") - return total / max(1, n_batches) if n_batches else float("inf") + return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf") def build_metadata(