Use precomputed exposure embeddings
This commit is contained in:
153
models.py
153
models.py
@@ -7,7 +7,6 @@ import torch.nn.functional as F
|
||||
from backbones import (
|
||||
AgeSinusoidalEncoding,
|
||||
GPTBlock,
|
||||
TimesNetExposureEncoder,
|
||||
)
|
||||
from targets import PAD_IDX
|
||||
|
||||
@@ -30,16 +29,7 @@ class DeepHealth(nn.Module):
|
||||
target_mode: str = "next_token", # "next_token" or "all_future"
|
||||
dist_mode: str = "exponential", # "exponential", "weibull" or "mixed"
|
||||
dropout: float = 0.0,
|
||||
use_exposure_encoder: bool = False,
|
||||
exposure_daily_input_dim: int = 4,
|
||||
exposure_monthly_input_dim: int = 2,
|
||||
exposure_d_model: int | None = 64,
|
||||
exposure_n_layers: int = 2,
|
||||
exposure_top_k: int = 2,
|
||||
exposure_n_backbone_blocks: int = 1,
|
||||
exposure_backbone_kernel_size: int = 5,
|
||||
exposure_backbone_expansion: float = 2.0,
|
||||
exposure_use_gate: bool = True,
|
||||
use_exposure_embeddings: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
if target_mode not in ["next_token", "all_future"]:
|
||||
@@ -53,26 +43,9 @@ class DeepHealth(nn.Module):
|
||||
2, n_embd) # Assuming binary gender
|
||||
self.target_mode = target_mode
|
||||
self.dist_mode = dist_mode
|
||||
self.use_exposure_encoder = use_exposure_encoder
|
||||
self.use_exposure_embeddings = bool(use_exposure_embeddings)
|
||||
self.n_embd = n_embd
|
||||
self.vocab_size = vocab_size
|
||||
self.exposure_encoder = (
|
||||
TimesNetExposureEncoder(
|
||||
n_embd=n_embd,
|
||||
daily_input_dim=exposure_daily_input_dim,
|
||||
monthly_input_dim=exposure_monthly_input_dim,
|
||||
d_model=exposure_d_model,
|
||||
n_layers=exposure_n_layers,
|
||||
top_k=exposure_top_k,
|
||||
n_backbone_blocks=exposure_n_backbone_blocks,
|
||||
backbone_kernel_size=exposure_backbone_kernel_size,
|
||||
backbone_expansion=exposure_backbone_expansion,
|
||||
dropout=dropout,
|
||||
use_gate=exposure_use_gate,
|
||||
)
|
||||
if use_exposure_encoder
|
||||
else None
|
||||
)
|
||||
nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.token_embedding.weight[0])
|
||||
nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
|
||||
@@ -121,95 +94,6 @@ class DeepHealth(nn.Module):
|
||||
dtype=dtype,
|
||||
).masked_fill(~valid, -1e4)[:, None, :, :]
|
||||
|
||||
def _encode_event_exposure(
|
||||
self,
|
||||
exposure_daily: torch.Tensor | None,
|
||||
exposure_monthly: torch.Tensor | None,
|
||||
exposure_daily_mask: torch.Tensor | None,
|
||||
exposure_monthly_mask: torch.Tensor | None,
|
||||
event_shape: tuple[int, int],
|
||||
) -> torch.Tensor | None:
|
||||
if self.exposure_encoder is None:
|
||||
if exposure_daily is not None or exposure_monthly is not None:
|
||||
raise ValueError(
|
||||
"Exposure tensors were provided but use_exposure_encoder=False"
|
||||
)
|
||||
return None
|
||||
if exposure_daily is None or exposure_monthly is None:
|
||||
raise ValueError(
|
||||
"exposure_daily and exposure_monthly are required when "
|
||||
"use_exposure_encoder=True"
|
||||
)
|
||||
|
||||
batch_size, event_len = event_shape
|
||||
if exposure_daily.shape[:2] != event_shape:
|
||||
raise ValueError(
|
||||
"exposure_daily must have shape (B, L, T_daily, C_daily), got "
|
||||
f"{tuple(exposure_daily.shape)} for event shape {event_shape}"
|
||||
)
|
||||
if exposure_monthly.shape[:2] != event_shape:
|
||||
raise ValueError(
|
||||
"exposure_monthly must have shape (B, L, T_monthly, C_monthly), got "
|
||||
f"{tuple(exposure_monthly.shape)} for event shape {event_shape}"
|
||||
)
|
||||
if exposure_daily.dim() != 4 or exposure_monthly.dim() != 4:
|
||||
raise ValueError(
|
||||
"exposure_daily and exposure_monthly must both be 4D tensors"
|
||||
)
|
||||
|
||||
def flatten_mask(mask: torch.Tensor | None, ref: torch.Tensor) -> torch.Tensor | None:
|
||||
if mask is None:
|
||||
return None
|
||||
if mask.shape[:2] != event_shape:
|
||||
raise ValueError(
|
||||
"exposure mask must start with shape (B, L), got "
|
||||
f"{tuple(mask.shape)}"
|
||||
)
|
||||
if mask.dim() not in {3, 4}:
|
||||
raise ValueError(
|
||||
"exposure mask must have shape (B, L, T) or (B, L, T, C), "
|
||||
f"got {tuple(mask.shape)}"
|
||||
)
|
||||
if mask.shape[2] != ref.shape[2]:
|
||||
raise ValueError(
|
||||
"exposure mask time dimension does not match exposure tensor"
|
||||
)
|
||||
if mask.dim() == 4 and mask.shape[3] != ref.shape[3]:
|
||||
raise ValueError(
|
||||
"exposure mask channel dimension does not match exposure tensor"
|
||||
)
|
||||
return mask.reshape(batch_size * event_len, *mask.shape[2:])
|
||||
|
||||
daily = exposure_daily.reshape(
|
||||
batch_size * event_len,
|
||||
exposure_daily.size(2),
|
||||
exposure_daily.size(3),
|
||||
)
|
||||
monthly = exposure_monthly.reshape(
|
||||
batch_size * event_len,
|
||||
exposure_monthly.size(2),
|
||||
exposure_monthly.size(3),
|
||||
)
|
||||
daily_mask = flatten_mask(exposure_daily_mask, exposure_daily)
|
||||
monthly_mask = flatten_mask(exposure_monthly_mask, exposure_monthly)
|
||||
|
||||
exposure_device = exposure_daily.device
|
||||
exposure_dtype = self.token_embedding.weight.dtype
|
||||
daily = daily.to(device=exposure_device, dtype=exposure_dtype)
|
||||
monthly = monthly.to(device=exposure_device, dtype=exposure_dtype)
|
||||
if daily_mask is not None:
|
||||
daily_mask = daily_mask.to(device=exposure_device)
|
||||
if monthly_mask is not None:
|
||||
monthly_mask = monthly_mask.to(device=exposure_device)
|
||||
|
||||
exposure_emb = self.exposure_encoder(
|
||||
daily=daily,
|
||||
monthly=monthly,
|
||||
daily_mask=daily_mask,
|
||||
monthly_mask=monthly_mask,
|
||||
)
|
||||
return exposure_emb.reshape(batch_size, event_len, self.n_embd)
|
||||
|
||||
def _forward_shared(
|
||||
self,
|
||||
event_seq: torch.LongTensor,
|
||||
@@ -218,10 +102,7 @@ class DeepHealth(nn.Module):
|
||||
mode: str,
|
||||
padding_mask: torch.Tensor | None = None,
|
||||
t_query: torch.FloatTensor | None = None,
|
||||
exposure_daily: torch.Tensor | None = None,
|
||||
exposure_monthly: torch.Tensor | None = None,
|
||||
exposure_daily_mask: torch.Tensor | None = None,
|
||||
exposure_monthly_mask: torch.Tensor | None = None,
|
||||
exposure_embedding: torch.Tensor | None = None,
|
||||
return_output: bool = False,
|
||||
**unused_kwargs,
|
||||
) -> torch.Tensor | DeepHealthOutput:
|
||||
@@ -240,16 +121,24 @@ class DeepHealth(nn.Module):
|
||||
|
||||
event_len = event_seq.size(1)
|
||||
h_disease = self.token_embedding(event_seq)
|
||||
h_exposure = self._encode_event_exposure(
|
||||
exposure_daily=exposure_daily,
|
||||
exposure_monthly=exposure_monthly,
|
||||
exposure_daily_mask=exposure_daily_mask,
|
||||
exposure_monthly_mask=exposure_monthly_mask,
|
||||
event_shape=(event_seq.size(0), event_len),
|
||||
)
|
||||
if h_exposure is not None:
|
||||
h_exposure = h_exposure.to(device=event_seq.device, dtype=h_disease.dtype)
|
||||
h_disease = h_disease + h_exposure
|
||||
if self.use_exposure_embeddings:
|
||||
if exposure_embedding is None:
|
||||
raise ValueError(
|
||||
"exposure_embedding is required when "
|
||||
"use_exposure_embeddings=True"
|
||||
)
|
||||
if exposure_embedding.shape != h_disease.shape:
|
||||
raise ValueError(
|
||||
"exposure_embedding must have shape "
|
||||
f"{tuple(h_disease.shape)}, got {tuple(exposure_embedding.shape)}"
|
||||
)
|
||||
h_disease = h_disease + exposure_embedding.to(
|
||||
device=h_disease.device, dtype=h_disease.dtype
|
||||
)
|
||||
elif exposure_embedding is not None:
|
||||
raise ValueError(
|
||||
"exposure_embedding provided but use_exposure_embeddings=False"
|
||||
)
|
||||
t_disease = time_seq
|
||||
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user