#run_experiment.py
#Copyright (c) 2020 Rachel Lea Ballantyne Draelos
#MIT License
#Permission is hereby granted, free of charge, to any person obtaining a copy
#of this software and associated documentation files (the "Software"), to deal
#in the Software without restriction, including without limitation the rights
#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#copies of the Software, and to permit persons to whom the Software is
#furnished to do so, subject to the following conditions:
#The above copyright notice and this permission notice shall be included in all
#copies or substantial portions of the Software.
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#SOFTWARE
import os
import timeit
import datetime
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
import torch, torch.nn as nn, torch.nn.functional as F
import torchvision
from torchvision import transforms, models, utils
import evaluate
from load_dataset import custom_datasets
#Set seeds
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed(0)
torch.cuda.manual_seed_all(0)
class DukeCTModel(object):
def __init__(self, descriptor, custom_net, custom_net_args,
loss, loss_args, num_epochs, patience, batch_size, device, data_parallel,
use_test_set, task, old_params_dir, dataset_class, dataset_args):
"""Variables:
<descriptor>: string describing the experiment
<custom_net>: class defining a model
<custom_net_args>: dictionary where keys correspond to custom net
input arguments, and values are the desired values
<loss>: 'bce' for binary cross entropy
<loss_args>: arguments to pass to the loss function if any
<num_epochs>: int for the maximum number of epochs to train
<patience>: number of epochs for which loss must fail to improve to
cause early stopping
<batch_size>: int for number of examples per batch
<device>: int specifying which device to use, or 'all' for all devices
<data_parallel>: if True then parallelize across available GPUs.
<use_test_set>: if True, then run model on the test set. If False, use
only the training and validation sets.
<task>:
'train_eval': train and evaluate a new model. 'evaluate' will
always imply use of the validation set. if <use_test_set> is
True, then 'evaluate' also includes calculation of test set
performance for the best validation epoch.
'predict_on_test': load a trained model and make predictions on
the test set using that model.
<old_params_dir>: this is only needed if <task>=='predict_on_test'. This
is the path to the parameters that will be loaded in to the model.
<dataset_class>: CT Dataset class for preprocessing the data
<dataset_args>: arguments for the dataset class specifying how
the data should be prepared."""
self.descriptor = descriptor
self.set_up_results_dirs()
self.custom_net = custom_net
self.custom_net_args = custom_net_args
self.loss = loss
self.loss_args = loss_args
self.num_epochs = num_epochs
self.batch_size = batch_size
print('self.batch_size=',self.batch_size)
#num_workers is number of threads to use for data loading
self.num_workers = int(batch_size*4) #batch_size 1 = num_workers 4. batch_size 2 = num workers 8. batch_size 4 = num_workers 16.
print('self.num_workers=',self.num_workers)
if self.num_workers == 1:
print('Warning: Using only one worker will slow down data loading')
#Set Device and Data Parallelism
if device in [0,1,2,3]: #i.e. if a GPU number was specified:
self.device = torch.device('cuda:'+str(device))
print('using device:',str(self.device),'\ndescriptor: ',self.descriptor)
elif device == 'all':
self.device = torch.device('cuda')
self.data_parallel = data_parallel
if self.data_parallel:
assert device == 'all' #use all devices when running data parallel
#Set Task
self.use_test_set = use_test_set
self.task = task
assert self.task in ['train_eval','predict_on_test']
if self.task == 'predict_on_test':
#overwrite the params dir that was created in the call to
#set_up_results_dirs() with the dir you want to load from
self.params_dir = old_params_dir
#Data and Labels
self.CTDatasetClass = dataset_class
self.dataset_args = dataset_args
#Get label meanings, a list of descriptive strings (list elements must
#be strings found in the column headers of the labels file)
self.set_up_label_meanings(self.dataset_args['label_meanings'])
if self.task == 'train_eval':
self.dataset_train = self.CTDatasetClass(setname = 'train', **self.dataset_args)
self.dataset_valid = self.CTDatasetClass(setname = 'valid', **self.dataset_args)
if self.use_test_set:
self.dataset_test = self.CTDatasetClass(setname = 'test', **self.dataset_args)
#Tracking losses and evaluation results
self.train_loss = np.zeros((self.num_epochs))
self.valid_loss = np.zeros((self.num_epochs))
self.eval_results_valid, self.eval_results_test = evaluate.initialize_evaluation_dfs(self.label_meanings, self.num_epochs)
#For early stopping
self.initial_patience = patience
self.patience_remaining = patience
self.best_valid_epoch = 0
self.min_val_loss = np.inf
#Run everything
self.run_model()
### Methods ###
def set_up_label_meanings(self,label_meanings):
if label_meanings == 'all': #get full list of all available labels
temp = custom_datasets.read_in_labels(self.dataset_args['label_type_ld'], 'valid')
self.label_meanings = temp.columns.values.tolist()
else: #use the label meanings that were passed in
self.label_meanings = label_meanings
print('label meanings ('+str(len(self.label_meanings))+' labels total):',self.label_meanings)
def set_up_results_dirs(self):
if not os.path.isdir('results'):
os.mkdir('results')
self.results_dir = os.path.join('results',datetime.datetime.today().strftime('%Y-%m-%d')+'_'+self.descriptor)
if not os.path.isdir(self.results_dir):
os.mkdir(self.results_dir)
self.params_dir = os.path.join(self.results_dir,'params')
if not os.path.isdir(self.params_dir):
os.mkdir(self.params_dir)
self.backup_dir = os.path.join(self.results_dir,'backup')
if not os.path.isdir(self.backup_dir):
os.mkdir(self.backup_dir)
def run_model(self):
if self.data_parallel:
self.model = nn.DataParallel(self.custom_net(**self.custom_net_args)).to(self.device)
else:
self.model = self.custom_net(**self.custom_net_args).to(self.device)
self.sigmoid = torch.nn.Sigmoid()
self.set_up_loss_function()
momentum = 0.99
print('Running with optimizer lr=1e-3, momentum='+str(round(momentum,2))+' and weight_decay=1e-7')
self.optimizer = torch.optim.SGD(self.model.parameters(), lr = 1e-3, momentum=momentum, weight_decay=1e-7)
train_dataloader = DataLoader(self.dataset_train, batch_size=self.batch_size, shuffle=True, num_workers = self.num_workers)
valid_dataloader = DataLoader(self.dataset_valid, batch_size=self.batch_size, shuffle=False, num_workers = self.num_workers)
if self.task == 'train_eval':
for epoch in range(self.num_epochs):
t0 = timeit.default_timer()
self.train(train_dataloader, epoch)
self.valid(valid_dataloader, epoch)
self.save_evals(epoch)
if self.patience_remaining <= 0:
print('No more patience (',self.initial_patience,') left at epoch',epoch)
print('--> Implementing early stopping. Best epoch was:',self.best_valid_epoch)
break
t1 = timeit.default_timer()
self.back_up_model_every_ten(epoch)
print('Epoch',epoch,'time:',round((t1 - t0)/60.0,2),'minutes')
if self.use_test_set: self.test(DataLoader(self.dataset_test, batch_size=self.batch_size, shuffle=False, num_workers = self.num_workers))
self.save_final_summary()
def set_up_loss_function(self):
if self.loss == 'bce':
self.loss_func = nn.BCEWithLogitsLoss() #includes application of sigmoid for numerical stability
def train(self, dataloader, epoch):
model = self.model.train()
epoch_loss, pred_epoch, gr_truth_epoch, volume_accs_epoch = self.iterate_through_batches(model, dataloader, epoch, training=True)
self.train_loss[epoch] = epoch_loss
self.plot_roc_and_pr_curves('train', epoch, pred_epoch, gr_truth_epoch)
print("{:5s} {:<3d} {:11s} {:.3f}".format('Epoch', epoch, 'Train Loss', epoch_loss))
def valid(self, dataloader, epoch):
model = self.model.eval()
with torch.no_grad():
epoch_loss, pred_epoch, gr_truth_epoch, volume_accs_epoch = self.iterate_through_batches(model, dataloader, epoch, training=False)
self.valid_loss[epoch] = epoch_loss
self.eval_results_valid = evaluate.evaluate_all(self.eval_results_valid, epoch,
self.label_meanings, gr_truth_epoch, pred_epoch)
self.early_stopping_check(epoch, pred_epoch, gr_truth_epoch, volume_accs_epoch)
print("{:5s} {:<3d} {:11s} {:.3f}".format('Epoch', epoch, 'Valid Loss', epoch_loss))
def early_stopping_check(self, epoch, val_pred_epoch, val_gr_truth_epoch, val_volume_accs_epoch):
"""Check whether criteria for early stopping are met and update
counters accordingly"""
val_loss = self.valid_loss[epoch]
if (val_loss < self.min_val_loss) or epoch==0: #then save parameters
self.min_val_loss = val_loss
check_point = {'params': self.model.state_dict(),
'optimizer': self.optimizer.state_dict()}
torch.save(check_point, os.path.join(self.params_dir, self.descriptor))
self.best_valid_epoch = epoch
self.patience_remaining = self.initial_patience
print('model saved, val loss',val_loss)
self.plot_roc_and_pr_curves('valid', epoch, val_pred_epoch, val_gr_truth_epoch)
self.save_all_pred_probs('valid', epoch, val_pred_epoch, val_gr_truth_epoch, val_volume_accs_epoch)
else:
self.patience_remaining -= 1
def back_up_model_every_ten(self, epoch):
"""Back up the model parameters every 10 epochs"""
if epoch % 10 == 0:
check_point = {'params': self.model.state_dict(),
'optimizer': self.optimizer.state_dict()}
torch.save(check_point, os.path.join(self.backup_dir, self.descriptor+'_ep_'+str(epoch)))
def test(self, dataloader):
epoch = self.best_valid_epoch
if self.data_parallel:
model = nn.DataParallel(self.custom_net(**self.custom_net_args)).to(self.device).eval()
else:
model = self.custom_net(**self.custom_net_args).to(self.device).eval()
params_path = os.path.join(self.params_dir,self.descriptor)
print('For test set predictions, loading model params from params_path=',params_path)
check_point = torch.load(params_path)
model.load_state_dict(check_point['params'])
with torch.no_grad():
epoch_loss, pred_epoch, gr_truth_epoch, volume_accs_epoch = self.iterate_through_batches(model, dataloader, epoch, training=False)
self.eval_results_test = evaluate.evaluate_all(self.eval_results_test, epoch,
self.label_meanings, gr_truth_epoch, pred_epoch)
self.plot_roc_and_pr_curves('test', epoch, pred_epoch, gr_truth_epoch)
self.save_all_pred_probs('test', epoch, pred_epoch, gr_truth_epoch, volume_accs_epoch)
print("{:5s} {:<3d} {:11s} {:.3f}".format('Epoch', epoch, 'Test Loss', epoch_loss))
def iterate_through_batches(self, model, dataloader, epoch, training):
epoch_loss = 0
#Initialize numpy arrays for storing results. examples x labels
#Do NOT use concatenation, or else you will have memory fragmentation.
num_examples = len(dataloader.dataset)
num_labels = len(self.label_meanings)
pred_epoch = np.zeros([num_examples,num_labels])
gr_truth_epoch = np.zeros([num_examples,num_labels])
volume_accs_epoch = np.empty(num_examples,dtype='U32') #need to use U32 to allow string of length 32
for batch_idx, batch in enumerate(dataloader):
data, gr_truth = self.move_data_to_device(batch)
self.optimizer.zero_grad()
if training:
out = model(data)
else:
with torch.set_grad_enabled(False):
out = model(data)
loss = self.loss_func(out, gr_truth)
if training:
loss.backward()
self.optimizer.step()
epoch_loss += loss.item()
torch.cuda.empty_cache()
#Save predictions and ground truth across batches
pred = self.sigmoid(out.data).detach().cpu().numpy()
gr_truth = gr_truth.detach().cpu().numpy()
start_row = batch_idx*self.batch_size
stop_row = min(start_row + self.batch_size, num_examples)
pred_epoch[start_row:stop_row,:] = pred #pred_epoch is e.g. [25355,80] and pred is e.g. [1,80] for a batch size of 1
gr_truth_epoch[start_row:stop_row,:] = gr_truth #gr_truth_epoch has same shape as pred_epoch
volume_accs_epoch[start_row:stop_row] = batch['volume_acc'] #volume_accs_epoch stores the volume accessions in the order they were used
#the following line to empty the cache is necessary in order to
#reduce memory usage and avoid OOM error:
torch.cuda.empty_cache()
return epoch_loss, pred_epoch, gr_truth_epoch, volume_accs_epoch
def move_data_to_device(self, batch):
"""Move data and ground truth to device."""
assert self.dataset_args['crop_type'] == 'single'
if self.dataset_args['crop_type'] == 'single':
data = batch['data'].to(self.device)
#Ground truth to device
gr_truth = batch['gr_truth'].to(self.device)
return data, gr_truth
def plot_roc_and_pr_curves(self, setname, epoch, pred_epoch, gr_truth_epoch):
outdir = os.path.join(self.results_dir,'curves')
if not os.path.isdir(outdir):
os.mkdir(outdir)
evaluate.plot_roc_curve_multi_class(label_meanings=self.label_meanings,
y_test=gr_truth_epoch, y_score=pred_epoch,
outdir = outdir, setname = setname, epoch = epoch)
evaluate.plot_pr_curve_multi_class(label_meanings=self.label_meanings,
y_test=gr_truth_epoch, y_score=pred_epoch,
outdir = outdir, setname = setname, epoch = epoch)
def save_all_pred_probs(self, setname, epoch, pred_epoch, gr_truth_epoch, volume_accs_epoch):
outdir = os.path.join(self.results_dir,'pred_probs')
if not os.path.isdir(outdir):
os.mkdir(outdir)
(pd.DataFrame(pred_epoch,columns=self.label_meanings,index=volume_accs_epoch.tolist())).to_csv(os.path.join(outdir, setname+'_predprob_ep'+str(epoch)+'.csv'))
(pd.DataFrame(gr_truth_epoch,columns=self.label_meanings,index=volume_accs_epoch.tolist())).to_csv(os.path.join(outdir, setname+'_grtruth_ep'+str(epoch)+'.csv'))
def save_evals(self, epoch):
evaluate.save(self.eval_results_valid, self.results_dir, self.descriptor+'_valid')
if self.use_test_set: evaluate.save(self.eval_results_test, self.results_dir, self.descriptor+'_test')
evaluate.plot_learning_curves(self.train_loss, self.valid_loss, self.results_dir, self.descriptor)
def save_final_summary(self):
evaluate.save_final_summary(self.eval_results_valid, self.best_valid_epoch, 'valid', self.results_dir)
if self.use_test_set: evaluate.save_final_summary(self.eval_results_test, self.best_valid_epoch, 'test', self.results_dir)
evaluate.clean_up_output_files(self.best_valid_epoch, self.results_dir)
