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

Refactor code structure for improved readability and maintainability

Browse Source
This commit is contained in:
Jiarui Li
2026-06-18 13:07:35 +08:00
parent 1ea72e9133
commit aa8ec5c3ac
7 changed files with 1005002 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
train_next_step.py
View File

@@ -38,6 +38,7 @@ from train_util import (
set_seed,
setup_logging,
split_dataset,
split_dataset_by_eid_files,
)
@@ -55,6 +56,9 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--train_ratio", type=float, default=0.7)
parser.add_argument("--val_ratio", type=float, default=0.15)
parser.add_argument("--test_ratio", type=float, default=0.15)
parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
parser.add_argument("--n_embd", type=int, default=120)
parser.add_argument("--n_head", type=int, default=10)
@@ -92,7 +96,11 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--device", type=str, default="cuda")
args = parser.parse_args()
if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
use_eid_split = all(
getattr(args, name)
for name in ("train_eid_file", "val_eid_file", "test_eid_file")
)
if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
if args.target_mode == "uts":
args.readout_name = args.readout_name or "same_time_group_end"
@@ -477,13 +485,29 @@ def main() -> None:
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
extra_info_types=args.extra_info_types,
)
train_subset, val_subset, test_subset = split_dataset(
dataset=dataset,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
train_subset, val_subset, test_subset = split_dataset_by_eid_files(
dataset=dataset,
train_eid_file=args.train_eid_file,
val_eid_file=args.val_eid_file,
test_eid_file=args.test_eid_file,
)
logger.info(
"Using eid split files: "
f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
)
else:
train_subset, val_subset, test_subset = split_dataset(
dataset=dataset,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
logger.info(
f"Using random ratio split: train={args.train_ratio}, "
f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
)
logger.info(
f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
)