diff --git a/evaluate_delphi.py b/evaluate_delphi.py deleted file mode 100644 index cdfc68f..0000000 --- a/evaluate_delphi.py +++ /dev/null @@ -1,1079 +0,0 @@ -#!/usr/bin/env python -# coding: utf-8 - -# # Delphi inference and evaluation -# -# Welcome to the Delphi evaluation notebook! To run this notebook, you need to have the Delphi model checkpoint. -# Refer to the README for instructions on how to train it on synthetic (or real!) data. -# -# Here, we show how to work with the model, load data and perform inference. We also reproduce some of the figures from the paper. -# -# Note that this notebook in its current state was executed using the original Delphi checkpoint and full UK biobank data. The small synthetic dataset we provide in this repository may not be sufficient to reproduce all the results. -# -# On Mac M1 Pro CPU using the synthetic dataset, the notebook takes ~10 minutes to run. -# -# ## Table of contents -# -# 1. Loading model -# 2. Data: structure and loading -# 3. Inference -# 4. Prediction of future disease rates -# 5. Checking calibration of predicted rates -# 6. Evaluation of AUC -# 7. Looking into attention patterns -# 8. Token embedding UMAP -# -# - - - -import os -import torch -from model import DelphiConfig, Delphi -from tqdm import tqdm -import pandas as pd -import numpy as np -import textwrap -import matplotlib.pyplot as plt -get_ipython().run_line_magic('config', "InlineBackend.figure_format='retina'") - -plt.rcParams['figure.facecolor'] = 'white' -plt.rcParams.update({'axes.grid': True, - 'grid.linestyle': ':', - 'axes.spines.bottom': False, - 'axes.spines.left': False, - 'axes.spines.right': False, - 'axes.spines.top': False}) -plt.rcParams['figure.dpi'] = 72 -plt.rcParams['pdf.fonttype'] = 42 - -#Green -light_male = '#BAEBE3' -normal_male = '#0FB8A1' -dark_male = '#00574A' - - -#Purple -light_female = '#DEC7FF' -normal_female = '#8520F1' -dark_female = '#7A00BF' - - -delphi_labels = pd.read_csv('delphi_labels_chapters_colours_icd.csv') - - - - -# Delphi is capable of predicting the disease risk for 1,256 diseases from ICD-10 plus death. -# For illustrative purposes, some of the plots will focus on a subset of 10 selected diseases - the same subset in used in the Delphi paper. - -diseases_of_interest = [46, 95, 1168, 1188, 374, 214, 305, 505, 603, 1269] -delphi_labels.iloc[diseases_of_interest][['name', 'ICD-10 Chapter (short)']] - - -# ## Load model - - - -out_dir = 'Delphi-2M' -device = 'cpu' # examples: 'cpu', 'cuda', 'cuda:0', 'mps', etc. -dtype ='float32' #'bfloat16' # 'float32' or 'bfloat16' or 'float16' -dtype = {'float32': torch.float32, 'float64': torch.float64, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype] -seed = 1337 - -torch.manual_seed(seed) -torch.cuda.manual_seed(seed) - - - - -ckpt_path = os.path.join(out_dir, 'ckpt.pt') -checkpoint = torch.load(ckpt_path, map_location=device) -conf = DelphiConfig(**checkpoint['model_args']) -model = Delphi(conf) -state_dict = checkpoint['model'] -model.load_state_dict(state_dict) - -model.eval() -model = model.to(device) - -checkpoint['model_args'] - - - - -# Let's try to use the loaded model to extrapolate a partial health trajectory. - -example_health_trajectory = [ - ('Male', 0), - ('B01 Varicella [chickenpox]',2), - ('L20 Atopic dermatitis',3), - ('No event', 5), - ('No event', 10), - ('No event', 15), - ('No event', 20), - ('G43 Migraine', 20), - ('E73 Lactose intolerance',21), - ('B27 Infectious mononucleosis',22), - ('No event', 25), - ('J11 Influenza, virus not identified',28), - ('No event', 30), - ('No event', 35), - ('No event', 40), - ('Smoking low', 41), - ('BMI mid', 41), - ('Alcohol low', 41), - ('No event', 42), -] -example_health_trajectory = [(a, b * 365.25) for a,b in example_health_trajectory] - - - - -max_new_tokens = 100 - -name_to_token_id = {row[1]['name']: row[1]['index'] for row in delphi_labels.iterrows()} - -events = [name_to_token_id[event[0]] for event in example_health_trajectory] -events = torch.tensor(events, device=device).unsqueeze(0) -ages = [event[1] for event in example_health_trajectory] -ages = torch.tensor(ages, device=device).unsqueeze(0) - -res = [] -with torch.no_grad(): - y,b,_ = model.generate(events, ages, max_new_tokens, termination_tokens=[1269]) - # Convert model outputs to readable format - events_data = zip(y.cpu().numpy().flatten(), b.cpu().numpy().flatten()/365.) - - print('Input trajectory:') - for i, (event_id, age_years) in enumerate(events_data): - if i == len(example_health_trajectory): - print('=====================') - print('Generated trajectory:') - event_name = delphi_labels.loc[event_id, 'name'] - print(f"{age_years:2.1f}: {event_name}") - - -# ## Load data -# -# The data include: -# -# Tokens include: -# - 1,257 different ICD-10 level 3 disease codes (e.g., E11 for Type 2 diabetes) -# - 9 lifestyle tokens (alcohol, smoking, BMI - each with 3 levels) -# - 2 sex tokens (male/female) -# -# The following technical tokens are added in the `get_batch` function: -# - 1 "no event" padding token -# - 1 non-informative padding token -# -# No-event tokens eliminate long time intervals without tokens, which are typical for younger ages, when people generally have fewer diseases and therefore less medical records. Transformers predict the text token probability distribution only at the time of currently observed tokens, hence, no-event tokens can also be inserted during inference to obtain the predicted disease risk at any given time of interest. - - - -from utils import get_batch, get_p2i - -train = np.fromfile('data/ukb_simulated_data/train.bin', dtype=np.uint32).reshape(-1,3) -val = np.fromfile('data/ukb_simulated_data/val.bin', dtype=np.uint32).reshape(-1,3) - -train = np.fromfile('../data/ukb_real_data/train.bin', dtype=np.uint32).reshape(-1,3) -val = np.fromfile('../data/ukb_real_data/val.bin', dtype=np.uint32).reshape(-1,3) - -train_p2i = get_p2i(train) # mapping trajectory id to its position in the dataset -val_p2i = get_p2i(val) - -dataset_subset_size = 2000 # len(val_p2i) # can be set to smaller number (e.g. 2048) for a quick run - - -# ## Calibration of predicted times - - - -# Fetch a bit of data and calculate future disease rates from it -d = get_batch(range(256), val, val_p2i, select='left', padding='random', block_size=128, device=device) -with torch.no_grad(): - p = model(*d)[0].cpu().detach().numpy().squeeze() -t = (d[3]-d[1])[:,:].cpu().numpy().squeeze() - - - - -from scipy.special import logsumexp - -# Calculate expected waiting times from model predictions using competing exponentials theory -# The expected time until any event occurs is 1/sum(λᵢ) = 1/exp(logsumexp(logits)) -# logsumexp provides numerical stability vs. calculating exp(logits) directly - -# Let's see how the predicted waiting times compare to the observed waiting times - -plt.figure(figsize=(4, 4)) -# Calculate expected time to next token (inverse of hazard rate) -expected_t = 1/np.exp(logsumexp(p, axis=-1)) - -# Define bin width for logarithmic binning -delta_log_t = 0.1 -log_range = np.arange(1.75, 4, delta_log_t) - -# Calculate average observed time for each logarithmic bin -observed_t = [] -for i in log_range: - # Create mask for current bin and valid times - bin_mask = (expected_t > 10**i) & (expected_t <= 10**(i+delta_log_t)) & (t > 0) - # Calculate mean for this bin - bin_mean = t[bin_mask].mean() if bin_mask.sum() > 0 else np.nan - observed_t.append(bin_mean) -plt.axes().set_aspect('equal') -plt.scatter(expected_t, t+0.5, marker=".", c='lightgrey', rasterized=True) -plt.xlabel('Expected days to next token') -plt.ylabel('Observed days to next token') -plt.plot(10**(np.arange(1.75,4,delta_log_t)+delta_log_t/2.),observed_t, label='average') -plt.yscale('log') -plt.xscale('log') -plt.legend() -plt.xlim(1,2e3) -plt.ylim(1,2e3) -plt.plot([0,1],[0,1], transform = plt.gca().transAxes, c='k' , ls=(0, (5, 5)), linewidth=0.7) - -plt.gca().tick_params(length=1.15, width=0.3, labelsize=8, grid_alpha=1, grid_linewidth=0.45, grid_linestyle=':') -plt.gca().tick_params(length=1.15, width=0.3, labelsize=8, grid_alpha=0.0, grid_linewidth=0.35, which='minor') - - -# ## Incidence - - - -## Load large chunk of data -# `get_batch` function reads health trajectories from the dataset for requested indices of individuals -# it also packes the trajectories into batches of size `block_size`, padding with padding token if needed -# finally, it randomly add no event tokens and returns a tuple of tensors: -# - `d[0]`: diseases -# - `d[1]`: corresponding age -# - `d[2]`: disease labels (same as `d[0]`, but shifted by 1) -# - `d[3]`: label age (same as `d[1]`, but shifted by 1) -subset_size = 10_000 -d = get_batch(range(subset_size), val, val_p2i, - select='left', block_size=128, - device=device, padding='random') - - - - -# 2 is female token, 3 is male token - -is_male = (d[0] == 3).any(axis=1).cpu().numpy() -is_female = (d[0] == 2).any(axis=1).cpu().numpy() -has_gender = is_male | is_female - - - - -# lets split the large data chanks to smaller batches and calculate the logits for the whole dataset -p = [] -model.to(device) -batch_size = 256 -subset_size = min(dataset_subset_size, 10_000) -with torch.no_grad(): - for d_batch in tqdm(zip(*map(lambda x: torch.split(x, batch_size), d)), total=d[0].shape[0]//batch_size+1): - p.append(model(*d_batch)[0].cpu().detach()) -p = torch.vstack(p) - -d = [d_.cpu() for d_ in d] - - -# ### Age-sex incidence baseline - - - -# calculate disease incidence rates for each disease, given age and sex - -females = train[np.isin(train[:,0], train[train[:,2]==1,0])] -males = train[np.isin(train[:,0], train[train[:,2]==2,0])] -n_females = (train[:,2]==1).sum() -n_males = (train[:,2]==2).sum() - -unique_male_indices = np.where(males[:-1,0] != males[1:,0])[0] -unique_female_indices = np.where(females[:-1,0] != females[1:,0])[0] - -def calc_age_distribution(data, unique_indices): - ages = np.maximum(40, np.round(data[unique_indices, 1]/365.25) + 1) - counts = np.histogram(ages, np.arange(100))[0] - cumulative = -np.cumsum(counts) - return cumulative - cumulative[-1] - -n_males = calc_age_distribution(males, unique_male_indices) -n_females = calc_age_distribution(females, unique_female_indices) - -ukb_condition = (males[:,2] > 2) & (males[:,2] <= 4) -males_in_ukb = np.cumsum(np.histogram((males[ukb_condition, 1]/365.25).astype('int'), np.arange(100))[0]) - -ukb_condition = (females[:,2] > 2) & (females[:,2] <= 4) -females_in_ukb = np.cumsum(np.histogram((females[ukb_condition, 1]/365.25).astype('int'), np.arange(100))[0]) - - -# ### Modelled age-incidence - -# ### Selected diseases -# -# Delphi predicts the disease rate. Let's plot Delphi-predicted rates for the selected diseases vs age and compare them with the reported incidence rates (population averages from UKB). -# -# Shown in the graph are population average disease rates (solid lines), Delphi-predicted rates for arbitrary timepoints (pale dots) and Delphi-predicted rates for the penultimate step before disease (bright dots). -# -# Bright dots are often located above the population average rates, which indicates that Delphi correctly captures the elevated disease risk for such participants. - - - -def plot_age_incidence(ix, d, p, highlight_idx=0): - """ - Plot age-specific incidence rates for selected diseases. - - Parameters: - ----------- - ix : list or array - Indices of diseases to plot - d : tuple - Tuple containing disease data: - - d[0]: disease history - - d[1]: age information - - d[2]: disease labels - - d[3]: additional time information - p : torch.Tensor - Probability tensor from Delphi model - highlight_idx : int, default=0 - Index of the case to highlight in the trajectory plot - - Returns: - -------- - None - Displays the plot but doesn't return any value - """ - # Calculate number of rows needed based on number of diseases - n_rows = (len(ix) - 1) // 5 + 1 - fig, ax = plt.subplots(n_rows, 5, figsize=(18, 3 * n_rows), sharex=False, sharey=True) - axf = ax.ravel() - - for i, k in enumerate(ix): - # Prepare data - x = d[1][:,:].detach().numpy() / 365.25 - y = np.exp(p.detach().numpy()[:,:,k]) * 365.25 - y = 1 - np.exp(-y) - - # Filter for cases without prior disease - no_prior_disease = ~np.isin(d[0], k).any(axis=1) - sub_sample = np.random.randint(0, len(x[has_gender * no_prior_disease].ravel()), 5000) - - # Plot background points - axf[i].scatter( - x[has_gender * no_prior_disease].ravel()[sub_sample], - y[has_gender * no_prior_disease].ravel()[sub_sample], - marker='.', - c=np.repeat(np.array([light_female, light_male])[0+is_male[has_gender * no_prior_disease]], x.shape[1]).ravel()[sub_sample], - edgecolors='white', - s=50, - label='Delphi, all time steps', - rasterized=True - ) - - # Plot points just before disease onset - has_k = np.where(d[2].detach().numpy()[has_gender] == k)[0] - before_k = d[2].detach().numpy()[has_gender].ravel() == k - axf[i].scatter( - x[has_gender].ravel()[before_k], - y[has_gender].ravel()[before_k], - marker='.', - c=np.array([dark_female, dark_male])[0+is_male[has_gender][has_k]], - edgecolors='white', - s=50, - label='Delphi, penultimate step before disease', - rasterized=True - ) - - # Plot selected case trajectory - j = np.where(np.isin(d[2], k).any(axis=1))[0][highlight_idx] - j0 = np.where(x[j] >= 0)[0][0] - jk = np.where(d[2][j,:].detach().numpy() == k)[0][0] - axf[i].plot( - x[j][j0:jk+1], - y[j][j0:jk+1], - ds='steps-post', - c='k', - ls="-", - marker='.', - markersize=8, - markeredgecolor='white', - markerfacecolor='k', - label='selected case' - ) - axf[i].scatter(x[j][jk], y[j][jk], marker='.', s=200, edgecolors='white', c='k', zorder=3) - - # Plot reported incidence rates - h, x = np.histogram(females[females[:,2]==k-1,1]/365.25, np.arange(100)) - axf[i].stairs(h/n_females, x, color=normal_female, lw=2, label='reported incidence, female') - - h, x = np.histogram(males[males[:,2]==k-1,1]/365.25, np.arange(100)) - axf[i].stairs(h/n_males, x, color=normal_male, lw=2, label='reported incidence, male') - - # Set plot properties - axf[i].set_ylim((1e-5, 1)) - axf[i].set_xlim((0, 80)) - axf[i].set_yscale('log') - axf[i].set_title("\n".join(textwrap.wrap(delphi_labels.loc[k,'name'], width=30)), - verticalalignment='top', fontsize=10, fontweight='bold') - - if i % ax.shape[1] == 0: - axf[i].set_ylabel('Rate per year') - - if i // ax.shape[1] == ax.shape[0] - 1: - axf[i].set_xlabel('Age') - - if i == len(ix) - 1: - axf[i].legend(loc='center left', bbox_to_anchor=(1.05, 0.5)) - - - - -plot_age_incidence(diseases_of_interest, d, p, highlight_idx=0) -plt.gcf().tight_layout(h_pad=0.5) -plt.show() - - -# ### Calibration -# Delphi predicts the absolude disease rate. In this section, we evaluate how well Delphi's predictions match the observed disease rates in the UKB dataset. -# -# The strategy for calibration assesment is the following: -# 1. Run Delphi for the entire dataset -# 2. Stratify all participants into sex & age groups -# 3. For each age-sex group, split all participants into bins according to the predicted disease risk -# 4. For each bin, calculate the observed and predicted disease rates -# 5. Plot the calibration curve -# - - - -def auc(x1, x2): - "Calcualte AUC, given x1 vector of disease risks for cases and x2 vector of disease risks for controls" - n1 = len(x1) - n2 = len(x2) - R1 = np.concatenate([x1,x2]).argsort().argsort()[:n1].sum() + n1 - U1 = n1*n2 + 0.5*n1*(n1+1) - R1 - if n1 == 0 or n2 == 0: - return np.nan - return U1 / n1 / n2 - - - - -d100k = get_batch(range(dataset_subset_size), val, val_p2i, - select='left', block_size=128, - device=device, padding='random') - - - - -p100k = [] -model.to(device) -batch_size=256 -with torch.no_grad(): - for dd in tqdm(zip(*map(lambda x: torch.split(x, batch_size), d100k)), total=d100k[0].shape[0]//batch_size+1): - p100k.append(model(*[x.to(device) for x in dd])[0].cpu().detach()[:,:,diseases_of_interest].numpy()) -p100k = np.vstack(p100k) - - - - -import scipy -import warnings - -def plot_calibration(disease_idx, data, logits, offset = 365.25, age_groups=range(45,85,5), n_samples=3, calibration = 'bins', binning='power', bins=10**np.arange(-6.,1.5,.5)): - """ - Plot calibration curves for disease predictions. - The selection of controls and cases in this function happens in the following way: - - For cases, we can just select the predicted disease rates corresponding to the moment before - occurrence of the disease (given the offset). - - For controls, there isn't a particular moment in time when the disease occurs, so we just - sample random moments of the trajectory. - - Args: - disease_idx: Index of disease in the dataset - data: Tuple of tensors containing input data (tokens, times, targets, target_times) - logits: Model prediction logits - offset: Time offset in days (default: 365.25) - age_groups: Range of age groups to analyze (default: range(45,85,5)) - n_samples: Number of samples (default: 3) - calibration: Calibration method, 'bins' or other (default: 'bins') - binning: Binning method, 'power' or 'deciles' (default: 'power') - bins: Bin edges for power binning (default: 10**np.arange(-6.,1.5,.5)) - - Returns: - List of calibration data for each age group - """ - - l = len(age_groups) - age_step = age_groups[1] - age_groups[0] - - fig, ax = plt.subplots(2, l, figsize=(20/8*l,3), sharex=True, sharey=False, height_ratios=[1, .5]) - # Indices of cases - wk = np.where(data[2].detach().numpy()==disease_idx) - - if len(wk[0])<2: - return np.repeat(np.nan, l) - - # Indices of controls - wc = np.where(data[2].detach().numpy()!=disease_idx) - - c_sub = range(wc[0].shape[0]) - wall = (np.concatenate([wk[0], wc[0][c_sub]]), np.concatenate([wk[1], wc[1][c_sub]])) - - pred_idx = (data[1][wall[0]] <= data[3][wall].reshape(-1,1) - offset).sum(1) -1 - z = data[1].detach().numpy()[(wall[0], pred_idx)] - z = z[pred_idx != -1] - - zk = data[3].detach().numpy()[wall] # Target ages, cases and controls - zk = zk[pred_idx != -1] - - x = np.exp(logits[(wall[0], pred_idx)]) * 365.25 # Disease rates - x = x[pred_idx != -1] - x = 1 - np.exp(-x * age_step) - - wk = (wk[0][pred_idx[:len(wk[0])] != -1], wk[1][pred_idx[:len(wk[0])] != -1]) - p_idx = wall[0][pred_idx!=-1] - - out = [] - - for i,aa in enumerate(age_groups): - ax_cal = ax[0, i] # Calibration plot - ax_box = ax[1, i] # Boxplot - - a = np.logical_and(z / 365.25 >= aa, z / 365.25 < aa+ age_step) - a *= zk - z < 365.25 #* age_step - a *= np.isin(np.arange(a.shape[0]),np.unique(p_idx * a, return_index=True)[1]) # Mask duplicated people in age bracket - ax_box.boxplot((x[len(wk[0]):][a[len(wk[0]):]], x[:len(wk[0])][a[:len(wk[0])]]), vert=False, sym='.', widths=.5, whis=(5,95), - flierprops=dict(marker='.', markeredgecolor='white', markerfacecolor='k')) - ax_box.set_xscale('log') - ax_box.set_xlim((1e-5, 1)) - ax_box.set_yticks((1,2), ['','']) - if i==0: - ax_cal.set_ylabel('Observed rate [1/yr]') - ax_box.set_yticks((1,2), (f'{["Healthy","Alive"][disease_idx==1268]}',f'{["Diseased","Deceased"][disease_idx==1268]}')) - y = auc(x[len(wk[0]):][a[len(wk[0]):]], x[:len(wk[0])][a[:len(wk[0])]]) - - foo =["dis'd","dec'd"] - ax_cal.text(0,.9, s= f'{len(x[len(wk[0]):][a[len(wk[0]):]])} {["healthy","alive"][disease_idx==1268]}\n{len(x[:len(wk[0])][a[:len(wk[0])]])} {foo[disease_idx==1268]}', - transform=ax_cal.transAxes, va='top') - ax_box.text(0.5, .8, s = f"AUC={y:.2}", transform=ax_box.transAxes, va='center', ha='center') - ax_box.set_xlabel('Predicted rate [1/yr]') - ax_box.set_ylim((0.5,3.5)) - ax_cal.text(0.5, 1, s = f'{aa}-{aa+age_step}yr', transform=ax_cal.transAxes, va='bottom', ha='center', weight='bold') - - - xa = x[a] - ya = np.concatenate([np.ones(len(wk[0])), np.zeros(x.shape[0] - len(wk[0]))])[a] - - if len(xa) == 0: - continue - - - if calibration == 'bins': - if binning == 'deciles': - bins = np.quantile(xa, np.arange(0,1.05,0.05)) - else: - bins = bins - bin_masks = [np.logical_and(xa > bins[b-1], xa <= bins[b]) for b in range(1,len(bins))] - # np.errstate doesn't suppress RuntimeWarning, need to use warnings module - with warnings.catch_warnings(): - warnings.filterwarnings('ignore', category=RuntimeWarning) - pred = np.array([xa[bin_mask].mean() for bin_mask in bin_masks]) - obs = np.array([ya[bin_mask].mean() for bin_mask in bin_masks]) - ci = np.array([scipy.stats.beta(0.1 + ya[bin_mask].sum(), 0.1 + (1-ya[bin_mask]).sum()).ppf([0.025,0.975]) for bin_mask in bin_masks]) - ax_cal.scatter(pred, obs + 1e-5, marker='.', c='k') - for j,pr in enumerate(pred): - if not np.isnan(obs[j]): - ax_cal.plot( np.repeat(pr,2),ci[j], c='k', lw=.5, ls=":") - wgt = np.array([[ya[bin_mask].sum(),bin_mask.sum()] for bin_mask in bin_masks]) - out.append([pred, obs, ci, wgt]) - else: - o = np.argsort(xa) - ax_cal.plot(xa[o], ya[o]/(ya.sum() - np.cumsum(ya[o]))/age_step, ds='steps') - out.append(np.nan) - - ax_cal.set_box_aspect(1) - ax_cal.scatter(xa.mean(), ya.mean(), c='r', ec='w') - ax_cal.set_yscale('log') - ax_cal.set_xscale('log') - ax_cal.set_ylim((1e-5, 1)) - ax_cal.set_xlim((1e-5, 1)) - ax_cal.plot([0, 1], [0, 1], transform=ax_cal.transAxes, lw=.5, c='k', ls="--") - - return out - - - - -out = [] -plt.rcParams.update({'figure.max_open_warning': 0}) - -is_male = (d100k[0] == 2).sum(1)>0 -is_female = (d100k[0] == 3).sum(1)>0 - -calibration_inputs = { - 'male': {'data': [d_[is_male].cpu() for d_ in d100k], 'logits': p100k[is_male.cpu()]}, - 'female': {'data': [d_[is_female].cpu() for d_ in d100k], 'logits': p100k[is_female.cpu()]} -} - -for j, k in enumerate(diseases_of_interest): - disease_name = delphi_labels.loc[k, 'name'] - for sex in ['male', 'female']: - out.append(plot_calibration(k,calibration_inputs[sex]['data'], calibration_inputs[sex]['logits'][..., j], age_groups=np.arange(40,80,5), offset=0.1)) - # plt.tight_layout() - plt.suptitle(disease_name, fontsize=10, weight='bold', ha='left', x=0.1, y=1.05) - plt.show() - - - - -# the same calibrations curves as above, but a more compact version - -from matplotlib.colors import LinearSegmentedColormap - -fig, ax = plt.subplots(2,5,figsize=(18,6), sharex=True, sharey=True) -ax=ax.ravel() -for i, calibration_data in enumerate(out): - j = i // 2 # Females and males - is_male = i % 2 == 0 - for age_bracket_idx, cal in enumerate(calibration_data): - if not isinstance(cal, list) and np.isnan(cal).all(): - continue - intensity = 0.15+age_bracket_idx/9*.55 - cmap = LinearSegmentedColormap.from_list('cmap',list(zip([0,.5,1],[['white',normal_male,dark_male], ['white',normal_female,dark_female]][is_male]))) - ax[j].plot(cal[0], cal[1]**2/cal[1], label=f"{40+5*age_bracket_idx}-{40+5*(age_bracket_idx+1)}yrs", c=cmap(intensity)) - ax[j].set_yscale('log') - ax[j].set_xscale('log') - ax[j].set_xlim(5e-5, 0.5) - ax[j].set_ylim(5e-5, 0.5) - ax[j].plot([0, 1], [0, 1], transform=ax[j].transAxes, lw=.5, c='k', ls="--") - ax[j].set_title("\n".join(textwrap.wrap(delphi_labels['name'].iloc[diseases_of_interest[j]],30)),verticalalignment='top', size=10, weight='bold') - ax[j].set_xlabel('Model rate [1/yr]') - ax[j].set_ylabel('Observed rate [1/yr]') - ax[j].label_outer() - -ax[j].legend(loc='lower left', ncol=2, bbox_to_anchor=(1.05, 0)) - -plt.gcf().tight_layout(h_pad=0.5) -plt.show() - - -# ## AUC of disease prediction -# -# For evaluation of the AUCs, we use a similar strategy as in the calibration assesment. -# -# 1. Run Delphi for the entire dataset -# 2. Stratify all participants into sex & age groups - this is needed to regress out the "baseline" disease rate change - it's not that difficult to predict that older people have higher disease risk (for most diseases) -# 3. For each age-sex group, select controls and cases -# 4. Calculate the AUC using Delphi disease rates as predictors -# 5. (Optional) Use DeLong's method (recommended) or bootstrap to calculate the variance of the AUC - - - -from evaluate_auc import get_calibration_auc, evaluate_auc_pipeline - -offset = 0.1 -pred_idx_precompute = (d100k[1][:, :, np.newaxis] < d100k[3][:, np.newaxis, :] - offset).to(torch.float32).sum(1) - 1 # float comvertion saves memory (somehow) -pred_idx_precompute = pred_idx_precompute.to(torch.int32) - -is_male = (d100k[0] == 2).sum(1)>0 -is_female = (d100k[0] == 3).sum(1)>0 - -auc_inputs = { - 'male': { - 'data': [d_[is_male].cpu().numpy() for d_ in d100k], - 'logits': p100k[is_male.cpu()], - 'pred_idx_precompute': pred_idx_precompute[is_male].cpu().numpy() - }, - 'female': { - 'data': [d_[is_female].cpu().numpy() for d_ in d100k], - 'logits': p100k[is_female.cpu()], - 'pred_idx_precompute': pred_idx_precompute[is_female].cpu().numpy() - } -} - - - - -all_aucs = [] - -for disease_idx_batch, disease_idx in tqdm(enumerate(diseases_of_interest), total=len(diseases_of_interest)): - for sex in ['male', 'female']: - - out = get_calibration_auc( - disease_idx_batch, - disease_idx, - auc_inputs[sex]['data'], - auc_inputs[sex]['logits'], - age_groups=np.arange(40, 80, 5), - offset=offset, - precomputed_idx=auc_inputs[sex]['pred_idx_precompute'], - use_delong=True, - ) - - if out is None: - continue - for out_item in out: - out_item["sex"] = sex - all_aucs.append(out_item) - - - - -# this df contains AUC calculations for all diseases, sexes, and age groups -# to get the AUC for a specific disease, one needs to aggregate over the age groups and sexes -# however, while for the mean AUC this is straightforward, it's a bit more complicated for the variance -# as Delong's method provides confidence intervals as the form of variance of a normal distribution -# we can use the closed form of the variance of the mean of a normal distributions to get the variance of the AUC -# let's defile a custom funciton for it - -def aggregate_normals(group): - # For normal distributions, when averaging them: - # The mean is the weighted average of means - # The variance of the sum is the sum of variances - # The variance of the average is the sum of variances divided by n^2 - n = len(group) - mean = group['auc_delong'].mean() - # Since we're taking the average, divide combined variance by n^2 - var = group['auc_variance_delong'].sum() / (n**2) - return pd.Series({ - 'auc': mean, - 'auc_variance_delong': var, - 'n_samples': n, - 'n_diseased': group['n_diseased'].sum(), - 'n_healthy': group['n_healthy'].sum(), - }) - - -auc_df_all_brackets = pd.DataFrame(all_aucs) -auc_df = auc_df_all_brackets.groupby(['token']).apply(aggregate_normals, include_groups=False).reset_index() -auc_df = auc_df.merge(delphi_labels[['name', 'index']], left_on='token', right_on='index', how='inner') -auc_df - - - - -plt.figure(figsize=(7, 5)) - -# Create the bar chart -bars = plt.bar(range(len(auc_df)), auc_df['auc'], color='skyblue') - -plt.errorbar( - range(len(auc_df)), - auc_df['auc'], - yerr=1.96 * np.sqrt(auc_df['auc_variance_delong']), - fmt='none', - color='black', - capsize=0, - linewidth=1.0, -) - -# Add labels and title -plt.xlabel('Disease') -plt.ylabel('AUC (sex & age stratified)') -plt.title('AUC for selected diseases with 95% confidence intervals') -plt.xticks(range(len(auc_df)), auc_df['name'], rotation=45, ha='right') -plt.gca().set_axisbelow(True) -plt.grid(axis='x', visible=False) -plt.axhline(0.5, color='k', linestyle='--', linewidth=0.75) -plt.ylim(0, 1.05) -plt.grid(axis='y', linestyle='--', alpha=0.7) - -plt.tight_layout() - - -# ## AUC for the entire disease set -# -# Auc can be evaluated for all diseases, but it would take about 30 minutes to run for 100k trajectories with a gpu; much londer with a cpu. -# -# Therefore, we well use precomputed results here. - - - -# df_auc_unpooled_merged, df_auc_merged = evaluate_auc_pipeline(model, -# d100k, -# output_path=None, -# delphi_labels=delphi_labels[13:].index.values, -# diseases_of_interest=diseases_of_interest, -# filter_min_total=100, # remove rare diseases -# device=device, -# ) - - - - -df_auc_all_diseases = pd.read_csv('supplementary/delphi_auc.csv') -df_auc_all_diseases['mean_auc'] = df_auc_all_diseases[['AUC Female, (no gap)', 'AUC Male, (no gap)']].mean(axis=1) - - - - -plt.figure(figsize=(7, 4)) -plt.scatter(df_auc_all_diseases['N tokens, training'], df_auc_all_diseases['mean_auc'], - c=df_auc_all_diseases['Colour'], s=24, edgecolor='white', linewidth=0.65) -plt.axhline(0.5, color='k', linestyle='--', linewidth=0.75) -plt.title('AUC vs number of tokens in training set') -plt.xscale('log') -plt.ylim(0, 1.05) -plt.xlabel('Number of tokens in training set') -plt.ylabel('AUC') -plt.show() - - - - -import matplotlib.pyplot as plt -import numpy as np -import matplotlib.patches as mpatches - -filtered_df = df_auc_all_diseases.dropna(subset=['mean_auc']) - -chapters = filtered_df['ICD-10 Chapter (short)'].unique() -chapter_data = {} - -for chapter in chapters: - if chapter not in ['Technical', 'Sex', 'Smoking, Alcohol and BMI']: # Skip non-disease chapters - chapter_data[chapter] = filtered_df[filtered_df['ICD-10 Chapter (short)'] == chapter]['mean_auc'].values - -fig, ax = plt.subplots(figsize=(8, 5)) - -chapter_colors = {} -for chapter in chapter_data.keys(): - chapter_rows = filtered_df[filtered_df['ICD-10 Chapter (short)'] == chapter] - chapter_colors[chapter] = chapter_rows['Colour'].iloc[0] - -positions = np.arange(1, len(chapter_data) + 1) -boxplots = [] - -for i, (chapter, values) in enumerate(chapter_data.items()): - bp = ax.boxplot(values, positions=[positions[i]], patch_artist=True, - widths=0.6, whis=[2.5, 97.5], showfliers=True, - boxprops={'linewidth': 1.25, 'facecolor': chapter_colors[chapter], - 'edgecolor': chapter_colors[chapter]}, - medianprops={'color': 'black', 'linewidth': 1.5}, - whiskerprops={'color': 'gray', 'linewidth': 1}, - capprops={'color': 'gray', 'linewidth': 1}, - flierprops={'marker': 'x', 'markerfacecolor': 'none', - 'markeredgecolor': 'black', 'markersize': 3, 'alpha': 0.3}) - - boxplots.append(bp) - -ax.set_xticks(positions) -ax.set_xticklabels([chapter for chapter in chapter_data.keys()], rotation=45, ha='right') - -ax.set_ylim(0, 1.025) -ax.axhline(0.5, color='black', linestyle='--', linewidth=0.75) - -ax.yaxis.grid(True, linestyle='--', alpha=0.7) -ax.set_axisbelow(True) - -ax.set_ylabel('AUC') -ax.set_xlabel('ICD-10 chapter') -ax.set_title('AUC, grouped by ICD-10 chapter', y=1.05) - -plt.tight_layout() -plt.grid(axis='x', visible=False) -plt.show() - - - - -import matplotlib.pyplot as plt -import numpy as np -import matplotlib.patches as mpatches - -# Filter out rows with NaN values in mean_auc -filtered_df = df_auc_all_diseases.dropna(subset=['mean_auc']) - -# Create separate data for males and females -male_data = filtered_df[filtered_df['AUC Male, (no gap)'].notna()]['AUC Male, (no gap)'].values -female_data = filtered_df[filtered_df['AUC Female, (no gap)'].notna()]['AUC Female, (no gap)'].values - -# Set up the figure -fig, ax = plt.subplots(figsize=(1.75, 4)) - -# Define colors for male and female -male_color = normal_male -female_color = normal_female - -# Create boxplots -positions = [1, 2] -boxplots = [] - -# Create boxplots for both sexes using a loop -sex_data = [female_data, male_data] -sex_colors = [female_color, male_color] -sex_labels = ['Female', 'Male'] - -for i in range(2): - bp = ax.boxplot(sex_data[i], positions=[positions[i]], patch_artist=True, - widths=0.6, whis=[2.5, 97.5], showfliers=True, - boxprops={'linewidth': 1.25, 'facecolor': sex_colors[i], - 'edgecolor': sex_colors[i]}, - medianprops={'color': 'black', 'linewidth': 1.5}, - whiskerprops={'color': 'gray', 'linewidth': 1}, - capprops={'color': 'gray', 'linewidth': 1}, - flierprops={'marker': 'x', 'markerfacecolor': 'none', - 'markeredgecolor': 'black', 'markersize': 3, 'alpha': 0.3}) - boxplots.append(bp) - -# Set x-axis labels -ax.set_xticks(positions) -ax.set_xticklabels(sex_labels) - -# Set y-axis limits and add reference line -ax.set_ylim(0, 1.025) -ax.axhline(0.5, color='black', linestyle='--', linewidth=0.75) - -# Add grid for y-axis only -ax.yaxis.grid(True, linestyle='--', alpha=0.7) -ax.set_axisbelow(True) - -# Add labels and title -ax.set_ylabel('AUC') -ax.set_title('AUC, grouped by sex', y=1.05) - -# Adjust layout -plt.tight_layout() -plt.grid(axis='x', visible=False) - -plt.show() - - -# ## Interpretability - -# ### Attention maps -# -# Being a transformer model, Delphi uses attention to aggregate information from the input tokens. Here, we plot the attention matrices for all heads and layers for a single trajectory. -# -# Note how different heads and layers attend to different parts of the input trajectory. -# -# Attention maps can be used for interpretability, however for a more robust interpretation, we suggest using SHAP values (`shap_analysis.ipynb`). - - - -d = get_batch([0], val, val_p2i, select='left', block_size=model.config.block_size, device=device) - -risk = np.exp(model(d[0],d[1])[0].cpu().detach().numpy().squeeze()) -att = model(d[0],d[1])[2].cpu().detach().numpy().squeeze() -fig, ax = plt.subplots(*att.shape[:2], figsize=(12,12), sharex=True, sharey=True) -for i in range(att.shape[0]): - for j in range(att.shape[1]): - ax[i,j].imshow(att[i,j], vmax=0.35) - if i==0: - ax[i,j].set_title(f"Head {j}") - if j==0: - ax[i,j].set_ylabel(f"Layer {i}") -plt.tight_layout() - - - - -d = get_batch(range(dataset_subset_size), val, val_p2i, - select='left', block_size=48, - device=device, padding='random') -w = np.where(torch.isin(d[2].cpu(), torch.tensor(diseases_of_interest)).sum(axis=1)) -w = (w[0][:3],) -att = model(*list(map(lambda x: x[w[0],:], d)))[2].cpu().detach().numpy().squeeze() -att.shape - - -# We can also plot average attention across all heads and layers to see which tokens are attended to most "on average". -# -# Generally, tokens tend to lose most of their importance pretty quickly. High attention for the most recent token in the trajectory is likely due to this tokens being used by the model to estimate the current age of the patient, which is a very important predictor for the overall disease risk. - - - -import textwrap - -d = [d_.cpu() for d_ in d] - -for i in range(len(w[0])): - print(i) - j = (d[0][w[0][i]]==0).sum() - plt.figure(figsize=(3 * (d[3][w[0][i],-1]-d[1][w[0][i],0])/365.25/70,9 * (48-j)/48)) - x = torch.concatenate([d[1][w[0][i]], d[3][w[0][i],[-1]]])/365.25 - plt.pcolormesh( x[j:],np.arange(j,49,1), att[0,i,:,j:,j:].max((0)).T, cmap='Blues') - _ = plt.yticks(np.arange(j,48)+.5, [f"{textwrap.shorten(delphi_labels.loc[i,'name'],50)}" if i > 1 else "" for i,t in zip(d[0][w[0][i],j:].detach().numpy().squeeze(),d[1][w[0][i],j:].detach().numpy().squeeze()/365.25)]) - plt.gca().invert_yaxis() - plt.xlabel('Age') - plt.show() - - -# ## Embeddings -# -# Lastly, it's interesting to look into the learned latent space of the model. -# -# Here, we plot the UMAP of the learned disease embeddings. -# -# We see that diseases cluster by their ICD-10 chapter - which is interesting, because the model had no knowledge about the ICD-10 hierarchy during training; all diseases were treated equally. - - - -import umap -import matplotlib as mpl - - - - -wte = model.transformer.wte.weight.cpu().detach().numpy() -seed = 1413 -t = umap.UMAP(random_state=seed, n_neighbors=30, min_dist=0.05, metric='cosine').fit(wte) - -u0 = t.transform(model.transformer.wte.weight.cpu().detach().numpy()) -u = u0 - np.median(u0, axis=0) -u = - u - - - - -def remove_ticks(ax): - ax.set_xticklabels([]) - ax.set_yticklabels([]) - - for tick in ax.xaxis.get_major_ticks(): - tick.tick1line.set_visible(False) - tick.tick2line.set_visible(False) - - for tick in ax.yaxis.get_major_ticks(): - tick.tick1line.set_visible(False) - tick.tick2line.set_visible(False) - - - - -labels_all = pd.read_csv('delphi_labels_chapters_colours_icd.csv') -labels_all['UMAP1'] = u[:,0] -labels_all['UMAP2'] = u[:,1] -labels_all = labels_all[labels_all['count'] > 20].reset_index(drop=True).reset_index() -labels_non_technical = labels_all[~labels_all['ICD-10 Chapter'].isin(['Technical', 'Sex', 'Smoking, Alcohol and BMI'])] -labels_non_technical = labels_non_technical[(labels_non_technical['UMAP1'].abs() < 5) & (labels_non_technical['UMAP2'].abs() < 5)] -short_names = labels_all['ICD-10 Chapter (short)'].unique() -short_names_present = [i for i in short_names if i in labels_non_technical['ICD-10 Chapter (short)'].unique()] -color_mapping_short = {k: v for k, v in labels_all[['ICD-10 Chapter (short)', 'color']].values} - - - - -import seaborn as sns - -fig, ax = plt.subplots(figsize=(8, 8)) - -sns.scatterplot(x='UMAP1', y='UMAP2', data=labels_non_technical, hue='ICD-10 Chapter (short)', - palette=color_mapping_short, - hue_order=short_names_present, size='count', sizes=(20, 200), - alpha=0.9, ax=ax, linewidth=0.15) - -ax.legend_.set_bbox_to_anchor((1.1, 0.85)) -ax.grid(None) -remove_ticks(ax) -ax.set_aspect('equal') -plt.title('UMAP of learned disease embeddings'); - - -# ## The End! -# -# If you want to learn more about Delphi, check out the `shap_analysis.ipynb` notebook next, where we use SHAP values to interpret the model's predictions. diff --git a/shap_analysis.py b/shap_analysis.py deleted file mode 100644 index f4616d1..0000000 --- a/shap_analysis.py +++ /dev/null @@ -1,601 +0,0 @@ -#!/usr/bin/env python -# coding: utf-8 - -# # Looking inside Delphi with SHAP values -# -# Welcome to the Delphi SHAP notebook! -# -# Delphi is a generative autoregressive model that not only predicts the future disease rates, but also sample entire disease trajectories one step at a time. -# -# -# Let's start by looking at what SHAP values mean: -# -# ## SHAP Values and Delphi -# -# SHAP (SHapley Additive exPlanations) values help us understand how a machine learning model makes its predictions by showing the contribution of each input feature. -# -# -# ### Example: Patient Trajectory -# -# Consider a simplified patient trajectory: -# `Male, Migraine, Common cold, Brain cancer` -# -# For this trajectory, Delphi would predict a very high mortality risk (aka high rate for the Death token being next). Say, 95% chance of death within a year. Why? Technically, we don't know, since neural networks are black boxes. -# -# Let's try masking several tokens and predicting the next token again. -# -# `Male, [Masked: Migraine], Common cold, Brain cancer` -> 95% chance of death within a year, no change -# -# `Male, Migraine, Common cold, [Masked: Brain cancer]` -> 5% chance of death within a year, risk drops significantly -# -# Without speaking about causality, we can assume that there is *some* connection between brain cancer and death risk. SHAP framework allows using such masking to systematically assess the contribution of each token to the prediction. We can perform this analysis for all trajectories in the dataset and evaluate how, on average, a given disease influences the risk of any other disease. -# - - - -import os -import pickle -import torch -from model import DelphiConfig, Delphi -from tqdm import tqdm -import pandas as pd -import numpy as np -import textwrap -import warnings - -import matplotlib.pyplot as plt -get_ipython().run_line_magic('config', "InlineBackend.figure_format='retina'") - -plt.rcParams['figure.facecolor'] = 'white' -plt.rcParams.update({'axes.grid': True, - 'grid.linestyle': ':', - 'axes.spines.bottom': False, - 'axes.spines.left': False, - 'axes.spines.right': False, - 'axes.spines.top': False}) -plt.rcParams['figure.dpi']= 72 -plt.rcParams['pdf.fonttype'] = 42 - -delphi_labels = pd.read_csv('delphi_labels_chapters_colours_icd.csv') - - -# ## Load model - - - -out_dir = 'Delphi-2M' -device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc. -dtype ='float32' #'bfloat16' # 'float32' or 'bfloat16' or 'float16' -seed = 1337 - -torch.manual_seed(seed) -torch.cuda.manual_seed(seed) -torch.backends.cuda.matmul.allow_tf32 = True -torch.backends.cudnn.allow_tf32 = True - -device_type = 'cuda' if 'cuda' in device else 'cpu' -dtype = {'float32': torch.float32, 'float64': torch.float64, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype] - - - - -ckpt_path = os.path.join(out_dir, 'ckpt.pt') -checkpoint = torch.load(ckpt_path, map_location=device) -conf = DelphiConfig(**checkpoint['model_args']) -model = Delphi(conf) -state_dict = checkpoint['model'] -model.load_state_dict(state_dict) - -model.eval() -model = model.to(device) - - -# ## Load data - - - -from utils import get_batch, get_p2i - -train = np.fromfile('data/ukb_simulated_data/train.bin', dtype=np.uint32).reshape(-1,3) -val = np.fromfile('data/ukb_simulated_data/val.bin', dtype=np.uint32).reshape(-1,3) - -train_p2i = get_p2i(train) -val_p2i = get_p2i(val) - - - - -# define a random example health trajectory - -person = [('Male', 0), - ('B01 Varicella [chickenpox]',2), - ('L20 Atopic dermatitis',3), - ('Healthy', 5), - ('Healthy', 10), - ('Healthy', 15), - ('Healthy', 20), - ('G43 Migraine', 20), - ('E73 Lactose intolerance', 21), - ('B27 Infectious mononucleosis', 22), - ('Healthy', 25), - ('J11 Influenza, virus not identified', 28), - ('Healthy', 30), - ('Healthy', 35), - ('C25 Malignant neoplasm of pancreas', 38), - ('Healthy', 40), - ('Smoking low', 41), - ('BMI mid', 41), - ('Alcohol high', 41), - ('Healthy', 42), -] -person = [(a, b * 365.25) for a,b in person] - - -# ### Individual SHAP values - - - -# define helper functions - -id_to_token = delphi_labels['name'].to_dict() -token_to_id = {v:k for k, v in id_to_token.items()} - -def tokens_to_ids(tokens): - return [token_to_id[t] for t in tokens] - -def ids_to_tokens(ids): - return [id_to_token[int(id_)] for id_ in ids] - -def split_person(p): - tokens = [i[0] for i in p] - ages = [i[1] for i in p] - return tokens, ages - -def get_person(idx): - x, y, _, time = get_batch([idx], val, val_p2i, - select='left', block_size=48, - device=device, padding='random') - x, y = x[y > -1], y[y > -1] - person = [] - for token_id, date in zip(x, y): - person.append((id_to_token[token_id.item()], date.item())) - return person, y, time[0][-1] - - - - -from utils import shap_custom_tokenizer, shap_model_creator -import shap - -# person_to_process = get_person(137)[0] -person_to_process = person -diseases_of_interest = [1269, 46, 95, 1168, 374, 173, 214, 305, 505, 584] - -person_tokens, person_ages = split_person(person_to_process) -person_tokens_ids = tokens_to_ids(person_tokens) - -masker = shap.maskers.Text(shap_custom_tokenizer, output_type='str', mask_token='10000', collapse_mask_token=False) -model_shap = shap_model_creator(model, diseases_of_interest, person_tokens_ids, person_ages, device) -explainer = shap.Explainer(model_shap, masker, output_names=delphi_labels['name'].values[diseases_of_interest]) - -shap_values = explainer([' '.join(list(map(lambda x: str(token_to_id[x]), person_tokens)))]) -shap_values.data = np.array([list(map(lambda x: f"{delphi_labels['name'].values[token_to_id[x[0]]]}({x[1]/365:.1f} years) ", person_to_process))]) - - - - -out = shap.plots.text(shap_values, display=True) # sometimes this interactive plot can't be rendered well (eg in VS Code, feel free to skip it) - - - - -# SHAP values can be interpreted as how much each input token changes predicted logit corresponding to a particular disease. -# As Delphi logits are log-disease rates, we can convent SHAP values to the disease-specific fold risk changes. - -# Shown below is a waterfall plot, showing SHAP values for the most "influential" diseases within -# a single trajectory. - -from plotting import waterfall - -with plt.style.context('default'): - plt.rcParams['pdf.fonttype'] = 42 - plt.rcParams['figure.dpi'] = 150 - plt.rcParams['font.size'] = 4 - waterfall(shap_values[0, ..., 0], max_display=7, show=False, ages=person_ages) - plt.gca().set_title('Impact of diseases on mortality', fontweight=1, size=18) - plt.show() - - -# ## Pre-computed many cases -# -# The small synthetic dataset is not enough to properly run following part; if you have access to the full dataset, run `shap-agg-eval.py` to evaluate SHAP values for the entire dataset. - - - -import pickle - -with open('shap_agg.pickle', 'rb') as f: - shap_pkl = pickle.load(f) - -all_tokens = shap_pkl['tokens'] -all_values = shap_pkl['values'] - - - - -import pandas as pd -import seaborn as sns -import numpy as np -import matplotlib.pyplot as plt - -df_shap = pd.DataFrame(all_values) -df_shap['token'] = all_tokens.astype('int') - - - - -token_count_dict = df_shap['token'].value_counts().sort_index().to_dict() - -N_min = 5 # we will only consider diseases that have at least 5 calculated SHAP values, otherwise they are too noisy - -columns_more_N = [c for c in df_shap.columns if c == 1269 or (c in token_count_dict and token_count_dict[c] >= N_min)] -df_shap_agg = df_shap[df_shap['token'].apply(lambda x: token_count_dict[x] > N_min)].groupby('token').mean() - - -# Since we now have calculated SHAP values for the entire dataset, we can use them to analyse "connections" between diseases. -# -# For every "predicted"-"predictor" pair, we average all SHAP values for the given pair. -# -# We can further analyse them in sevaral directions: -# - For a given disease in the past medical history, which disease rates are most increased by it? -# - For a given potential future disease, having which disease in the past medical history would most increase its rate? -# -# Let's see which diseases increase disease risk the most and also which diseases are most influenced by being a heavy smoker. - - - -import matplotlib.pyplot as plt -import numpy as np - -def plot_shap_distribution(df_melted, y_axis_labels, group_by_col_name, title, x_lim_tuple, highlight_last_dot=True): - """ - Generates a plot showing median and quartile ranges of SHAP values for different tokens. - """ - plt.figure(figsize=(3, 6), facecolor='w') - ax = plt.gca() - - for i, label_for_y_tick in enumerate(y_axis_labels): - data_for_label = df_melted[df_melted[group_by_col_name] == label_for_y_tick]['value'] - - if not data_for_label.empty: - median = np.median(data_for_label) - quartiles = np.percentile(data_for_label, [25, 75]) - - dot_color = 'red' if highlight_last_dot and i == len(y_axis_labels) - 1 else 'black' - ax.plot(median, i, 'o', color=dot_color, zorder=3) - ax.hlines(i, quartiles[0], quartiles[1], color='gray', linestyles='solid', linewidth=1) - - plt.title(title) - plt.yticks(range(len(y_axis_labels)), y_axis_labels) - plt.xticks(rotation=25, ha='right') - plt.xscale('log') - plt.xlim(*x_lim_tuple) - plt.xlabel('Risk increase, folds', size=11, labelpad=10) - plt.show() - -target_token = 1269 -n_first = 20 -plot = True - -selected_context_tokens1 = df_shap_agg[target_token].nlargest(n_first).index[::-1] - -df_plot_source = df_shap[df_shap['token'].isin(selected_context_tokens1)] -df_plot_melted = df_plot_source[[target_token, 'token']].reset_index(drop=True).melt(id_vars=['token'], value_vars=[target_token]) -df_plot_melted['context_token_label'] = df_plot_melted['token'].map(id_to_token) -df_plot_melted['value'] = np.exp(df_plot_melted['value']) - -y_axis_labels1 = [id_to_token[token] for token in selected_context_tokens1] -title1 = 'Mortality factors' -xlim1 = (1, 1000) - -plot_shap_distribution( - df_melted=df_plot_melted, - y_axis_labels=y_axis_labels1, - group_by_col_name='context_token_label', - title=title1, - x_lim_tuple=xlim1 -) - - - - -target_token = 9 -n_first = 20 - -shap_values_for_context = df_shap_agg.loc[target_token] -selected_feature_tokens2 = shap_values_for_context.sort_values(ascending=False).index[:n_first][::-1] - -df_plot_source = df_shap[df_shap['token'] == target_token] -df_plot_melted = df_plot_source[[*selected_feature_tokens2, 'token']].reset_index(drop=True).melt( - id_vars=['token'], - value_vars=selected_feature_tokens2, - var_name='feature_token_id', - value_name='raw_shap_value' -) -df_plot_melted['feature_label'] = df_plot_melted['feature_token_id'].map(id_to_token) -df_plot_melted['value'] = np.exp(df_plot_melted['raw_shap_value']) - -y_axis_labels = [id_to_token[token] for token in selected_feature_tokens2] -title = 'Consequences of\nsmoking heavily' -xlim = (1, 11) - -plot_shap_distribution( - df_melted=df_plot_melted, - y_axis_labels=y_axis_labels, - group_by_col_name='feature_label', - title=title, - x_lim_tuple=xlim -) - - -# ### Time-resolved SHAP analysis -# -# Before, we aggregated the calculated SHAP values in a fairly simple way: just averaged them within all "predictor-predicted" pairs. This is an oversimplification, since the context in which these two diseases occur is also important. -# -# For instance, the amount of time passed since the "predictor" disease occured is important, since some acute conditions may have vastly different effects compared to their chronic forms. -# -# Now, we will aggregate SHAP values within the pairs, additionally separating them by the time between the "predictor" and "predicted" diseases. - - - -d = get_batch(range(len(np.unique(shap_pkl['people']))), val, val_p2i, - select='left', block_size=48, - device='cpu', padding='regular') - - - - -has_gender = torch.isin(d[0], torch.tensor([2, 3])).any(dim=1).numpy() -is_male = torch.isin(d[0], torch.tensor([3])).any(dim=1).numpy() -is_female = torch.isin(d[0], torch.tensor([2])).any(dim=1).numpy() - - - - -def get_person(idx): - x, y, _, time = get_batch([idx], val, val_p2i, - select='left', block_size=64, - device=device, padding='random', - cut_batch=True) - - x, y = x[y > -1], y[y > -1] - person = [] - for token_id, date in zip(x, y): - person.append((id_to_token[token_id.item()], date.item())) - return person, y, time[0][-1] - - - - -# the shap result pickle does not contain time, so we need to add it - -persons_lengths = [] -ages = [] -reg_times = [] - -for p in tqdm(np.unique(shap_pkl['people'])): - pers = get_person(p) - - reg_time_idx = np.where(np.isin(tokens_to_ids(np.array(pers[0])[:, 0]), np.arange(4, 13)))[0] - if len(reg_time_idx) > 0: - reg_time = pers[0][reg_time_idx[0]][1] - else: - reg_time = -1 - - reg_times += [reg_time] * len(pers[0]) - persons_lengths += [p] * len(pers[0]) - ages += [pers[-1].item()] * len(pers[0]) - -assert len(ages) == len(df_shap) - - - - -all_tokens = shap_pkl['tokens'] -all_values = shap_pkl['values'] -all_times = shap_pkl['times'] - -df_shap = pd.DataFrame(all_values) -df_shap['token'] = all_tokens -df_shap['time'] = all_times -df_shap['person'] = shap_pkl['people'] -df_shap['age'] = np.array(ages) / 365.25 -df_shap['reg_time_years'] = np.array(reg_times) / 365.25 - -df_shap['Time, years'] = df_shap['time'] / 365.25 -df_shap['age_at_token'] = df_shap['age'] - df_shap['time'] / 365.25 - -df_shap = df_shap[df_shap['reg_time_years'] > 0] - -token_count_dict = df_shap['token'].value_counts().sort_index().to_dict() - - - - -import numpy as np - -def bins_avg(x, y, grid_size=3): - '''Filter out regions wiht few data points''' - x, y = np.array(x), np.array(y) - - bin_edges = np.arange(np.min(x), np.max(x), grid_size) - - bin_indices = np.digitize(x, bin_edges) - bin_avgs = np.array([y[bin_indices == i].mean() for i in range(1, len(bin_edges)+1)]) - - return bin_edges, bin_avgs - - - - -tokens_of_interest = [46, 95, 1168, 1188, 173, 214, 305, 505, 584] -n_groups = len(tokens_of_interest) // 5 + 1 - -palette_faint = [sns.color_palette("Paired")[0], sns.color_palette("Paired")[2], sns.color_palette("Paired")[4]] -palette_bright = [sns.color_palette("Paired")[1], sns.color_palette("Paired")[3], sns.color_palette("Paired")[5]] - -for num_g, token_group in enumerate(np.array_split(tokens_of_interest, n_groups)): - - fig, axs = plt.subplots(1, 5, figsize=(12, 2), sharey=True) - - for num, (ax, token_id) in enumerate(zip(axs.flatten(), token_group)): - df_trait = df_shap[df_shap['token'] == token_id].copy() - df_trait[1269] = np.exp(df_trait[1269].values) - df_trait['Time, years'] = df_trait['time'] / 365.25 - df_trait = df_trait.head(2000) - if len(df_trait) < 2: - continue - - sns.scatterplot(data=df_trait, x='Time, years', y=1269, ax=ax, color=palette_faint[0], alpha=0.7, rasterized=True) - - x, y = df_trait['Time, years'], df_trait[1269] - n = 3 - - with warnings.catch_warnings(): - warnings.simplefilter('ignore') - x, y = bins_avg(x, y, grid_size=n) - - ax.plot(x, y, color=palette_bright[0], linewidth=1.5) - - ax.set_ylim(0.5, 500) - ax.set_xlim(0.1, 10) - ax.set_ylabel('Impact on mortality'); - ax.set_title(textwrap.fill(id_to_token[token_id], width=15) + f' {token_id}', size=9) - # ax.set_xscale('log') - - ax.set_yscale('log') - plt.show() - - -# As shown is the graph above, some diseases (pancteatic cancer, miocardial infarction, septiceamia) have a much higher impact on mortality if they occur in the recent past, while others (diabetis, depression) don't have a clear time-dependence. - -# ## Interaction heatmap -# -# To analyse the interactions between diseases more systematically, we can plot a heatmap of the SHAP values for all "predictor-predicted" pairs, sorted by ICD-10 chapter. -# -# Let's plot two separate heatmaps, one for the cases where the "predictor" disease occured in the past 5 years (with the "predicted disease being the reference) and one for the cases where it occured more than 10 years ago. - - - -N_min = 5 - -token_count_dict_below_5y = df_shap[df_shap['Time, years'] < 5]['token'].value_counts().sort_index().to_dict() -token_count_dict_over_10y = df_shap[df_shap['Time, years'] > 10]['token'].value_counts().sort_index().to_dict() - -for d in [token_count_dict_below_5y, token_count_dict_over_10y]: - for i in range(1300): - if i not in d: - d[i] = 0 - -columns_more_N = [c for c in df_shap.columns if c == 1269 or(c in token_count_dict_below_5y and token_count_dict_below_5y[c] >= N_min and - c in token_count_dict_over_10y and token_count_dict_over_10y[c] >= N_min)] -df_shap_agg_below_5y = df_shap[df_shap['token'].apply(lambda x: x in columns_more_N) & (df_shap['Time, years'] < 5)].groupby('token').mean()[columns_more_N] -df_shap_agg_over_10y = df_shap[df_shap['token'].apply(lambda x: x in columns_more_N) & (df_shap['Time, years'] > 10)].groupby('token').mean()[columns_more_N] - - - - -from matplotlib.colors import LogNorm - -to_exclude_predicted = ['Technical', 'Smoking, Alcohol and BMI', 'Sex', 'XVI. Perinatal Conditions'] -to_exclude_predictor = ['Technical', 'Smoking, Alcohol and BMI', 'Sex', 'XVI. Perinatal Conditions', 'Death'] - -chapter_order = ['I. Infectious Diseases', 'II. Neoplasms', - 'III. Blood & Immune Disorders', 'IV. Metabolic Diseases', - 'V. Mental Disorders', 'VI. Nervous System Diseases', - 'VII. Eye Diseases', 'VIII. Ear Diseases', 'IX. Circulatory Diseases', 'X. Respiratory Diseases', - 'XI. Digestive Diseases', 'XII. Skin Diseases', - 'XIII. Musculoskeletal Diseases', 'XIV. Genitourinary Diseases', - 'XV. Pregnancy & Childbirth', 'XVI. Perinatal Conditions', - 'XVII. Congenital Abnormalities', 'Death'] - -def get_tick_coords(arr): - return np.where(arr[1:] != arr[:-1])[0] - -def plot_full_shap_heatmap(cur_df, title): - new_death_rows = 10 - - for c in range(1269+1, 1269+new_death_rows+1): - cur_df[c] = cur_df[1269] - - delphi_labels = pd.read_csv("delphi_labels_chapters_colours_icd.csv", index_col=0) - death_df = delphi_labels[delphi_labels['ICD-10 Chapter (short)']=="Death"].sample(new_death_rows, replace=True) - death_df['name'] = death_df['name'].apply(lambda x: x + str(np.random.randint(0, 100000))) - death_df.index = pd.Index(range(1269+1, 1269+new_death_rows+1)) - delphi_labels = pd.concat([delphi_labels, death_df]) - - to_exclude_predicted_idx = delphi_labels[~delphi_labels['ICD-10 Chapter (short)'].isin(to_exclude_predicted)].index - to_exclude_predictor_idx = delphi_labels[~delphi_labels['ICD-10 Chapter (short)'].isin(to_exclude_predictor)].index - - to_exclude_predicted_idx = to_exclude_predicted_idx[to_exclude_predicted_idx.isin(cur_df.columns)] - to_exclude_predicted_idx = sorted(to_exclude_predicted_idx, key=lambda x: (chapter_order.index(delphi_labels.loc[x, 'ICD-10 Chapter (short)']), x)) - to_exclude_predicted_idx = pd.Index(to_exclude_predicted_idx) - - to_exclude_predictor_idx = to_exclude_predictor_idx[to_exclude_predictor_idx.isin(cur_df.index.values)] - to_exclude_predictor_idx = sorted(to_exclude_predictor_idx, key=lambda x: (chapter_order.index(delphi_labels.loc[x, 'ICD-10 Chapter (short)']), x)) - to_exclude_predictor_idx = pd.Index(to_exclude_predictor_idx) - - cur_df = cur_df.loc[to_exclude_predictor_idx, to_exclude_predicted_idx] - - row_colors = delphi_labels.iloc[cur_df.index.values]['color'].to_numpy() - col_colors = delphi_labels.iloc[cur_df.columns]['color'].to_numpy() - - y_tick_coords = get_tick_coords(delphi_labels.iloc[cur_df.index.values]['color'].to_numpy()) - x_tick_coords = get_tick_coords(delphi_labels.iloc[cur_df.columns]['color'].to_numpy()) - - g = sns.clustermap(np.exp(cur_df.values), row_cluster=False, col_cluster=False, - row_colors=row_colors, col_colors=col_colors, - # norm=LogNorm(vmin=5e-2, vmax=2e1), - norm=LogNorm(vmin=1e-1, vmax=1e1), - cmap='RdBu_r', - figsize=(8.5, 8.5), - rasterized=True, - ) - - g.ax_heatmap.set_xticks(x_tick_coords) - g.ax_heatmap.set_yticks(y_tick_coords) - g.ax_heatmap.tick_params(length=0, width=0.5, labelsize=8, grid_alpha=0.6, grid_linewidth=0.35, grid_color='gray') - g.ax_cbar.tick_params(length=0.5, width=0.6, labelsize=8, grid_alpha=0.45, grid_linewidth=0.45) - - for ch, color in delphi_labels[['ICD-10 Chapter (short)', 'color']].drop_duplicates('color').values: - col_loc = np.where(col_colors == color)[0].mean() if (col_colors == color).any() else np.nan - g.ax_heatmap.text(col_loc - 10, -60, ch, va='bottom', rotation=90, ha='left', size=8) - - row_loc = np.where(row_colors == color)[0].mean() if (row_colors == color).any() else np.nan - g.ax_heatmap.text(-70, row_loc, ch, va='center', ha='right', size=9) - - from matplotlib.patches import Patch - - # Create legend for chapter colors - chapter_color_map = delphi_labels[['ICD-10 Chapter (short)', 'color']].drop_duplicates('color') - chapter_color_map = chapter_color_map[~chapter_color_map['ICD-10 Chapter (short)'].isin(to_exclude_predicted)] - handles = [Patch(facecolor=color) for color in chapter_color_map['color']] - plt.legend(handles, chapter_color_map['ICD-10 Chapter (short)'], title='ICD-10 Chapters', - bbox_transform=plt.gcf().transFigure, loc='center left', bbox_to_anchor=(.96, 0.5)) - - plt.suptitle(title, y=1.1, size=10, x=0.5) - plt.show() - - - - -plot_full_shap_heatmap(df_shap_agg_below_5y, 'Influence of tokens from below 5 years,\nrisk increase, folds') - - - - -plot_full_shap_heatmap(df_shap_agg_over_10y, 'Influence of tokens from above 10 years,\nrisk increase, folds') - - -# -# The second ("above 10 years") heatmap is also more pale, meaning that most of the disease-disease interactions get weaker over time. diff --git a/train_iter_multigpu.py b/train_iter_multigpu.py deleted file mode 100644 index 8c9119b..0000000 --- a/train_iter_multigpu.py +++ /dev/null @@ -1,247 +0,0 @@ -import torch -import torch.nn as nn -from torch.optim import AdamW -from torch.utils.data import DataLoader -import numpy as np -import math -import tqdm -import matplotlib.pyplot as plt -import json -import itertools -import os -import torch.distributed as dist -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data.distributed import DistributedSampler - -from models import TimeAwareGPT2, CombinedLoss -from utils import PatientEventDataset - -# --- Configuration --- -class TrainConfig: - # Data parameters - train_data_path = 'ukb_real_train.bin' - val_data_path = 'ukb_real_val.bin' - block_length = 48 # Sequence length - - # Model parameters - n_embd = 120 - n_layer = 12 - n_head = 12 - pdrop = 0.1 - token_pdrop = 0.1 - - # Training parameters - max_iter = 200000 - batch_size = 128 # Per GPU - lr_initial = 6e-4 - lr_final = 6e-5 - weight_decay = 2e-1 - warmup_iter = 1000 - - # Loss parameters - # 0 = padding, 1 = "no event" - ignored_token_ids = [0, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] # Example ignored token IDs - - # System parameters - device = 'cuda' - -# --- DDP Setup --- -def setup_ddp(): - """Initializes the distributed data parallel environment.""" - dist.init_process_group(backend='nccl') - rank = dist.get_rank() - local_rank = int(os.environ['LOCAL_RANK']) - torch.cuda.set_device(local_rank) - return rank, local_rank - -def cleanup_ddp(): - """Cleans up the distributed data parallel environment.""" - dist.destroy_process_group() - -# --- Main Training Script --- -def main(): - rank, local_rank = setup_ddp() - is_main_process = (rank == 0) - - config = TrainConfig() - config.device = f'cuda:{local_rank}' - - if is_main_process: - model_filename = f"best_model_n_embd_{config.n_embd}_n_layer_{config.n_layer}_n_head_{config.n_head}_iter_multigpu.pt" - # --- 0. Save Configuration --- - config_filename = f"config_n_embd_{config.n_embd}_n_layer_{config.n_layer}_n_head_{config.n_head}_iter_multigpu.json" - config_dict = {k: v for k, v in vars(config).items() if not k.startswith('__')} - with open(config_filename, 'w') as f: - json.dump(config_dict, f, indent=4) - print(f"Configuration saved to {config_filename}") - - # --- 1. Data Loading --- - if is_main_process: - print(f"Loading data from {config.train_data_path} and {config.val_data_path}...") - train_data_arr = np.memmap(config.train_data_path, dtype=np.uint32, mode='r').reshape(-1, 3) - val_data_arr = np.memmap(config.val_data_path, dtype=np.uint32, mode='r').reshape(-1, 3) - - vocab_size = int(max(train_data_arr[:, 2].max(), val_data_arr[:, 2].max())) + 1 - if is_main_process: - print(f"Inferred vocabulary size: {vocab_size}") - - train_dataset = PatientEventDataset(train_data_arr, config.block_length) - val_dataset = PatientEventDataset(val_data_arr, config.block_length) - - train_sampler = DistributedSampler(train_dataset) - val_sampler = DistributedSampler(val_dataset, shuffle=False) - - train_loader = DataLoader(train_dataset, batch_size=config.batch_size, sampler=train_sampler, num_workers=4, pin_memory=True) - val_loader = DataLoader(val_dataset, batch_size=config.batch_size, sampler=val_sampler, num_workers=4, pin_memory=True) - train_iter_loader = iter(itertools.cycle(train_loader)) - - # --- 2. Model, Optimizer, and Loss Initialization --- - if is_main_process: - print(f"Initializing model on {config.device}...") - model = TimeAwareGPT2( - vocab_size=vocab_size, - n_embd=config.n_embd, - n_layer=config.n_layer, - n_head=config.n_head, - pdrop=config.pdrop, - token_pdrop=config.token_pdrop - ).to(config.device) - model = DDP(model, device_ids=[local_rank]) - - if is_main_process: - print(f"Model initialized with {model.module.get_num_params():.2f}M trainable parameters.") - - loss_fn = CombinedLoss(config.ignored_token_ids) - optimizer = AdamW(model.parameters(), lr=config.lr_initial, weight_decay=config.weight_decay) - - # --- 3. Training Loop --- - train_losses_ce, train_losses_surv, train_losses_total = [], [], [] - - if is_main_process: - print("Starting training...") - pbar = tqdm.tqdm(range(1, config.max_iter + 1), desc="Training", disable=not is_main_process) - for iter_num in pbar: - # --- Learning Rate Scheduling --- - if iter_num < config.warmup_iter: - lr = config.lr_initial - else: - progress = (iter_num - config.warmup_iter) / (config.max_iter - config.warmup_iter) - lr = config.lr_final + 0.5 * (config.lr_initial - config.lr_final) * (1 + math.cos(math.pi * progress)) - - for param_group in optimizer.param_groups: - param_group['lr'] = lr - - # --- Training Step --- - model.train() - - event_seq, time_seq = next(train_iter_loader) - event_seq, time_seq = event_seq.to(config.device), time_seq.to(config.device) - - input_events = event_seq[:, :-1] - input_times = time_seq[:, :-1] - target_events = event_seq[:, 1:] - target_wait_times = (time_seq[:, 1:] - time_seq[:, :-1]).float() - - logits = model(input_events, input_times) - loss_ce, loss_survival = loss_fn(logits, target_events, target_wait_times) - loss = loss_ce + loss_survival - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - if is_main_process: - train_losses_ce.append(loss_ce.item()) - train_losses_surv.append(loss_survival.item()) - train_losses_total.append(loss.item()) - pbar.set_postfix({'loss_ce': f'{loss_ce.item():.4f}', 'loss_surv': f'{loss_survival.item():.4f}', 'lr': f'{lr:.2e}'}) - - if is_main_process: - print("\nTraining finished.") - - # --- 4. Final Validation --- - if is_main_process: - print("Running final validation...") - model.eval() - val_loss_ce_acc, val_loss_surv_acc = 0.0, 0.0 - val_steps = 0 - - with torch.no_grad(): - pbar_val = tqdm.tqdm(val_loader, desc="Final Validation", disable=not is_main_process) - for event_seq, time_seq in pbar_val: - event_seq, time_seq = event_seq.to(config.device), time_seq.to(config.device) - - input_events = event_seq[:, :-1] - input_times = time_seq[:, :-1] - target_events = event_seq[:, 1:] - target_wait_times = (time_seq[:, 1:] - time_seq[:, :-1]).float() - - logits = model(input_events, input_times) - loss_ce, loss_survival = loss_fn(logits, target_events, target_wait_times) - - val_loss_ce_acc += loss_ce.item() - val_loss_surv_acc += loss_survival.item() - val_steps += 1 - if is_main_process: - pbar_val.set_postfix({'loss_ce': f'{loss_ce.item():.4f}', 'loss_surv': f'{loss_survival.item():.4f}'}) - - 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 - - if is_main_process: - print(f"Final Validation Summary: \n" - f" Val Loss: {total_val_loss:.4f} (CE: {avg_val_loss_ce:.4f}, Surv: {avg_val_loss_surv:.4f})") - - # --- 5. Save Model --- - print(f"Saving final model to {model_filename}") - torch.save(model.module.state_dict(), model_filename) - - # --- 6. Save and Plot Losses --- - losses_filename = f"losses_n_embd_{config.n_embd}_n_layer_{config.n_layer}_n_head_{config.n_head}_iter_multigpu.txt" - with open(losses_filename, 'w') as f: - f.write("iteration,train_loss_ce,train_loss_surv,train_loss_total\n") - for i in range(len(train_losses_total)): - f.write(f"{i+1},{train_losses_ce[i]},{train_losses_surv[i]},{train_losses_total[i]}\n") - print(f"\nLosses saved to {losses_filename}") - - # Plot and Save Loss Curves - iterations = range(1, len(train_losses_total) + 1) - - plt.figure(figsize=(18, 5)) - - # Plot CE Loss - plt.subplot(1, 3, 1) - plt.plot(iterations, train_losses_ce, label='Train CE') - plt.title('Cross-Entropy Loss') - plt.xlabel('Iterations') - plt.ylabel('Loss') - plt.legend() - plt.grid(True) - - # Plot Survival Loss - plt.subplot(1, 3, 2) - plt.plot(iterations, train_losses_surv, label='Train Survival') - plt.title('Survival Loss') - plt.xlabel('Iterations') - plt.ylabel('Loss') - plt.legend() - plt.grid(True) - - # Plot Total Loss - plt.subplot(1, 3, 3) - plt.plot(iterations, train_losses_total, label='Train Total') - plt.title('Total Loss') - plt.xlabel('Iterations') - plt.ylabel('Loss') - plt.legend() - plt.grid(True) - - plt.tight_layout() - plt.savefig('loss_curves_iter_multigpu.png') - print("\nLoss curves saved to loss_curves_iter_multigpu.png") - - cleanup_ddp() - -if __name__ == '__main__': - main() diff --git a/train_multigpu.py b/train_multigpu.py deleted file mode 100644 index 8ce817e..0000000 --- a/train_multigpu.py +++ /dev/null @@ -1,273 +0,0 @@ -import torch -import torch.nn as nn -from torch.optim import AdamW -from torch.utils.data import DataLoader -import numpy as np -import math -import tqdm -import matplotlib.pyplot as plt -import json -import os -import torch.distributed as dist -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data.distributed import DistributedSampler - -from models import TimeAwareGPT2, CombinedLoss -from utils import PatientEventDataset - -# --- Configuration --- -class TrainConfig: - # Data parameters - train_data_path = 'ukb_real_train.bin' - val_data_path = 'ukb_real_val.bin' - block_length = 48 # Sequence length - - # Model parameters - n_embd = 256 - n_layer = 16 - n_head = 16 - pdrop = 0.1 - token_pdrop = 0.1 - - # Training parameters - max_epoch = 200 - batch_size = 128 - lr_initial = 6e-4 - lr_final = 6e-5 - weight_decay = 2e-1 - warmup_epochs = 10 - early_stopping_patience = 10 - - # Loss parameters - ignored_token_ids = [0, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - - # System parameters (will be updated by DDP setup) - device = 'cuda' - -# --- DDP Setup --- -def setup_ddp(): - """Initializes the distributed data parallel environment.""" - dist.init_process_group(backend='nccl') - rank = dist.get_rank() - local_rank = int(os.environ['LOCAL_RANK']) - torch.cuda.set_device(local_rank) - return rank, local_rank - -def cleanup_ddp(): - """Cleans up the distributed data parallel environment.""" - dist.destroy_process_group() - -# --- Main Training Script --- -def main(): - rank, local_rank = setup_ddp() - is_main_process = (rank == 0) - - config = TrainConfig() - config.device = f'cuda:{local_rank}' - - if is_main_process: - model_filename = f"best_model_n_embd_{config.n_embd}_n_layer_{config.n_layer}_n_head_{config.n_head}.pt" - checkpoint_filename = f"best_model_checkpoint_n_embd_{config.n_embd}_n_layer_{config.n_layer}_n_head_{config.n_head}.pt" - - # --- 0. Save Configuration --- - config_filename = f"config_n_embd_{config.n_embd}_n_layer_{config.n_layer}_n_head_{config.n_head}.json" - config_dict = {k: v for k, v in vars(config).items() if not k.startswith('__')} - with open(config_filename, 'w') as f: - json.dump(config_dict, f, indent=4) - print(f"Configuration saved to {config_filename}") - - # --- 1. Data Loading --- - if is_main_process: - print(f"Loading data from {config.train_data_path} and {config.val_data_path}...") - train_data_arr = np.memmap(config.train_data_path, dtype=np.uint32, mode='r').reshape(-1, 3) - val_data_arr = np.memmap(config.val_data_path, dtype=np.uint32, mode='r').reshape(-1, 3) - - vocab_size = int(max(train_data_arr[:, 2].max(), val_data_arr[:, 2].max())) + 1 - if is_main_process: - print(f"Inferred vocabulary size: {vocab_size}") - - train_dataset = PatientEventDataset(train_data_arr, config.block_length) - val_dataset = PatientEventDataset(val_data_arr, config.block_length) - - train_sampler = DistributedSampler(train_dataset) - val_sampler = DistributedSampler(val_dataset, shuffle=False) - - train_loader = DataLoader(train_dataset, batch_size=config.batch_size, sampler=train_sampler, num_workers=4, pin_memory=True) - val_loader = DataLoader(val_dataset, batch_size=config.batch_size, sampler=val_sampler, num_workers=4, pin_memory=True) - - # --- 2. Model, Optimizer, and Loss Initialization --- - if is_main_process: - print(f"Initializing model on {config.device}...") - model = TimeAwareGPT2( - vocab_size=vocab_size, - n_embd=config.n_embd, - n_layer=config.n_layer, - n_head=config.n_head, - pdrop=config.pdrop, - token_pdrop=config.token_pdrop - ).to(config.device) - model = DDP(model, device_ids=[local_rank]) - - if is_main_process: - print(f"Model initialized with {model.module.get_num_params():.2f}M trainable parameters.") - - loss_fn = CombinedLoss(config.ignored_token_ids) - optimizer = AdamW(model.parameters(), lr=config.lr_initial, weight_decay=config.weight_decay, betas=(0.9, 0.99)) - - # --- 3. Training Loop --- - best_val_loss = float('inf') - patience_counter = 0 - - train_losses_ce, train_losses_surv, train_losses_total = [], [], [] - val_losses_ce, val_losses_surv, val_losses_total = [], [], [] - - if is_main_process: - print("Starting training...") - for epoch in range(config.max_epoch): - train_sampler.set_epoch(epoch) # Important for shuffling - - if epoch < config.warmup_epochs: - lr = config.lr_initial - else: - progress = (epoch - config.warmup_epochs) / (config.max_epoch - config.warmup_epochs) - lr = config.lr_final + 0.5 * (config.lr_initial - config.lr_final) * (1 + math.cos(math.pi * progress)) - - for param_group in optimizer.param_groups: - param_group['lr'] = lr - - model.train() - train_loss_ce_acc, train_loss_surv_acc = 0.0, 0.0 - train_steps = 0 - - pbar = tqdm.tqdm(train_loader, desc=f"Epoch {epoch+1}/{config.max_epoch} [Train]", disable=not is_main_process) - for event_seq, time_seq in pbar: - event_seq, time_seq = event_seq.to(config.device), time_seq.to(config.device) - - input_events, input_times = event_seq[:, :-1], time_seq[:, :-1] - target_events, target_wait_times = event_seq[:, 1:], (time_seq[:, 1:] - time_seq[:, :-1]).float() - - logits = model(input_events, input_times) - loss_ce, loss_survival = loss_fn(logits, target_events, target_wait_times) - loss = loss_ce + loss_survival - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - train_loss_ce_acc += loss_ce.item() - train_loss_surv_acc += loss_survival.item() - train_steps += 1 - if is_main_process: - pbar.set_postfix({'loss_ce': f'{loss_ce.item():.4f}', 'loss_surv': f'{loss_survival.item():.4f}', 'lr': f'{lr:.2e}'}) - - 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) - - model.eval() - val_loss_ce_acc, val_loss_surv_acc = 0.0, 0.0 - val_steps = 0 - - with torch.no_grad(): - pbar_val = tqdm.tqdm(val_loader, desc=f"Epoch {epoch+1}/{config.max_epoch} [Val]", disable=not is_main_process) - for event_seq, time_seq in pbar_val: - event_seq, time_seq = event_seq.to(config.device), time_seq.to(config.device) - - input_events, input_times = event_seq[:, :-1], time_seq[:, :-1] - target_events, target_wait_times = event_seq[:, 1:], (time_seq[:, 1:] - time_seq[:, :-1]).float() - - logits = model(input_events, input_times) - loss_ce, loss_survival = loss_fn(logits, target_events, target_wait_times) - - val_loss_ce_acc += loss_ce.item() - val_loss_surv_acc += loss_survival.item() - val_steps += 1 - if is_main_process: - pbar_val.set_postfix({'loss_ce': f'{loss_ce.item():.4f}', 'loss_surv': f'{loss_survival.item():.4f}'}) - - 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) - - if is_main_process: - 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" - f" Val Loss: {total_val_loss:.4f} (CE: {avg_val_loss_ce:.4f}, Surv: {avg_val_loss_surv:.4f})\n" - f" Learning Rate: {lr:.6f}") - - 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}. Saving checkpoint...") - torch.save(model.module.state_dict(), checkpoint_filename) - else: - 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: - if is_main_process: - print("\nEarly stopping triggered due to no improvement in validation loss.") - break - - if is_main_process: - if best_val_loss != float('inf'): - print(f"\nTraining finished. Loading best model from checkpoint with validation loss {best_val_loss:.4f}.") - # Load the best weights into the module before saving the final model - model.module.load_state_dict(torch.load(checkpoint_filename)) - print(f"Saving final best model to {model_filename}") - torch.save(model.module.state_dict(), model_filename) - else: - print("\nTraining finished. No best model to save as validation loss never improved.") - - losses_filename = f"losses_n_embd_{config.n_embd}_n_layer_{config.n_layer}_n_head_{config.n_head}.txt" - with open(losses_filename, 'w') as f: - f.write("epoch,train_loss_ce,train_loss_surv,train_loss_total,val_loss_ce,val_loss_surv,val_loss_total\n") - for i in range(len(train_losses_total)): - f.write(f"{i+1},{train_losses_ce[i]},{train_losses_surv[i]},{train_losses_total[i]},{val_losses_ce[i]},{val_losses_surv[i]},{val_losses_total[i]}\n") - print(f"\nLosses saved to {losses_filename}") - - num_epochs = len(train_losses_total) - epochs = range(1, num_epochs + 1) - plt.figure(figsize=(18, 5)) - - 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) - - 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) - - 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") - - cleanup_ddp() - -if __name__ == '__main__': - main()