import os
import random
import subprocess
import argparse
import time
import numpy as np
import pandas as pd
from sksurv.metrics import concordance_index_censored
import torch
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import Dataset
from model import HECTOR
from im4MEC import Im4MEC
from utils import *
from utils_loss import NLLSurvLoss
def set_seed():
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
def seed_worker(worker_id):
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
def evaluate_model(epoch, model, model_mol, device, loader, n_bins, writer, loss_fn, bins_values, train_BS, test_BS):
model.eval()
eval_loss = 0.
all_survival_probs = np.zeros((len(loader), n_bins))
all_risk_scores = np.zeros((len(loader))) # This is the computed risk score.
all_censorships = np.zeros((len(loader))) # This is the binary censorship status: 1 censored; 0 uncensored (reccured).
all_event_times = np.zeros((len(loader)))
with torch.no_grad():
for batch_idx, (data, features_flattened, label, event_time, censorship, stage, _) in enumerate(loader):
data, label, censorship, stage = data.to(device), label.to(device), censorship.to(device), stage.to(device)
_, _, Y_hat, _, _ = model_mol(features_flattened.to(device))
hazards_prob, survival_prob, Y_hat, _, _ = model(data, stage, Y_hat.squeeze(1)) # Returns hazards, survival, Y_hat, A_raw, M.
# We can emphasize on the contribution of uncensored patient cases only in training by minimizing a weighted sum of the 2 losses
loss = loss_fn(hazards=hazards_prob, S=survival_prob, Y=label, c=censorship, alpha=0)
eval_loss += loss.item()
risk = -torch.sum(survival_prob, dim=1).cpu().numpy()
all_risk_scores[batch_idx] = risk
all_censorships[batch_idx] = censorship.cpu().numpy()
all_event_times[batch_idx] = event_time
all_survival_probs[batch_idx] = survival_prob.cpu().numpy()
eval_loss /= len(loader)
# Compute a few survival metrics.
c_index = concordance_index_censored(
event_indicator=(1-all_censorships).astype(bool),
event_time=all_event_times,
estimate=all_risk_scores, tied_tol=1e-08)[0]
# Years of interest can be adapted in utils.py
(BS, years_of_interest), (IBS, yearI_of_interest, yearF_of_interest), (_, meanAUC), (c_index_ipcw) = compute_surv_metrics_eval(bins_values, all_survival_probs, all_risk_scores, train_BS, test_BS)
print(f'Eval epoch: {epoch}, loss: {eval_loss}, c_index: {c_index}, BS at each {years_of_interest}Y: {BS}, IBS and mean cumAUC from {yearI_of_interest}Y to {yearF_of_interest}Y: {IBS} and {meanAUC}')
writer.add_scalar("Loss/eval", eval_loss, epoch)
writer.add_scalar("C_index/eval", c_index, epoch)
for i in range(len(years_of_interest)):
writer.add_scalar(f"eval_metrics/BS_{str(years_of_interest[i])}Y", BS[i], epoch)
writer.add_scalar(f"eval_metrics/IBS_{str(yearI_of_interest)}Y-{str(yearF_of_interest)}Y", IBS, epoch)
writer.add_scalar(f"eval_metrics/meanAUC_{str(yearI_of_interest)}Y-{str(yearF_of_interest)}Y", meanAUC, epoch)
return eval_loss, c_index, (BS, IBS, meanAUC, c_index_ipcw)
def train_one_epoch(epoch, model, model_mol, device, train_loader, optimizer, n_bins, writer, loss_fn):
model.train()
epoch_start_time = time.time()
train_loss = 0.
all_risk_scores = np.zeros((len(train_loader))) # Computed risk score.
all_censorships = np.zeros((len(train_loader))) # Binary censorship status: 1 censored; 0 uncensored.
all_event_times = np.zeros((len(train_loader))) # Real t event time or last follow-up.
batch_start_time = time.time()
for batch_idx, (data, features_flattened, label, event_time, censorship, stage, _) in enumerate(train_loader):
data_load_duration = time.time() - batch_start_time
data, label, censorship, stage = data.to(device), label.to(device), censorship.to(device), stage.to(device)
# To get the image-based molecular class, non-merged features were used as this model was trained with way.
# Merged features could be used alternatively.
_, _, Y_hat, _, _ = model_mol(features_flattened.to(device))
# Returns hazards, survival, Y_hat, A_raw, M.
hazards_prob, survival_prob, Y_hat, _, _ = model(data, stage, Y_hat.squeeze(1))
# Loss.
loss = loss_fn(hazards=hazards_prob, S=survival_prob, Y=label, c=censorship)
train_loss += loss.item()
# Store outputs.
risk = -torch.sum(survival_prob, dim=1).detach().cpu().numpy()
all_risk_scores[batch_idx] = risk
all_censorships[batch_idx] = censorship.item()
all_event_times[batch_idx] = event_time
# Backward pass.
loss.backward()
# Step.
optimizer.step()
optimizer.zero_grad()
batch_duration = time.time() - batch_start_time
batch_start_time = time.time()
writer.add_scalar("duration/data_load", data_load_duration, epoch)
writer.add_scalar("duration/batch", batch_duration, epoch)
epoch_duration = time.time() - epoch_start_time
print(f"Finished training on epoch {epoch} in {epoch_duration:.2f}s")
train_loss /= len(train_loader)
train_c_index = concordance_index_censored(
event_indicator=(1-all_censorships).astype(bool),
event_time=all_event_times,
estimate=all_risk_scores, tied_tol=1e-08)[0]
print(f'Epoch: {epoch}, epoch_duration : {epoch_duration}, train_loss: {train_loss}, train_c_index: {train_c_index}')
filepath = os.path.join(writer.log_dir, f"{epoch}_checkpoint.pt")
print(f"Saving model to {filepath}")
torch.save(model.state_dict(), filepath)
writer.add_scalar("duration/epoch", epoch_duration, epoch)
writer.add_scalar("LR", get_lr(optimizer), epoch)
writer.add_scalar("Loss/train", train_loss, epoch)
writer.add_scalar("C_index/train", train_c_index, epoch)
def run_train_eval_loop(train_loader, val_loader, loss_fn, hparams, run_id, BS_data, checkpoint_model_molecular):
writer = SummaryWriter(os.path.join("./runs", run_id))
device = torch.device("cuda")
n_bins = hparams["n_bins"]
model = HECTOR(
input_feature_size=hparams["input_feature_size"],
precompression_layer=hparams["precompression_layer"],
feature_size_comp=hparams["feature_size_comp"],
feature_size_attn=hparams["feature_size_attn"],
postcompression_layer=hparams["postcompression_layer"],
feature_size_comp_post=hparams["feature_size_comp_post"],
dropout=True,
p_dropout_fc=hparams["p_dropout_fc"],
p_dropout_atn=hparams["p_dropout_atn"],
n_classes=n_bins,
input_stage_size=hparams["input_stage_size"],
embedding_dim_stage=hparams["embedding_dim_stage"],
depth_dim_stage=hparams["depth_dim_stage"],
act_fct_stage=hparams["act_fct_stage"],
dropout_stage=hparams["dropout_stage"],
p_dropout_stage=hparams["p_dropout_stage"],
input_mol_size=4,
embedding_dim_mol=hparams["embedding_dim_mol"],
depth_dim_mol=hparams["depth_dim_mol"],
act_fct_mol=hparams["act_fct_mol"],
dropout_mol=hparams["dropout_mol"],
p_dropout_mol=hparams["p_dropout_mol"],
fusion_type=hparams["fusion_type"],
use_bilinear=hparams["use_bilinear"],
gate_hist=hparams["gate_hist"],
gate_stage=hparams["gate_stage"],
gate_mol=hparams["gate_mol"],
scale=hparams["scale"],
).to(device)
print('model')
print_model(model)
# This model is instance with the trained weights towards molecular classification and will be used in inference mode only.
# NOTE: it is important that the molecular model, here im4MEC, has been trained on the same patients as training to avoid patient-level information leakage.
model_mol = Im4MEC(
input_feature_size=hparams["input_feature_size"],
precompression_layer=True,
feature_size_comp=hparams["feature_size_comp_molecular"],
feature_size_attn=hparams["feature_size_attn_molecular"],
n_classes=hparams["n_classes_molecular"],
dropout=True, # Not used in inference.
p_dropout_fc=0.25,
p_dropout_atn=0.25,
).to(device)
msg = model_mol.load_state_dict(torch.load(checkpoint_model_molecular, map_location=device), strict=True)
print(msg)
for p in model_mol.parameters():
p.requires_grad = False
print(f"HECTOR and plugged-in im4MEC are built and checkpoints loaded")
model_mol.eval()
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()),
lr=hparams["initial_lr"],
weight_decay=hparams["weight_decay"],
)
# Using a multi-step LR decay routine.
milestones = [int(x) for x in hparams["milestones"].split(",")]
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=milestones, gamma=hparams["gamma_lr"]
)
monitor_tracker = MonitorBestModelEarlyStopping(
patience=hparams["earlystop_patience"],
min_epochs=hparams["earlystop_min_epochs"],
saving_checkpoint=True,
)
for epoch in range(hparams["max_epochs"]):
train_one_epoch(epoch, model, model_mol, device, train_loader, optimizer, n_bins, writer, loss_fn)
# Evaluation on validation set.
print("Evaluating model on validation set...")
eval_loss, eval_cindex, eval_other_metrics = evaluate_model(epoch, model, model_mol, device, val_loader, n_bins, writer, loss_fn, hparams["bins_values"], *BS_data)
monitor_tracker(epoch, eval_loss, eval_cindex, eval_other_metrics, model, writer.log_dir)
# Update LR decay.
scheduler.step()
if monitor_tracker.early_stop:
print(f"Early stop criterion reached. Broke off training loop after epoch {epoch}.")
break
# Log the hyperparameters of the experiments.
runs_history = {
"run_id" : run_id,
"best_epoch_CI" : monitor_tracker.best_epoch_CI,
"best_CI_score" : monitor_tracker.best_CI_score,
"best_epoch_loss": monitor_tracker.best_epoch_loss,
"best_evalLoss" : monitor_tracker.eval_loss_min,
"BS" : monitor_tracker.best_metrics_score[0],
"IBS" : monitor_tracker.best_metrics_score[1],
"cumMeanAUC" : monitor_tracker.best_metrics_score[2],
"CI_ipwc" : monitor_tracker.best_metrics_score[3],
**hparams,
}
with open('runs_history.txt', 'a') as filehandle:
for _, value in runs_history.items():
filehandle.write('%s;' % value)
filehandle.write('\n')
writer.close()
def prepare_datasets(args):
df = pd.read_csv(args.manifest)
n_bins = len(df['disc_label'].unique())
assert n_bins == args.n_bins, 'mismatch between the number of bins passed in args and classes in dataset'
bins_values = get_bins_time_value(df, n_bins, time_col_name='recurrence_years', label_time_col_name='disc_label')
assert len(bins_values)==n_bins
print(f'Read {args.manifest} dataset containing {len(df)} samples with {n_bins} bins of following values {bins_values}')
# NOTE: you may need to use the two lines below depending on how the category is listed in the csv file.
#df.stage = df.stage.apply(lambda x : 'III' if 'III' in x else ('II' if 'II' in x else 'I')).astype("category")
#df.stage = pd.Categorical(df['stage'], categories=['I', 'II', 'III'], ordered=True).codes
print(f'stage taxonomy used: {df.stage.unique()}')
try:
training_set = df[df["split"] == "training"]
validation_set = df[df["split"] == "validation"]
except:
raise Exception(
f"Could not find training and validation splits in {args.manifest}"
)
train_split = FeatureBagsDataset(df=training_set,
data_dir=args.data_dir,
input_feature_size=args.input_feature_size,
stage_class=len(training_set.stage.unique()))
val_split = FeatureBagsDataset(df=validation_set,
data_dir=args.data_dir,
input_feature_size=args.input_feature_size,
stage_class=len(validation_set.stage.unique()))
# To compute the Brier score (BS), you need a specific format of censorship and times.
_, train_BS = get_survival_data_for_BS(training_set, time_col_name='recurrence_years')
_, test_BS = get_survival_data_for_BS(validation_set, time_col_name='recurrence_years')
return train_split, val_split, train_BS, test_BS, bins_values, len(df.stage.unique())
def main(args):
# Set random seed for some degree of reproducibility. See PyTorch docs on this topic for caveats.
# https://pytorch.org/docs/stable/notes/randomness.html#reproducibility
set_seed()
if not torch.cuda.is_available():
raise Exception(
"No CUDA device available. Training without one is not feasible."
)
git_sha = subprocess.check_output(["git", "describe", "--always"]).strip().decode("utf-8")
train_run_id = f"{git_sha}_hp{args.hp}_{time.strftime('%Y%m%d-%H%M')}"
train_split, val_split, train_BS, test_BS, bins_values, stage_taxonomy = prepare_datasets(args)
print(f"=> Run ID {train_run_id}")
print(f"=> Training on {len(train_split)} samples")
print(f"=> Validating on {len(val_split)} samples")
base_hparams = dict(
# Preprocessing settings. This should be changed with the dataset called accordingly.
# Storing values here for readibility.
n_bins=args.n_bins, # Partion on the continuous time scale.
bins_values=bins_values,
input_feature_size=args.input_feature_size,
features_extraction=os.path.dirname(args.data_dir),
# Settings that be changed in the loop:
# Training.
sampling_method="random",
max_epochs=100,
earlystop_warmup=0,
earlystop_patience=30,
earlystop_min_epochs=30,
# Loss.
alpha_surv = 0.0,
# Optimizer.
initial_lr=0.00003,
milestones="2, 5, 15, 25",
gamma_lr=0.1,
weight_decay=0.00001,
# Model architecture parameters. See model class for details.
precompression_layer=True,
feature_size_comp=512,
feature_size_attn=256,
postcompression_layer=True,
feature_size_comp_post=128,
p_dropout_fc=0.25,
p_dropout_atn=0.25,
# Model of molecular classification. In our case only inference is used.
n_classes_molecular=args.n_classes_molecular,
feature_size_comp_molecular=args.feature_size_comp_molecular,
feature_size_attn_molecular=args.feature_size_attn_molecular,
# Fusion parameters.
input_stage_size=stage_taxonomy,
embedding_dim_stage=16,
depth_dim_stage=1,
act_fct_stage='elu',
dropout_stage=True,
p_dropout_stage=0.25,
embedding_dim_mol=16,
depth_dim_mol=1,
act_fct_mol='elu',
dropout_mol=True,
p_dropout_mol=0.25,
fusion_type='bilinear',
use_bilinear=[True,True,True],
gate_hist=True,
gate_stage=True,
gate_mol=True,
scale=[2,1,1],
)
hparam_sets = [
{
**base_hparams,
},
]
hps = hparam_sets[args.hp]
train_loader, val_loader = define_data_sampling(
train_split,
val_split,
method=hps["sampling_method"],
workers=args.workers,
)
run_train_eval_loop(
train_loader=train_loader,
val_loader=val_loader,
loss_fn = NLLSurvLoss(alpha=hps["alpha_surv"]), # Used the Negative log likelihood loss.
hparams=hps,
run_id=train_run_id,
BS_data = (train_BS, test_BS),
checkpoint_model_molecular=args.checkpoint_model_molecular,
)
print("Finished training.")
def get_args_parser():
parser = argparse.ArgumentParser('Training script', add_help=False)
parser.add_argument(
"--manifest",
type=str,
help="CSV file listing all slides, their labels, and which split (train/test/val) they belong to.",
)
parser.add_argument(
"--n_bins",
type=int,
help="Number of time intervals used to create the time labels. It should be the same as the manifest.",
)
parser.add_argument(
"--data_dir",
type=str,
help="Directory where all *_features.h5 files are stored",
)
parser.add_argument(
"--input_feature_size",
help="The size of the input features from the feature bags. Recommend going by blocks from these output size [96, 96, 192, 192, 384, 384, 384, 384, 768, 768]",
type=int,
required=True,
)
parser.add_argument(
"--checkpoint_model_molecular",
type=str,
default='',
help="Path to checkpoint of im4MEC",
)
parser.add_argument(
"--n_classes_molecular",
type=int,
required=True,
help="",
)
parser.add_argument(
"--feature_size_comp_molecular",
type=int,
required=True,
help="Size of the model of the trained im4MEC. See in im4MEC.py",
)
parser.add_argument(
"--feature_size_attn_molecular",
type=int,
required=True,
help="Size of the model of the trained im4MEC. See in im4MEC.py",
)
parser.add_argument(
"--workers",
help="The number of workers to use for the data loaders.",
type=int,
default=4,
)
parser.add_argument(
"--hp",
type=int,
required=True,
)
return parser
if __name__ == "__main__":
parser = argparse.ArgumentParser('Training script', parents=[get_args_parser()])
args = parser.parse_args()
main(args)
Datasets
Models
