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