File size: 15,778 Bytes
a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 6524cf7 a23d562 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 | #!/usr/bin/env python3
"""
Train one ResNet on one cross validation fold.
Paths may be given as arguments or come from the site config (see config.example.yaml):
python scripts/train/train.py resnet18_abrv -b 32 -n 60 -c 0 -d 0 -t -e
"""
import math
import os
import pickle
import tracemalloc
import torch.utils.data
import wandb
from torch import nn, optim
from torch.amp import GradScaler
import auto_detect_breast_mri.evaluation.metrics as eva
from auto_detect_breast_mri.config import get_config, resolve_path
from auto_detect_breast_mri.data import loaders
from auto_detect_breast_mri.data.transforms import (intensityAugmentation, minimalAugmentation,
pre_image_shape)
from auto_detect_breast_mri.models.checkpoints import load_pretrained_model, save_model
from auto_detect_breast_mri.models.resnets import model_names
from auto_detect_breast_mri.training.cli import generate_parser
from auto_detect_breast_mri.training.lr_scheduler import LRScheduler_Manager
from auto_detect_breast_mri.training.loops import (CPU, GPU, display_top, eval_epoch_subjects,
train_epoch_subjects)
from auto_detect_breast_mri.data.metadata import get_uka_metatensor
if __name__ == '__main__':
torch.manual_seed(31)
tracemalloc.start()
DEVICE = GPU
if not torch.cuda.is_available():
DEVICE = CPU
values = None # get_evaluation_metrics()
parser = generate_parser()
args = parser.parse_args()
print(args)
model_name = args.model_name
# Paths come from the command line; anything omitted is taken from the site config
# (see config.example.yaml and src/config.py).
path_base = resolve_path(args.data_path, "data_root", "root folder of the NIfTI data")
feature_path = resolve_path(args.feature_path, "metadata_file", "metadata export")
split_files_folder = resolve_path(args.split_files_folder, "split_root",
"folder holding the split files")
resulting_folder = resolve_path(args.output_path, "output_root", "output folder")
os.makedirs(resulting_folder, exist_ok=True)
# Read optional arguments if given
### a - augmentation
if args.augmentation and args.augmentation == 'intensity':
transform = intensityAugmentation
t = args.augmentation
else:
transform = minimalAugmentation
t = 'basic'
### b - batch_size
batch_size = 16
if args.batch_size and args.batch_size > 0:
batch_size = args.batch_size
### c, d - outer, inner cross validation fold
fold = None
subfold = None
if args.test_fold is not None:
fold = args.test_fold
if args.validation_fold is not None:
subfold = args.validation_fold
if args.dataloader is not None and args.dataloader != "subjects":
parser.error(f"Dataloadertype {args.dataloader} is not supported. Only 'subjects' is available.")
### e - eval
do_eval = args.evaluation
### f - fraction
if args.fraction_of_data and 0 < args.fraction_of_data < 1:
fraction = args.fraction_of_data
else:
fraction = -1.0
### l - learning rate
learning_rate = 0.0001
if args.learning_rate and 0 < args.learning_rate:
learning_rate = args.learning_rate
### m - mixed precision
use_mixed_precision = args.mixed_precision
### n - number of epochs
num_epochs = 60
if args.number_of_epochs:
num_epochs = args.number_of_epochs
### o - output path
model_path = None
if args.model_path:
if len(args.model_path) > 4 and args.model_path.endswith('.pth'):
model_path = args.model_path
print("use model: " + model_path)
else:
raise ValueError("Invalid model path: " + args.model_path)
else:
model_path = None
### p - patience
patience = None # number of epochs to wait without improvement before stopping
if args.patience and 0 < args.patience <= num_epochs:
patience = args.patience
### s - scheduler for learning rate
scheduler = None
scheduler_name = "No scheduler"
use_scheduler = False
if args.scheduler:
print("Scheduler parameter is activated.")
use_scheduler = True
### t - do training
train = args.train
## TRAIN each model
model = model_names.get(model_name)
print('Prepare model {} for training file under {}'.format(model_name, split_files_folder))
# possible options: model_abrv, model_full
if "abrv" in model_name:
protocol = 'abbreviated'
elif "sub" in model_name:
protocol = ['Sub_1']
elif "full" in model_name:
protocol = 'full'
else:
protocol = 'abbreviated'
print("Protocol not supported: " + model_name)
# Dataset and DataLoader
train_loader = []
eval_loader = []
feature_dataframe = get_uka_metatensor(0, feature_path)
if fold is not None and subfold is not None and train and do_eval:
print("using subjects loader")
train_loader, eval_loader, test_loader = loaders.get_multiple_subjects_dataloader(path_base,
feature_dataframe,
pre_image_shape,
protocol=protocol,
set_folder=split_files_folder,
transform=transform,
batch_size=batch_size,
stratified=True,
fold=fold,
subfold=subfold,
fraction=fraction)
else:
print("get individual dataloader")
split_file_path_sceleton = split_files_folder + "fold{}/stratified_{}"
if train:
subfold = 0 if subfold is None else subfold
split_file_path = split_file_path_sceleton.format(fold, f"training_set")
train_loader = loaders.get_subjects_dataloader(path_base, feature_dataframe,
pre_image_shape,
transform,
protocol,
split_file_path,
batch_size,
fraction,
fold,
subfold)
if do_eval:
split_file_path = split_file_path_sceleton.format(fold, f"evaluation_set")
# augmentation belongs to the training loader only; None falls back to the
# deterministic default transform, so validation metrics are not computed on
# randomly flipped and noised volumes
eval_loader = loaders.get_subjects_dataloader(path_base, feature_dataframe,
pre_image_shape,
None,
protocol,
split_file_path,
batch_size,
fraction,
fold,
subfold)
split_file_path = split_file_path_sceleton.format(fold, f"test_set")
test_loader = loaders.get_subjects_dataloader(path_base, feature_dataframe,
pre_image_shape,
None, # no augmentation at test time
protocol,
split_file_path,
batch_size,
1,
fold)
n_train_sample = len(train_loader.dataset) if train else 0
n_eval_sample = len(eval_loader.dataset) if do_eval else 0
snapshot = tracemalloc.take_snapshot()
display_top(snapshot=snapshot, title="after dataset creation")
weight_decay = 1e-2
# ------------ Initialize Model, Loss Function, Optimizer ------------
print("Use device: " + DEVICE)
model.to(DEVICE)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
print("use scheduler? " + str(use_scheduler))
if use_scheduler:
schedule_manager = LRScheduler_Manager(optimizer, num_epochs)
scheduler, scheduler_name = schedule_manager.get_scheduler(args.scheduler)
print(f"scheduler: {scheduler_name}")
if model_path is not None:
# load pretrained model:
checkpoint = load_pretrained_model(model, model_path, DEVICE)
model.load_state_dict(checkpoint)
lr = round(learning_rate, 7)
# Project/entity/mode come from the site config; no local path is logged, since run configs
# are visible to everybody with access to the wandb project.
wandb.init(**get_config().wandb_init_kwargs(), config={
"architecture": model_name,
"number of training samples": n_train_sample,
"number of evaluation samples": n_eval_sample,
"number of test samples": len(test_loader.dataset),
"protocol": protocol,
"optimizer": str(optimizer),
"weight decay": weight_decay,
"Augmentation": t + str(transform),
"batch size": batch_size,
"Fold": fold,
"train fraction": fraction,
"learning-rate": lr,
"Mixed Precision": use_mixed_precision,
"epochs": num_epochs,
"Pretrained": (model_path is not None),
"patience": patience,
"scheduler": scheduler_name,
"Machine": "HPC"
}, name="{}_b={}_l={}_n={}_t={}_split={}".format(model_name, batch_size, lr, num_epochs, t, fold))
print('Start TRAINING model {} with splits under {}'.format(model_name, split_files_folder))
### ----------------- Main loop ----------------- ###
scaler = None
if use_mixed_precision:
scaler = GradScaler(device=DEVICE)
best_auc = -math.inf
best_epoch = 0
frac_string = str(fraction).replace('-', '') # -1.0 should be 1.0
if train:
for epoch in range(num_epochs):
title_suffix = 'EPOCH {}'.format(epoch)
wandb.log({'Learning rate': optimizer.param_groups[0]['lr']})
# --- Training loop ---
train_loss = train_epoch_subjects(model, optimizer, criterion, train_loader, DEVICE, values,
show_figs=False, scaler=scaler, title_suffix=title_suffix)
if use_scheduler and not scheduler_name.startswith("ReduceLROnPlateau"):
scheduler.step()
# Visualization of evaluation measures
wandb.log({'Epoch': epoch, "average_loss/train": train_loss})
# --- Validation loop ---
if do_eval:
eval_loss, results_dict = eval_epoch_subjects(model, criterion, eval_loader, DEVICE, values, 'eval_',
scaler=scaler, title_suffix=title_suffix)
y_true, y_pred, y_probs = results_dict["y_true"], results_dict["y_pred"], results_dict["y_prob"]
sens, spec = eva.compute_sensitivity_specificity(y_true, y_pred, y_scores=y_probs,
prefix='Eval')
roc_results = eva.compute_roc_curve(y_true, y_probs, plot=False, title_suffix='Evaluation ')
roc_auc = roc_results['roc_auc_val'] if "roc_auc_val" in roc_results.keys() else math.nan
wandb.log({"sensitivity/Eval": sens,
"specificity/Eval": spec,
"AUC_ROC/Eval": float(roc_auc)})
print(f"Evaluation AUC(Youden) {roc_auc}.")
if use_scheduler and scheduler_name.startswith("ReduceLROnPlateau"):
scheduler.step(eval_loss)
if patience is not None:
if (roc_auc - best_auc) > 0:
# reached a better auc than before
best_auc = roc_auc
best_epoch = epoch
elif (epoch - best_epoch) > patience:
# patience is reached
print("Early stopping: {} was best AUC at epoch {}".format(best_auc, best_epoch))
break
if epoch % 10 == 0:
save_model(model, resulting_folder,
"{}_fold={}_subfold={}_frac={}_n={}".format(model_name, fold, subfold, frac_string, epoch))
if do_eval:
with open(
f'{resulting_folder}/results_{model_name}_fold{fold}_subfold{subfold}_frac{frac_string}-{epoch}',
'wb') as result_file:
pickle.dump(results_dict, result_file)
else:
title_suffix = 'Final (Evaluation) EPOCH'
# check performance on test set
test_loss, final_results_dict = eval_epoch_subjects(model, criterion, test_loader, DEVICE, values, 'Test_',
plot_tpr_fpr=resulting_folder + f"roc_{model_name}_{fold}.png",
scaler=scaler, title_suffix=title_suffix)
y_true, y_pred, y_probs = final_results_dict["y_true"], final_results_dict["y_pred"], final_results_dict["y_prob"]
sens, spec = eva.compute_sensitivity_specificity(y_true, y_pred, y_scores=y_probs, prefix='Test')
roc_results = eva.compute_roc_curve(y_true, y_probs, plot=True, title_suffix='Evaluation ')
#roc = roc_results['roc'], fpr = roc_results['fpr'], tpr = roc_results['tpr'], thresholds = roc_results['thresholds'], best_index = roc_results['best_index_val']
roc_auc = roc_results['roc_auc_val'] if "roc_auc_val" in roc_results.keys() else math.nan
wandb.log({"sensitivity/Test": sens,
"specificity/Test": spec,
"AUC_ROC/Test": float(roc_auc)})
print(f"Finished with test AUC (Youden) of {roc_auc}.")
# save model parameters and predictions
save_model(model, resulting_folder,
"{}_fold={}_subfold={}_frac={}_FINAL".format(model_name, fold, subfold, frac_string))
with open(f'{resulting_folder}/results_{model_name}_fold={fold}_frac={frac_string}', 'wb') as result_file:
pickle.dump(final_results_dict, result_file)
|