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

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(