From 4181ead03a6d2c39b3a5397bcb4ca565244fed5f Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Thu, 16 Oct 2025 15:57:27 +0800 Subject: [PATCH] Refactor: Improve attention mechanism and early stopping - Refactor the self-attention mechanism in `models.py` to use `nn.MultiheadAttention` for better performance and clarity. - Disable early stopping check during warmup epochs in `train.py` to improve training stability. --- models.py | 59 ++++++++++----------------------------------- train.py | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 80 insertions(+), 51 deletions(-) diff --git a/models.py b/models.py index f03ebe7..32021b2 100644 --- a/models.py +++ b/models.py @@ -2,57 +2,15 @@ import torch import torch.nn as nn from torch.nn import functional as F from typing import Tuple -import math - -class CausalSelfAttention(nn.Module): - """ - A vanilla multi-head masked self-attention layer with a projection at the end. - """ - - def __init__(self, n_embd: int, n_head: int, pdrop: float): - super().__init__() - assert n_embd % n_head == 0 - # key, query, value projections for all heads - self.c_attn = nn.Linear(n_embd, 3 * n_embd) - # output projection - self.c_proj = nn.Linear(n_embd, n_embd) - # regularization - self.attn_dropout = nn.Dropout(pdrop) - self.resid_dropout = nn.Dropout(pdrop) - self.n_head = n_head - self.n_embd = n_embd - - def forward(self, x: torch.Tensor, custom_mask: torch.Tensor) -> torch.Tensor: - B, L, D = x.size() # batch size, sequence length, embedding dimensionality (n_embd) - - # calculate query, key, values for all heads in batch and move head forward to be the batch dim - q, k, v = self.c_attn(x).split(self.n_embd, dim=2) - k = k.view(B, L, self.n_head, D // self.n_head).transpose(1, 2) # (B, nh, L, hs) - q = q.view(B, L, self.n_head, D // self.n_head).transpose(1, 2) # (B, nh, L, hs) - v = v.view(B, L, self.n_head, D // self.n_head).transpose(1, 2) # (B, nh, L, hs) - - # causal self-attention; Self-attend: (B, nh, L, hs) x (B, nh, hs, L) -> (B, nh, L, L) - att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) - - # Apply the time-based causal mask - att = att.masked_fill(custom_mask.unsqueeze(1) == 0, float('-inf')) - - att = F.softmax(att, dim=-1) - att = self.attn_dropout(att) - y = att @ v # (B, nh, L, L) x (B, nh, L, hs) -> (B, nh, L, hs) - y = y.transpose(1, 2).contiguous().view(B, L, D) # re-assemble all head outputs side by side - - # output projection - y = self.resid_dropout(self.c_proj(y)) - return y class Block(nn.Module): """ an unassuming Transformer block """ def __init__(self, n_embd: int, n_head: int, pdrop: float): super().__init__() + self.n_head = n_head self.ln_1 = nn.LayerNorm(n_embd) - self.attn = CausalSelfAttention(n_embd, n_head, pdrop) + self.attn = nn.MultiheadAttention(n_embd, n_head, dropout=pdrop, batch_first=True) self.ln_2 = nn.LayerNorm(n_embd) self.mlp = nn.ModuleDict(dict( c_fc = nn.Linear(n_embd, 4 * n_embd), @@ -62,9 +20,16 @@ class Block(nn.Module): )) m = self.mlp self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x)))) # MLP forward + self.resid_dropout = nn.Dropout(pdrop) def forward(self, x: torch.Tensor, custom_mask: torch.Tensor) -> torch.Tensor: - x = x + self.attn(self.ln_1(x), custom_mask=custom_mask) + normed_x = self.ln_1(x) + + attn_mask = ~custom_mask + attn_mask = attn_mask.repeat_interleave(self.n_head, dim=0) + + attn_output, _ = self.attn(normed_x, normed_x, normed_x, attn_mask=attn_mask, need_weights=False) + x = x + self.resid_dropout(attn_output) x = x + self.mlpf(self.ln_2(x)) return x @@ -190,13 +155,13 @@ class TimeAwareGPT2(nn.Module): # 5. Generate attention mask # The attention mask combines two conditions: - # a) Time-based causality: A token i can attend to a token j only if time_seq[j] < time_seq[i]. + # a) Time-based causality: A token i can attend to a token j only if time_seq[j] <= time_seq[i]. # b) Padding mask: Do not attend to positions where the event token is 0. # a) Time-based causal mask t_i = time_seq.unsqueeze(-1) # (B, L, 1) t_j = time_seq.unsqueeze(1) # (B, 1, L) - time_mask = (t_j < t_i) + time_mask = (t_j <= t_i) # b) Padding mask (prevents attending to key positions that are padding) padding_mask = (event_seq != 0).unsqueeze(1) # Shape: (B, 1, L) diff --git a/train.py b/train.py index fbcfb2a..92fa286 100644 --- a/train.py +++ b/train.py @@ -5,6 +5,7 @@ from torch.utils.data import DataLoader import numpy as np import math import tqdm +import matplotlib.pyplot as plt from models import TimeAwareGPT2, CombinedLoss from utils import PatientEventDataset @@ -14,7 +15,7 @@ class TrainConfig: # Data parameters train_data_path = 'ukb_real_train.bin' val_data_path = 'ukb_real_val.bin' - block_length = 256 # Sequence length + block_length = 24 # Sequence length # Model parameters n_embd = 256 @@ -76,6 +77,11 @@ def main(): # --- 3. Training Loop --- best_val_loss = float('inf') patience_counter = 0 + + # Lists to store losses + train_losses_ce, train_losses_surv, train_losses_total = [], [], [] + val_losses_ce, val_losses_surv, val_losses_total = [], [], [] + print("Starting training...") for epoch in range(config.max_epoch): # --- Learning Rate Scheduling --- @@ -120,6 +126,9 @@ def main(): avg_train_loss_ce = train_loss_ce_acc / train_steps avg_train_loss_surv = train_loss_surv_acc / train_steps + train_losses_ce.append(avg_train_loss_ce) + train_losses_surv.append(avg_train_loss_surv) + train_losses_total.append(avg_train_loss_ce + avg_train_loss_surv) # --- Validation Phase --- model.eval() @@ -147,6 +156,9 @@ def main(): avg_val_loss_ce = val_loss_ce_acc / val_steps avg_val_loss_surv = val_loss_surv_acc / val_steps total_val_loss = avg_val_loss_ce + avg_val_loss_surv + val_losses_ce.append(avg_val_loss_ce) + val_losses_surv.append(avg_val_loss_surv) + val_losses_total.append(total_val_loss) print(f"Epoch {epoch+1} Summary: \n" f" Train Loss: {avg_train_loss_ce + avg_train_loss_surv:.4f} (CE: {avg_train_loss_ce:.4f}, Surv: {avg_train_loss_surv:.4f})\n" @@ -157,14 +169,66 @@ def main(): if total_val_loss < best_val_loss: best_val_loss = total_val_loss patience_counter = 0 - print(f"Validation loss improved to {best_val_loss:.4f}. Resetting patience.") + print(f"Validation loss improved to {best_val_loss:.4f}. Saving checkpoint...") + torch.save(model.state_dict(), 'best_model_checkpoint.pt') else: - patience_counter += 1 - print(f"Validation loss did not improve. Patience: {patience_counter}/{config.early_stopping_patience}") + if epoch >= config.warmup_epochs: + patience_counter += 1 + print(f"Validation loss did not improve. Patience: {patience_counter}/{config.early_stopping_patience}") if patience_counter >= config.early_stopping_patience: print("\nEarly stopping triggered due to no improvement in validation loss.") break + # --- Save Best Model at the End --- + if best_val_loss != float('inf'): + print(f"\nTraining finished. Loading best model from checkpoint with validation loss {best_val_loss:.4f}.") + model.load_state_dict(torch.load('best_model_checkpoint.pt')) + print("Saving final best model to best_model.pt") + torch.save(model.state_dict(), 'best_model.pt') + else: + print("\nTraining finished. No best model to save as validation loss never improved.") + + # --- Plot and Save Loss Curves --- + num_epochs = len(train_losses_total) + epochs = range(1, num_epochs + 1) + + plt.figure(figsize=(18, 5)) + + # Plot CE Loss + plt.subplot(1, 3, 1) + plt.plot(epochs, train_losses_ce, label='Train CE') + plt.plot(epochs, val_losses_ce, label='Val CE') + plt.title('Cross-Entropy Loss') + plt.xlabel('Epochs') + plt.ylabel('Loss') + plt.legend() + plt.grid(True) + + # Plot Survival Loss + plt.subplot(1, 3, 2) + plt.plot(epochs, train_losses_surv, label='Train Survival') + plt.plot(epochs, val_losses_surv, label='Val Survival') + plt.title('Survival Loss') + plt.xlabel('Epochs') + plt.ylabel('Loss') + plt.legend() + plt.grid(True) + + # Plot Total Loss + plt.subplot(1, 3, 3) + plt.plot(epochs, train_losses_total, label='Train Total') + plt.plot(epochs, val_losses_total, label='Val Total') + plt.title('Total Loss') + plt.xlabel('Epochs') + plt.ylabel('Loss') + plt.legend() + plt.grid(True) + + plt.tight_layout() + plt.savefig('loss_curves.png') + print("\nLoss curves saved to loss_curves.png") + + if __name__ == '__main__': main()