Expose TimesNet internal dimension

README.md

@@ -89,7 +89,8 @@ The end-to-end next-step trainer supports the same DDP launch pattern:
 ```bash
 torchrun --standalone --nproc_per_node=4 train_next_step.py \
   --exposure_cache_dir ukb_exposure_cache \
-  --batch_size 128
+  --batch_size 128 \
+  --d_model 64
 ```
 Training-channel statistics are cached at
 `<exposure_cache_dir>/train_channel_stats.npz`; use

evaluate_auc.py

@@ -324,7 +324,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
         dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
         use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)),
-        exposure_d_model=cfg_get(args, cfg, "exposure_d_model", 64),
+        exposure_d_model=cfg_get(args, cfg, "d_model", 64),
         exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)),
         exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 2)),
         exposure_n_backbone_blocks=int(cfg_get(args, cfg, "exposure_n_backbone_blocks", 1)),

train_next_step.py

@@ -225,7 +225,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--dropout", type=float, default=0.0)
     parser.add_argument("--exposure_cache_dir", type=str, default=None)
     parser.add_argument("--mask_onset_exposure", action="store_true")
-    parser.add_argument("--exposure_d_model", type=int, default=64)
+    parser.add_argument(
+        "--d_model",
+        type=int,
+        default=64,
+        help="Internal TimesNet channel dimension for exposure encoding.",
+    )
     parser.add_argument("--exposure_n_layers", type=int, default=2)
     parser.add_argument("--exposure_top_k", type=int, default=2)
     parser.add_argument("--exposure_n_backbone_blocks", type=int, default=1)
@@ -306,6 +311,8 @@ def parse_args() -> argparse.Namespace:
         raise ValueError("prefetch_factor must be positive when num_workers > 0")
     if args.exposure_locality_buffer_size < 0:
         raise ValueError("exposure_locality_buffer_size must be non-negative")
+    if args.d_model <= 0:
+        raise ValueError("--d_model must be positive")
     if args.target_mode == "uts":
         args.readout_name = args.readout_name or "same_time_group_end"
         args.include_no_event_in_uts_target = True
@@ -434,7 +441,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         dist_mode="exponential",
         dropout=args.dropout,
         use_exposure_encoder=args.exposure_cache_dir is not None,
-        exposure_d_model=args.exposure_d_model,
+        exposure_d_model=args.d_model,
         exposure_n_layers=args.exposure_n_layers,
         exposure_top_k=args.exposure_top_k,
         exposure_n_backbone_blocks=args.exposure_n_backbone_blocks,
@@ -627,7 +634,7 @@ def build_metadata(
         "use_exposure_encoder": args.exposure_cache_dir is not None,
         "exposure_cache_dir": args.exposure_cache_dir,
         "mask_onset_exposure": bool(args.mask_onset_exposure),
-        "exposure_d_model": args.exposure_d_model,
+        "d_model": int(args.d_model),
         "exposure_n_layers": int(args.exposure_n_layers),
         "exposure_top_k": int(args.exposure_top_k),
         "exposure_n_backbone_blocks": int(args.exposure_n_backbone_blocks),