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

Refactor loss computation and model input handling for improved clarity and efficiency

Browse Source
This commit is contained in:
Jiarui Li
2026-06-20 11:26:03 +08:00
parent c3cac6dcea
commit e5ecb714ba
4 changed files with 220 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
losses.py
View File

@@ -99,29 +99,33 @@ class Delphi2MLoss(nn.Module):
}
return total_loss
logits_valid = logits[valid_mask]
target_events_valid = target_events[valid_mask]
target_times_valid = target_times[valid_mask]
current_times_valid = current_times[valid_mask]
logits_safe = torch.nan_to_num(
logits,
logits_valid,
nan=0.0,
posinf=self.max_exp_input,
neginf=-self.max_exp_input,
)
loss_ce = F.cross_entropy(
logits_safe.reshape(-1, vocab_size)[valid_mask.reshape(-1)],
target_events.reshape(-1)[valid_mask.reshape(-1)],
logits_safe,
target_events_valid,
reduction="mean",
)
logits_for_lse = logits_safe
if self.exclude_ignored_from_intensity:
ignore_mask = _make_ignore_mask(vocab_size, self.ignored_tokens, logits.device)
logits_for_lse = logits_safe.clone()
logits_for_lse[:, :, ignore_mask] = float("-inf")
logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))
log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
log_lambda_total = -torch.log(torch.exp(-log_lambda_total) + self.t_min)
dt = torch.clamp(target_times - current_times, min=self.t_min)
dt = torch.clamp(target_times_valid - current_times_valid, min=self.t_min)
log_dt_inv = -torch.log(dt + self.t_min)
loss_dt = -(
log_lambda_total
@@ -129,7 +133,7 @@ class Delphi2MLoss(nn.Module):
torch.clamp(log_lambda_total - log_dt_inv, max=self.max_exp_input)
)
)
loss_dt = loss_dt[valid_mask].mean()
loss_dt = loss_dt.mean()
total_loss = self.ce_weight * loss_ce + self.time_weight * loss_dt
if return_components:
@@ -181,17 +185,9 @@ class UniqueTimeSetExponentialLoss(nn.Module):
if readout_mask.shape != (bsz, seq_len):
raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}")
logits_safe = torch.nan_to_num(
logits,
nan=0.0,
posinf=self.max_exp_input,
neginf=-self.max_exp_input,
)
ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device)
target_valid = target_multi_hot.float().clone()
target_valid[:, :, ignore_mask] = 0.0
num_targets = target_valid.sum(dim=-1)
num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1)
valid_mask = readout_mask.bool() & (num_targets > 0)
if not valid_mask.any():
@@ -204,34 +200,35 @@ class UniqueTimeSetExponentialLoss(nn.Module):
}
return total_loss
log_lambda_targets = logits_safe * target_valid
log_lambda_targets = torch.where(
target_valid.bool(),
log_lambda_targets,
torch.zeros_like(log_lambda_targets),
logits_safe = torch.nan_to_num(
logits[valid_mask],
nan=0.0,
posinf=self.max_exp_input,
neginf=-self.max_exp_input,
)
target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype)
target_valid[:, ignore_mask] = 0.0
observed_term = log_lambda_targets.sum(dim=-1)
penalty_scale = num_targets
observed_term = (logits_safe * target_valid).sum(dim=-1)
penalty_scale = target_valid.sum(dim=-1)
logits_for_lse = logits_safe
if self.exclude_ignored_from_intensity:
logits_for_lse = logits_safe.clone()
logits_for_lse[:, :, ignore_mask] = float("-inf")
logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))
dt_clamped = torch.clamp(target_dt_unique, min=self.t_min)
dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min)
log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
log_penalty = log_lambda_total + dt_clamped.log()
penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input))
observed_loss = -observed_term
penalty_loss = penalty_scale * penalty
total_loss = (observed_loss + penalty_loss)[valid_mask].mean()
total_loss = (observed_loss + penalty_loss).mean()
if return_components:
return total_loss, {
"observed": observed_loss[valid_mask].mean().detach(),
"penalty": penalty_loss[valid_mask].mean().detach(),
"observed": observed_loss.mean().detach(),
"penalty": penalty_loss.mean().detach(),
"total": total_loss.detach(),
}
return total_loss

69
models.py
View File

@@ -260,12 +260,27 @@ class DeepHealth(nn.Module):
other_time: torch.Tensor,
other_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
batch_size, _n_other, n_embd = h_other.shape
group_counts = [
int(torch.unique(other_time[b, other_mask[b]], sorted=True).numel())
for b in range(batch_size)
]
max_groups = max(group_counts, default=0)
batch_size, n_other, n_embd = h_other.shape
if n_other == 0:
empty_h = h_other.new_zeros(batch_size, 0, n_embd)
empty_t = other_time.new_zeros(batch_size, 0)
empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device)
return empty_h, empty_t, empty_m
masked_time = other_time.masked_fill(~other_mask, float("inf"))
_sorted_time_with_pad, order = masked_time.sort(dim=1)
sorted_time = other_time.gather(1, order)
sorted_mask = other_mask.gather(1, order)
sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd))
group_start = torch.zeros_like(sorted_mask)
group_start[:, 0] = sorted_mask[:, 0]
group_start[:, 1:] = sorted_mask[:, 1:] & (
sorted_time[:, 1:] != sorted_time[:, :-1]
)
group_id = group_start.long().cumsum(dim=1) - 1
max_groups = int(group_start.sum(dim=1).max().item())
pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd)
pooled_time = other_time.new_zeros(batch_size, max_groups)
pooled_mask = torch.zeros(
@@ -277,35 +292,29 @@ class DeepHealth(nn.Module):
if max_groups == 0:
return pooled_h, pooled_time, pooled_mask
for b in range(batch_size):
valid_time = other_time[b, other_mask[b]]
if valid_time.numel() == 0:
continue
valid_h = h_other[b, other_mask[b]]
unique_time, inverse = torch.unique(
valid_time,
sorted=True,
return_inverse=True,
)
n_groups = unique_time.numel()
group_h = valid_h.new_zeros(n_groups, n_embd)
group_h.scatter_add_(
0,
inverse[:, None].expand(-1, n_embd),
valid_h,
safe_group_id = group_id.clamp_min(0)
pooled_h.scatter_add_(
1,
safe_group_id.unsqueeze(-1).expand_as(sorted_h),
sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype),
)
if self.extra_pool_reduce == "mean":
counts = valid_h.new_zeros(n_groups, 1)
counts = h_other.new_zeros(batch_size, max_groups, 1)
counts.scatter_add_(
0,
inverse[:, None],
torch.ones_like(valid_h[:, :1]),
1,
safe_group_id.unsqueeze(-1),
sorted_mask.unsqueeze(-1).to(h_other.dtype),
)
group_h = group_h / counts.clamp_min(1.0)
pooled_h = pooled_h / counts.clamp_min(1.0)
pooled_h[b, :n_groups] = group_h
pooled_time[b, :n_groups] = unique_time
pooled_mask[b, :n_groups] = True
pooled_time.scatter_add_(
1,
safe_group_id,
sorted_time * group_start.to(sorted_time.dtype),
)
group_count = group_start.sum(dim=1)
arange_groups = torch.arange(max_groups, device=h_other.device)
pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1)
return pooled_h, pooled_time, pooled_mask
def _forward_shared(

37
train_all_future.py
View File

@@ -44,6 +44,19 @@ from train_util import (
)
MODEL_INPUT_KEYS = (
"event_seq",
"time_seq",
"sex",
"padding_mask",
"t_query",
"other_type",
"other_value",
"other_value_kind",
"other_time",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Train DeepHealth with all-future supervision")
@@ -87,6 +100,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--min_lr_ratio", type=float, default=0.1)
parser.add_argument("--num_workers", type=int, default=4)
parser.add_argument("--device", type=str, default="cuda")
parser.add_argument("--progress_interval", type=int, default=20)
args = parser.parse_args()
if args.min_history_events < 1:
@@ -169,7 +183,14 @@ def compute_all_future_loss(
batch: Dict[str, torch.Tensor],
device: torch.device,
) -> torch.Tensor:
batch = move_batch_to_device(batch, device)
required_keys = set(MODEL_INPUT_KEYS)
required_keys.update(("future_targets", "exposure"))
if args.dist_mode in {"weibull", "mixed"}:
required_keys.add("future_dt")
batch = move_batch_to_device(
{key: batch[key] for key in required_keys},
device,
)
hidden = model(
event_seq=batch["event_seq"],
@@ -224,10 +245,11 @@ def run_epoch(
is_train: bool,
) -> float:
model.train(is_train)
total = 0.0
total = torch.zeros((), device=device)
n_batches = 0
skipped = 0
desc = "train" if is_train else "val"
progress_interval = max(1, int(args.progress_interval))
progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
for batch_idx, batch in enumerate(progress):
@@ -242,10 +264,15 @@ def run_epoch(
clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step()
total += float(loss.detach().cpu())
total = total + loss.detach()
n_batches += 1
if (batch_idx + 1) % progress_interval == 0:
avg = total / max(1, n_batches)
progress.set_postfix(loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{avg:.4f}", skipped=skipped)
progress.set_postfix(
loss=f"{float(loss.detach().cpu()):.4f}",
avg=f"{float(avg.detach().cpu()):.4f}",
skipped=skipped,
)
except RuntimeError as exc:
if "Loss is not finite" not in str(exc):
raise
@@ -254,7 +281,7 @@ def run_epoch(
if skipped:
logger.info(f"Skipped {skipped} batches due to non-finite loss")
return total / max(1, n_batches) if n_batches else float("inf")
return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
def build_metadata(

143
train_next_step.py
View File

@@ -42,6 +42,18 @@ from train_util import (
)
MODEL_INPUT_KEYS = (
"event_seq",
"time_seq",
"sex",
"padding_mask",
"other_type",
"other_value",
"other_value_kind",
"other_time",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Train DeepHealth with next-token/point supervision")
@@ -94,6 +106,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--min_lr_ratio", type=float, default=0.1)
parser.add_argument("--num_workers", type=int, default=4)
parser.add_argument("--device", type=str, default="cuda")
parser.add_argument("--progress_interval", type=int, default=20)
args = parser.parse_args()
use_eid_split = all(
@@ -181,92 +194,111 @@ def build_next_step_loss(args: argparse.Namespace):
def build_augmented_next_step_targets(
batch: Dict[str, torch.Tensor],
batch_cpu: Dict[str, torch.Tensor],
model_out: DeepHealthOutput,
include_uts_targets: bool,
) -> Dict[str, torch.Tensor]:
hidden_len = model_out.hidden.size(1)
event_len = int(model_out.event_len)
extra_len = hidden_len - event_len
if extra_len <= 0:
return {
"target_event_seq": batch["target_event_seq"],
"target_time_seq": batch["target_time_seq"],
"readout_mask": batch["readout_mask"],
"target_dt_unique": batch["target_dt_unique"],
"target_multi_hot": batch["target_multi_hot"],
}
device = model_out.hidden.device
bsz, _seq_len, vocab_size = batch["target_multi_hot"].shape
extra_mask = model_out.padding_mask[:, event_len:]
extra_time = model_out.time_seq[:, event_len:]
non_blocking = device.type == "cuda"
if extra_len <= 0:
targets = {
"target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
"target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
"readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
}
if include_uts_targets:
targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
device, non_blocking=non_blocking
)
targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
device, non_blocking=non_blocking
)
return targets
bsz = batch_cpu["target_event_seq"].size(0)
vocab_size = (
batch_cpu["target_multi_hot"].size(2)
if include_uts_targets
else None
)
other_valid = batch_cpu["other_type"] > 0
extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
for b in range(bsz):
unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True)
n_time = min(int(unique_time.numel()), extra_len)
if n_time > 0:
extra_time[b, :n_time] = unique_time[:n_time]
extra_mask[b, :n_time] = True
target_event_seq = torch.cat(
[
batch["target_event_seq"],
batch_cpu["target_event_seq"],
torch.full(
(bsz, extra_len),
PAD_IDX,
dtype=batch["target_event_seq"].dtype,
device=device,
dtype=batch_cpu["target_event_seq"].dtype,
),
],
dim=1,
)
target_time_seq = torch.cat(
[
batch["target_time_seq"],
batch_cpu["target_time_seq"],
torch.zeros(
bsz,
extra_len,
dtype=batch["target_time_seq"].dtype,
device=device,
dtype=batch_cpu["target_time_seq"].dtype,
),
],
dim=1,
)
readout_mask = torch.cat([batch["readout_mask"], extra_mask], dim=1)
readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
target_dt_unique = None
target_multi_hot = None
if include_uts_targets:
target_dt_unique = torch.cat(
[
batch["target_dt_unique"],
batch_cpu["target_dt_unique"],
torch.zeros(
bsz,
extra_len,
dtype=batch["target_dt_unique"].dtype,
device=device,
dtype=batch_cpu["target_dt_unique"].dtype,
),
],
dim=1,
)
target_multi_hot = torch.cat(
[
batch["target_multi_hot"],
batch_cpu["target_multi_hot"],
torch.zeros(
bsz,
extra_len,
vocab_size,
dtype=batch["target_multi_hot"].dtype,
device=device,
dtype=batch_cpu["target_multi_hot"].dtype,
),
],
dim=1,
)
for b in range(bsz):
valid_event = batch["padding_mask"][b].bool()
valid_event = batch_cpu["padding_mask"][b].bool()
if not valid_event.any():
continue
n_event = int(valid_event.sum().item())
events = torch.cat(
[
batch["event_seq"][b, :n_event],
batch["target_event_seq"][b, n_event - 1:n_event],
batch_cpu["event_seq"][b, :n_event],
batch_cpu["target_event_seq"][b, n_event - 1:n_event],
]
)
times = torch.cat(
[
batch["time_seq"][b, :n_event],
batch["target_time_seq"][b, n_event - 1:n_event],
batch_cpu["time_seq"][b, :n_event],
batch_cpu["target_time_seq"][b, n_event - 1:n_event],
]
)
valid_full = events > PAD_IDX
@@ -291,6 +323,9 @@ def build_augmented_next_step_targets(
target_event_seq[b, pos] = next_event
target_time_seq[b, pos] = next_time
if not include_uts_targets:
continue
same_next_time = times == next_time
next_events = events[same_next_time]
valid_next_events = next_events[
@@ -302,13 +337,15 @@ def build_augmented_next_step_targets(
target_multi_hot[b, pos, valid_next_events] = True
target_dt_unique[b, pos] = next_time - t
return {
"target_event_seq": target_event_seq,
"target_time_seq": target_time_seq,
"readout_mask": readout_mask,
"target_dt_unique": target_dt_unique,
"target_multi_hot": target_multi_hot,
targets = {
"target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
"target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
"readout_mask": readout_mask.to(device, non_blocking=non_blocking),
}
if include_uts_targets:
targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking)
targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
return targets
def compute_next_step_loss(
@@ -319,7 +356,11 @@ def compute_next_step_loss(
batch: Dict[str, torch.Tensor],
device: torch.device,
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
batch = move_batch_to_device(batch, device)
batch_cpu = batch
batch = move_batch_to_device(
{key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
device,
)
model_out = model(
event_seq=batch["event_seq"],
time_seq=batch["time_seq"],
@@ -334,7 +375,11 @@ def compute_next_step_loss(
)
if not isinstance(model_out, DeepHealthOutput):
raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
targets = build_augmented_next_step_targets(batch=batch, model_out=model_out)
targets = build_augmented_next_step_targets(
batch_cpu=batch_cpu,
model_out=model_out,
include_uts_targets=args.target_mode == "uts",
)
readout_out = readout(
hidden=model_out.hidden,
time_seq=model_out.time_seq,
@@ -380,11 +425,12 @@ def run_epoch(
) -> float:
model.train(is_train)
readout.train(is_train)
total = 0.0
total = torch.zeros((), device=device)
n_batches = 0
skipped = 0
parts_sum: Dict[str, float] = {}
parts_sum: Dict[str, torch.Tensor] = {}
desc = "train" if is_train else "val"
progress_interval = max(1, int(args.progress_interval))
progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
for batch_idx, batch in enumerate(progress):
@@ -399,14 +445,19 @@ def run_epoch(
clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step()
total += float(loss.detach().cpu())
total = total + loss.detach()
n_batches += 1
for name, value in parts.items():
parts_sum[name] = parts_sum.get(name, 0.0) + float(value.detach().cpu())
parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
if (batch_idx + 1) % progress_interval == 0:
avg = total / max(1, n_batches)
postfix = {"loss": f"{float(loss.detach().cpu()):.4f}", "avg": f"{avg:.4f}", "skipped": skipped}
postfix = {
"loss": f"{float(loss.detach().cpu()):.4f}",
"avg": f"{float(avg.detach().cpu()):.4f}",
"skipped": skipped,
}
for name, value in parts_sum.items():
postfix[name] = f"{value / max(1, n_batches):.4f}"
postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}"
progress.set_postfix(postfix)
except RuntimeError as exc:
if "Loss is not finite" not in str(exc):
@@ -416,7 +467,7 @@ def run_epoch(
if skipped:
logger.info(f"Skipped {skipped} batches due to non-finite loss")
return total / max(1, n_batches) if n_batches else float("inf")
return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
def build_metadata(