Refactor loss computation and model input handling for improved clarity and efficiency

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(