Improve all-future first-onset training

This commit is contained in:
2026-08-21 13:48:59 +08:00
parent 9ccc6b56ec
commit d728d8585c
9 changed files with 1002 additions and 84 deletions

View File

@@ -66,7 +66,7 @@ from evaluate_auc_v2 import (
)
from losses import build_loss
from model_architectures import resolve_model_architecture
from targets import PAD_IDX, RESERVED_IDX
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
from train_util import load_eid_file
@@ -1483,7 +1483,7 @@ def build_calibration_summary(metrics: pd.DataFrame) -> pd.DataFrame:
def _build_point_process_criterion(
dist_mode: str,
) -> Any:
ignored = {PAD_IDX, RESERVED_IDX}
ignored = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
if dist_mode == "exponential":
return build_loss("exponential", ignored_idx=ignored)
if dist_mode == "weibull":
@@ -1545,6 +1545,8 @@ def evaluate_point_process_nll(
logits=logits,
targets=batch_device["future_targets"],
exposure=batch_device["exposure"],
dt=batch_device["future_dt"],
history=batch_device["event_seq"],
)
elif dist_mode == "weibull":
loss = criterion(
@@ -1553,6 +1555,7 @@ def evaluate_point_process_nll(
targets=batch_device["future_targets"],
dt=batch_device["future_dt"],
exposure=batch_device["exposure"],
history=batch_device["event_seq"],
)
else:
raise ValueError(f"Unsupported dist_mode: {dist_mode!r}")
@@ -1564,6 +1567,7 @@ def evaluate_point_process_nll(
query_count += batch_size
valid_targets = batch["future_targets"] > PAD_IDX
valid_targets &= batch["future_targets"] != RESERVED_IDX
valid_targets &= batch["future_targets"] != NO_EVENT_IDX
future_event_count += int(valid_targets.sum().item())
exposure_sum += float(batch["exposure"].sum().item())