Refactor DeepHealth model and related components
- 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:
@@ -24,7 +24,7 @@ from tqdm.auto import tqdm
|
||||
|
||||
from dataset import HealthDataset, collate_fn
|
||||
from losses import build_loss
|
||||
from models import DeepHealth
|
||||
from models import DeepHealth, DeepHealthOutput
|
||||
from readouts import build_readout
|
||||
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
||||
from train_util import (
|
||||
@@ -61,6 +61,8 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--n_hist_layer", type=int, default=12)
|
||||
parser.add_argument("--n_tab_layer", type=int, default=4)
|
||||
parser.add_argument("--n_bins", type=int, default=16)
|
||||
parser.add_argument("--extra_pool_reduce", type=str, default="mean",
|
||||
choices=["mean", "sum"])
|
||||
parser.add_argument("--time_mode", type=str, default="relative",
|
||||
choices=["relative", "absolute"])
|
||||
parser.add_argument("--dropout", type=float, default=0.0)
|
||||
@@ -135,6 +137,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
|
||||
n_categories=dataset.n_categories,
|
||||
cont_type_ids=dataset.cont_type_ids,
|
||||
n_bins=args.n_bins,
|
||||
extra_pool_reduce=args.extra_pool_reduce,
|
||||
target_mode="next_token",
|
||||
time_mode=args.time_mode,
|
||||
dist_mode="exponential",
|
||||
@@ -169,6 +172,137 @@ def build_next_step_loss(args: argparse.Namespace):
|
||||
)
|
||||
|
||||
|
||||
def build_augmented_next_step_targets(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
model_out: DeepHealthOutput,
|
||||
) -> 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:]
|
||||
|
||||
target_event_seq = torch.cat(
|
||||
[
|
||||
batch["target_event_seq"],
|
||||
torch.full(
|
||||
(bsz, extra_len),
|
||||
PAD_IDX,
|
||||
dtype=batch["target_event_seq"].dtype,
|
||||
device=device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
target_time_seq = torch.cat(
|
||||
[
|
||||
batch["target_time_seq"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
dtype=batch["target_time_seq"].dtype,
|
||||
device=device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
readout_mask = torch.cat([batch["readout_mask"], extra_mask], dim=1)
|
||||
target_dt_unique = torch.cat(
|
||||
[
|
||||
batch["target_dt_unique"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
dtype=batch["target_dt_unique"].dtype,
|
||||
device=device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
target_multi_hot = torch.cat(
|
||||
[
|
||||
batch["target_multi_hot"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
vocab_size,
|
||||
dtype=batch["target_multi_hot"].dtype,
|
||||
device=device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
for b in range(bsz):
|
||||
valid_event = batch["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],
|
||||
]
|
||||
)
|
||||
times = torch.cat(
|
||||
[
|
||||
batch["time_seq"][b, :n_event],
|
||||
batch["target_time_seq"][b, n_event - 1:n_event],
|
||||
]
|
||||
)
|
||||
valid_full = events > PAD_IDX
|
||||
events = events[valid_full]
|
||||
times = times[valid_full]
|
||||
if events.numel() == 0:
|
||||
continue
|
||||
|
||||
for j in range(extra_len):
|
||||
if not bool(extra_mask[b, j]):
|
||||
continue
|
||||
pos = event_len + j
|
||||
t = extra_time[b, j]
|
||||
future = times > t
|
||||
if not future.any():
|
||||
readout_mask[b, pos] = False
|
||||
continue
|
||||
|
||||
first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
|
||||
next_time = times[first_idx]
|
||||
next_event = events[first_idx]
|
||||
target_event_seq[b, pos] = next_event
|
||||
target_time_seq[b, pos] = next_time
|
||||
|
||||
same_next_time = times == next_time
|
||||
next_events = events[same_next_time]
|
||||
valid_next_events = next_events[
|
||||
(next_events > PAD_IDX) & (next_events < vocab_size)
|
||||
].long()
|
||||
if valid_next_events.numel() == 0:
|
||||
readout_mask[b, pos] = False
|
||||
continue
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
def compute_next_step_loss(
|
||||
args: argparse.Namespace,
|
||||
model: DeepHealth,
|
||||
@@ -178,7 +312,7 @@ def compute_next_step_loss(
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
batch = move_batch_to_device(batch, device)
|
||||
hidden = model(
|
||||
model_out = model(
|
||||
event_seq=batch["event_seq"],
|
||||
time_seq=batch["time_seq"],
|
||||
sex=batch["sex"],
|
||||
@@ -188,12 +322,16 @@ def compute_next_step_loss(
|
||||
other_value_kind=batch["other_value_kind"],
|
||||
other_time=batch["other_time"],
|
||||
target_mode="next_token",
|
||||
return_output=True,
|
||||
)
|
||||
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)
|
||||
readout_out = readout(
|
||||
hidden=hidden,
|
||||
time_seq=batch["time_seq"],
|
||||
padding_mask=batch["padding_mask"],
|
||||
readout_mask=batch["readout_mask"]
|
||||
hidden=model_out.hidden,
|
||||
time_seq=model_out.time_seq,
|
||||
padding_mask=model_out.padding_mask,
|
||||
readout_mask=targets["readout_mask"]
|
||||
if args.readout_name == "same_time_group_end"
|
||||
else None,
|
||||
)
|
||||
@@ -202,17 +340,17 @@ def compute_next_step_loss(
|
||||
if args.target_mode == "delphi2m":
|
||||
loss, parts = criterion(
|
||||
logits=logits,
|
||||
target_events=batch["target_event_seq"],
|
||||
target_times=batch["target_time_seq"],
|
||||
current_times=batch["time_seq"],
|
||||
target_events=targets["target_event_seq"],
|
||||
target_times=targets["target_time_seq"],
|
||||
current_times=model_out.time_seq,
|
||||
padding_mask=readout_out.readout_mask,
|
||||
return_components=True,
|
||||
)
|
||||
else:
|
||||
loss, parts = criterion(
|
||||
logits=logits,
|
||||
target_multi_hot=batch["target_multi_hot"],
|
||||
target_dt_unique=batch["target_dt_unique"],
|
||||
target_multi_hot=targets["target_multi_hot"],
|
||||
target_dt_unique=targets["target_dt_unique"],
|
||||
readout_mask=readout_out.readout_mask,
|
||||
return_components=True,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user