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

Lighten TimesNet exposure backbone

Browse Source
This commit is contained in:
Jiarui Li
2026-07-09 14:23:28 +08:00
parent 54fedc620b
commit f7fb6b7718
5 changed files with 95 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
train_exposure_autoencoder.py
View File

@@ -1,4 +1,4 @@
"""Pretrain a TimesNet + ConvNeXtV2 autoencoder on training-set exposure."""
"""Pretrain a lightweight TimesNet autoencoder on training-set exposure."""
from __future__ import annotations
import argparse
@@ -46,7 +46,7 @@ class ExposureWindowDataset(Dataset):
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Pretrain a TimesNet + ConvNeXtV2 exposure autoencoder"
description="Pretrain a lightweight TimesNet exposure autoencoder"
)
parser.add_argument("--exposure_cache_dir", required=True)
parser.add_argument("--train_eid_file", default="ukb_train_eid.csv")
@@ -54,12 +54,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--runs_root", default="runs")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--n_embd", type=int, default=120)
parser.add_argument("--d_model", type=int, default=None)
parser.add_argument("--d_model", type=int, default=64)
parser.add_argument("--n_layers", type=int, default=2)
parser.add_argument("--top_k", type=int, default=3)
parser.add_argument("--n_convnext_blocks", type=int, default=2)
parser.add_argument("--conv_kernel_size", type=int, default=7)
parser.add_argument("--mlp_ratio", type=float, default=4.0)
parser.add_argument("--top_k", type=int, default=2)
parser.add_argument("--n_backbone_blocks", type=int, default=1)
parser.add_argument("--backbone_kernel_size", type=int, default=5)
parser.add_argument("--backbone_expansion", type=float, default=2.0)
parser.add_argument("--dropout", type=float, default=0.0)
parser.add_argument("--mask_ratio", type=float, default=0.25)
parser.add_argument("--batch_size", type=int, default=16)
@@ -288,8 +288,9 @@ def main() -> None:
)
model = TimesNetExposureAutoencoder(
n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers,
top_k=args.top_k, n_convnext_blocks=args.n_convnext_blocks,
conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio,
top_k=args.top_k, n_backbone_blocks=args.n_backbone_blocks,
backbone_kernel_size=args.backbone_kernel_size,
backbone_expansion=args.backbone_expansion,
dropout=args.dropout,
).to(device)
model = maybe_wrap_data_parallel(model, args, device, logger)
@@ -349,8 +350,8 @@ def main() -> None:
"model_config": {
key: config[key] for key in (
"n_embd", "d_model", "n_layers", "top_k",
"n_convnext_blocks", "conv_kernel_size",
"mlp_ratio", "dropout",
"n_backbone_blocks", "backbone_kernel_size",
"backbone_expansion", "dropout",
)
},
"normalization": {