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

Refactor DeepHealth model and related components

Browse Source 
- Removed BaselineEncoder and CrossAttention classes from models.py.
- Introduced OtherInfoTokenizer for handling additional token types.
- Updated DeepHealth class to integrate OtherInfoTokenizer and manage extra pooling logic.
- Added support for extra_pool_reduce parameter to control pooling behavior.
- Modified forward methods to return structured output using DeepHealthOutput dataclass.
- Updated training scripts to accommodate changes in model architecture and output handling.
- Enhanced error handling and validation for input shapes and types.
This commit is contained in:
Jiarui Li
2026-06-17 11:05:10 +08:00
parent 27aefb2f90
commit 1757bcd25b
7 changed files with 435 additions and 493 deletions
Download Patch File Download Diff File Expand all files Collapse all files

352
backbones.py
View File

@@ -10,7 +10,7 @@ class TimeRoPE(nn.Module):
super().__init__()
assert dim % 2 == 0, "RoPE dim must be even"
self.dim = dim
# inv_freq: (dim // 2,) — not trainable, but should move with device
# inv_freq is not trainable, but should move with device.
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -73,7 +73,7 @@ class GaussianRBFTimeBasis(nn.Module):
def precompute_cache(self, tau: torch.Tensor) -> torch.Tensor:
time_coord = tau.float() # (B, L)
# Pairwise signed difference: query_i − key_j
# Pairwise signed difference: query_i - key_j.
diff = time_coord.unsqueeze(
2) - time_coord.unsqueeze(1) # (B, L_q, L_k)
# Gaussian RBF: exp(-0.5 * ((diff - c) / w)^2)
@@ -297,354 +297,6 @@ class TokenAutoDiscretization(nn.Module):
return torch.einsum("nb,nbd->nd", probs, e)
class BaselineEncoder(nn.Module):
PAD_KIND = 0
CONT_KIND = 1
CATE_KIND = 2
def __init__(
self,
n_embd: int,
n_head: int,
n_types: int,
n_cont_types: int,
n_categories: int,
cont_type_ids: list[int],
n_value_kinds: int = 3,
n_bins: int = 16,
n_tab_layer: int = 2,
dropout: float = 0.0,
):
super().__init__()
if len(cont_type_ids) != n_cont_types:
raise ValueError(
"cont_type_ids length must match n_cont_types, got "
f"{len(cont_type_ids)} vs {n_cont_types}"
)
if n_types <= 0:
raise ValueError(f"n_types must include PAD and be > 0, got {n_types}")
if n_categories <= 0:
raise ValueError(
f"n_categories must include PAD and be > 0, got {n_categories}"
)
if n_value_kinds <= self.CATE_KIND:
raise ValueError(
f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
)
self.n_embd = n_embd
self.cls_token = nn.Parameter(torch.zeros(1, 1, n_embd))
self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0)
self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
self.cont_value_encoder = (
TokenAutoDiscretization(
n_cont_types=n_cont_types,
n_bins=n_bins,
n_embd=n_embd,
)
if n_cont_types > 0
else None
)
self.cate_value_emb = nn.Embedding(
n_categories,
n_embd,
padding_idx=0,
)
cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
for idx, type_id in enumerate(cont_type_ids):
if type_id <= 0 or type_id >= n_types:
raise ValueError(
f"continuous type id {type_id} must be in [1, {n_types})"
)
cont_type_index[type_id] = idx
self.register_buffer(
"cont_type_index",
cont_type_index,
persistent=False,
)
self.blocks = nn.ModuleList([
GPTBlock(
n_embd=n_embd,
n_head=n_head,
use_time_rope=False,
use_rbf_bias=False,
mlp_dropout=dropout,
) for _ in range(n_tab_layer)
])
self.ln = nn.LayerNorm(n_embd)
self.reset_parameters()
def reset_parameters(self) -> None:
nn.init.normal_(self.cls_token, mean=0.0, std=0.02)
nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
nn.init.zeros_(self.type_emb.weight[0])
nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02)
nn.init.zeros_(self.kind_emb.weight[0])
nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02)
nn.init.zeros_(self.cate_value_emb.weight[0])
def _make_attn_mask(self, mask: torch.Tensor, dtype: torch.dtype):
return torch.zeros(
mask.size(0),
1,
1,
mask.size(1),
device=mask.device,
dtype=dtype,
).masked_fill(~mask[:, None, None, :], -1e4)
def _make_time_attn_mask(
self,
mask: torch.Tensor,
time: torch.Tensor,
dtype: torch.dtype,
):
valid_key = mask[:, None, :]
visible_by_time = time[:, None, :] <= time[:, :, None]
valid = valid_key & visible_by_time
return torch.zeros(
valid.shape,
device=valid.device,
dtype=dtype,
).masked_fill(~valid, -1e4)[:, None, :, :]
def forward(
self,
other_type: torch.LongTensor, # (B, K), 0 = padding
other_value: torch.Tensor, # (B, K), cate stores global id
other_value_kind: torch.LongTensor, # (B, K), 0=PAD, 1=CONT, 2=CATE
other_time: torch.Tensor | None = None, # (B, K)
cls_time: torch.Tensor | None = None, # (B,)
):
if other_type.shape != other_value.shape:
raise ValueError(
"other_type and other_value must have the same shape, got "
f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}"
)
if other_type.shape != other_value_kind.shape:
raise ValueError(
"other_type and other_value_kind must have the same shape, got "
f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}"
)
other_valid = other_type > 0
type_emb = self.type_emb(other_type)
kind_emb = self.kind_emb(other_value_kind)
value_emb = torch.zeros_like(type_emb)
cont_pos = other_valid & (other_value_kind == self.CONT_KIND)
if cont_pos.any():
if self.cont_value_encoder is None:
raise ValueError("continuous tokens found but n_cont_types is 0")
cont_idx = self.cont_type_index[other_type[cont_pos]]
if (cont_idx < 0).any():
bad_type = other_type[cont_pos][cont_idx < 0][0].item()
raise ValueError(
f"type_id={bad_type} is marked continuous but is not in "
"cont_type_ids"
)
value_emb[cont_pos] = self.cont_value_encoder(
cont_type_idx=cont_idx,
value=other_value[cont_pos].to(type_emb.dtype),
)
cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
if cate_pos.any():
cate_id = other_value[cate_pos].long()
value_emb[cate_pos] = self.cate_value_emb(cate_id)
f = type_emb + kind_emb + value_emb
f = f * other_valid.unsqueeze(-1).to(f.dtype)
cls = self.cls_token.expand(f.size(0), -1, -1)
f = torch.cat([cls, f], dim=1)
cls_valid = torch.ones(
other_valid.size(0),
1,
device=other_valid.device,
dtype=torch.bool,
)
full_valid = torch.cat([cls_valid, other_valid], dim=1)
if other_time is None:
attn_mask = self._make_attn_mask(full_valid, f.dtype)
else:
if other_time.shape != other_type.shape:
raise ValueError(
"other_time must have the same shape as other_type, got "
f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
)
if cls_time is None:
raise ValueError("cls_time is required when other_time is provided")
full_time = torch.cat(
[
cls_time.to(device=other_time.device, dtype=other_time.dtype)[:, None],
other_time,
],
dim=1,
)
attn_mask = self._make_time_attn_mask(full_valid, full_time, f.dtype)
for block in self.blocks:
f = block(f, attn_mask=attn_mask)
f = f * full_valid.unsqueeze(-1).to(f.dtype)
h = self.ln(f)
h = h * full_valid.unsqueeze(-1).to(h.dtype)
cls_summary = h[:, 0, :]
token_h = h[:, 1:, :]
token_h = token_h * other_valid.unsqueeze(-1).to(token_h.dtype)
return token_h, other_valid, cls_summary
class CrossAttention(nn.Module):
def __init__(
self,
n_embd: int,
n_head: int,
dropout: float = 0.0,
n_rbf_bases: int = 16,
max_time_diff: float = 40.0,
):
super().__init__()
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
self.n_head = n_head
self.d_head = n_embd // n_head
self.scale = 1.0 / math.sqrt(self.d_head)
self.dropout = dropout
self.mask_value = -1e4
self.q_proj = nn.Linear(n_embd, n_embd, bias=False)
self.k_proj = nn.Linear(n_embd, n_embd, bias=False)
self.v_proj = nn.Linear(n_embd, n_embd, bias=False)
self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
self.time_rope = TimeRoPE(self.d_head)
self.rbf_time_basis = GaussianRBFTimeBasis(
n_bases=n_rbf_bases,
max_time_diff=max_time_diff,
)
self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
self.resid_drop = nn.Dropout(dropout)
self.ln = nn.LayerNorm(n_embd)
self.out_ln = nn.LayerNorm(n_embd)
self.reset_parameters()
def reset_parameters(self) -> None:
nn.init.normal_(self.q_proj.weight, mean=0.0, std=0.02)
nn.init.normal_(self.k_proj.weight, mean=0.0, std=0.02)
nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02)
nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
nn.init.normal_(self.rbf_proj.weight, mean=0.0, std=0.02)
def _make_attn_mask(
self,
token_mask: torch.Tensor, # (B, K), True = valid
t_disease: torch.Tensor, # (B, L)
t_token: torch.Tensor, # (B, K)
dtype: torch.dtype,
):
valid_token = token_mask[:, None, :] # (B, 1, K)
visible_by_time = t_token[:, None, :] <= t_disease[:, :, None]
valid = visible_by_time & valid_token # (B, L, K)
has_any_valid = valid.any(dim=-1, keepdim=True)
safe_valid = valid.clone()
safe_valid[..., 0:1] = torch.where(
has_any_valid,
safe_valid[..., 0:1],
torch.ones_like(safe_valid[..., 0:1]),
)
attn_mask = torch.zeros(
safe_valid.shape,
device=safe_valid.device,
dtype=dtype,
).masked_fill(~safe_valid, self.mask_value)
return attn_mask[:, None, :, :], valid
def _cross_rbf_cache(
self,
t_disease: torch.Tensor,
t_token: torch.Tensor,
) -> torch.Tensor:
tau = torch.cat([t_disease, t_token], dim=1)
rbf_cache = self.rbf_time_basis.precompute_cache(tau)
n_disease = t_disease.size(1)
return rbf_cache[:, :n_disease, n_disease:, :]
def forward(
self,
h_disease: torch.Tensor, # (B, L, D)
t_disease: torch.Tensor, # (B, L)
h_token: torch.Tensor, # (B, K, D)
t_token: torch.Tensor, # (B, K)
token_mask: torch.Tensor, # (B, K), True = valid
need_weights: bool = False,
):
B, L, _ = h_disease.shape
K = h_token.size(1)
H, Dh = self.n_head, self.d_head
if K == 0:
empty_weights = h_disease.new_zeros(B, L, 0)
if need_weights:
return h_disease, empty_weights
return h_disease
attn_mask, valid = self._make_attn_mask(
token_mask=token_mask.to(device=h_disease.device, dtype=torch.bool),
t_disease=t_disease,
t_token=t_token,
dtype=h_disease.dtype,
)
q = self.q_proj(h_disease).reshape(B, L, H, Dh).transpose(1, 2)
k = self.k_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2)
v = self.v_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2)
q_cos, q_sin = self.time_rope.precompute_cache(t_disease)
k_cos, k_sin = self.time_rope.precompute_cache(t_token)
q = q * q_cos + TimeRoPE._rotate_half(q) * q_sin
k = k * k_cos + TimeRoPE._rotate_half(k) * k_sin
rbf_cache = self._cross_rbf_cache(t_disease, t_token) # (B, L, K, R)
time_bias = self.rbf_proj(rbf_cache).permute(0, 3, 1, 2)
time_bias = self.time_bias_scale.tanh() * time_bias
attn_bias = attn_mask.to(time_bias.dtype) + time_bias
attn_out = F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=attn_bias,
dropout_p=self.dropout if self.training else 0.0,
is_causal=False,
scale=self.scale,
)
attn_out = attn_out.transpose(1, 2).reshape(B, L, H * Dh)
attn_out = self.resid_drop(self.out_proj(attn_out))
has_any_valid = valid.any(dim=-1) # (B, L)
attn_out = attn_out * has_any_valid.unsqueeze(-1).to(attn_out.dtype)
attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
attn_scores = attn_scores + attn_bias
attn_weights = torch.softmax(attn_scores, dim=-1).mean(dim=1)
attn_weights = attn_weights * valid.to(attn_weights.dtype)
weight_sum = attn_weights.sum(dim=-1, keepdim=True).clamp_min(1e-12)
norm_weights = attn_weights / weight_sum
norm_weights = norm_weights * has_any_valid.unsqueeze(-1).to(
norm_weights.dtype
)
out = self.out_ln(h_disease + attn_out)
out = torch.where(has_any_valid.unsqueeze(-1), out, h_disease)
if need_weights:
return out, norm_weights
return out
class AgeSinusoidalEncoding(nn.Module):