JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Expose TimesNet internal dimension

Browse Source
This commit is contained in:
Jiarui Li
2026-07-09 16:22:34 +08:00
parent 8078fcb835
commit 552e09aa01
3 changed files with 13 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
train_next_step.py
View File

@@ -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),