import dataset
import utils
from utils import EarlyStopping, LRScheduler
import os
import pandas as pd
import argparse
import torch.backends.cudnn as cudnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
import time
parser = argparse.ArgumentParser(description='PET lymphoma classification')
#I/O PARAMS
parser.add_argument('--output', type=str, default='results', help='name of output folder (default: "results")')
#MODEL PARAMS
parser.add_argument('--normalize', action='store_true', default=False, help='normalize images')
parser.add_argument('--checkpoint', default='', type=str, help='model checkpoint if any (default: none)')
parser.add_argument('--resume', action='store_true', default=False, help='resume from checkpoint')
#OPTIMIZATION PARAMS
parser.add_argument('--optimizer', default='sgd', type=str, help='The optimizer to use (default: sgd)')
parser.add_argument('--lr', type=float, default=1e-3, help='learning rate (default: 1e-3)')
parser.add_argument('--lr_anneal', type=int, default=15, help='period for lr annealing (default: 15). Only works for SGD')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum (default: 0.9)')
parser.add_argument('--wd', default=1e-4, type=float, help='weight decay (default: 1e-4)')
#TRAINING PARAMS
parser.add_argument('--split_index', default=0, type=int, metavar='INT', choices=list(range(0,20)),help='which split index (default: 0)')
parser.add_argument('--run', default=1, type=int, metavar='INT', help='repetition run with same settings (default: 1)')
parser.add_argument('--batch_size', type=int, default=50, help='how many images to sample per slide (default: 50)')
parser.add_argument('--nepochs', type=int, default=40, help='number of epochs (default: 40)')
parser.add_argument('--workers', default=10, type=int, help='number of data loading workers (default: 10)')
parser.add_argument('--augm', default=0, type=int, choices=[0,1,2,3,12,4,5,14,34,45], help='augmentation procedure 0=none,1=flip,2=rot,3=flip LR, 12=flip+rot, 4=scale, 5=noise, 14=flip+scale, 34=flipLR+scale, 45=scale+noise (default: 0)')
parser.add_argument('--balance', action='store_true', default=False, help='balance dataset (balance loss)')
parser.add_argument('--lr_scheduler', action='store_true',default=False, help='decrease LR on platau')
parser.add_argument('--early_stopping', action='store_true',default=False, help='use early stopping')
def main():
### Get user input
global args
args = parser.parse_args()
print(args)
best_auc = 0.
### Output directory and files
if not os.path.isdir(args.output):
try:
os.mkdir(args.output)
except OSError:
print ('Creation of the output directory "{}" failed.'.format(args.output))
else:
print ('Successfully created the output directory "{}".'.format(args.output))
### Get model
model = utils.get_model()
if args.checkpoint:
ch = torch.load(args.checkpoint)
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in ch['state_dict'].items() if k in model_dict}
print('Loaded [{}/{}] keys from checkpoint'.format(len(pretrained_dict),len(model_dict)))
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
if args.resume:
ch = torch.load( os.path.join(args.output,'checkpoint_split'+str(args.split_index)+'_run'+str(args.run)+'.pth') )
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in ch['state_dict'].items() if k in model_dict}
print('Loaded [{}/{}] keys from checkpoint'.format(len(pretrained_dict),len(model_dict)))
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
### Set optimizer
optimizer = utils.create_optimizer(model, args.optimizer, args.lr, args.momentum, args.wd)
if args.resume and 'optimizer' in ch:
optimizer.load_state_dict(ch['optimizer'])
print('Loaded optimizer state')
cudnn.benchmark = True
### Augmentations
flipHorVer = dataset.RandomFlip()
flipLR = dataset.RandomFlipLeftRight()
rot90 = dataset.RandomRot90()
scale = dataset.RandomScale()
noise = dataset.RandomNoise()
if args.augm==0:
transform = None
elif args.augm==1:
transform = transforms.Compose([flipHorVer])
elif args.augm==2:
transform = transforms.Compose([rot90])
elif args.augm==3:
transform = transforms.Compose([flipLR])
elif args.augm==12:
transform = transforms.Compose([flipHorVer,rot90])
elif args.augm==4:
transform = transforms.Compose([scale])
elif args.augm==5:
transform = transforms.Compose([noise])
elif args.augm==14:
transform = transforms.Compose([flip,scale])
elif args.augm==34:
transform = transforms.Compose([flipLR,scale])
elif args.augm==45:
transform = transforms.Compose([scale,noise])
### Set datasets
train_dset,trainval_dset,val_dset,_,balance_weight_neg_pos = dataset.get_datasets_singleview(transform,args.normalize,args.balance,args.split_index)
print('Datasets train:{}, val:{}'.format(len(train_dset.df),len(val_dset.df)))
### Set loss criterion
if args.balance:
w = torch.Tensor(balance_weight_neg_pos)
print('Balance loss with weights:',balance_weight_neg_pos)
criterion = nn.BCEWithLogitsLoss(pos_weight=w).cuda()
else:
criterion = nn.BCEWithLogitsLoss().cuda()
### Early stopping
if args.lr_scheduler:
print('INFO: Initializing learning rate scheduler')
lr_scheduler = LRScheduler(optimizer)
if args.resume and 'lr_scheduler' in ch:
lr_scheduler.lr_scheduler.load_state_dict(ch['lr_scheduler'])
print('Loaded lr_scheduler state')
if args.early_stopping:
print('INFO: Initializing early stopping')
early_stopping = EarlyStopping()
if args.resume and 'early_stopping' in ch:
early_stopping.best_loss = ch['early_stopping']['best_loss']
early_stopping.counter = ch['early_stopping']['counter']
early_stopping.min_delta = ch['early_stopping']['min_delta']
early_stopping.patience = ch['early_stopping']['patience']
early_stopping.early_stop = ch['early_stopping']['early_stop']
print('Loaded early_stopping state')
### Set loaders
train_loader = torch.utils.data.DataLoader(train_dset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers)
trainval_loader = torch.utils.data.DataLoader(trainval_dset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)
val_loader = torch.utils.data.DataLoader(val_dset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)
### Set output files
convergence_name = 'convergence_split'+str(args.split_index)+'_run'+str(args.run)+'.csv'
if not args.resume:
fconv = open(os.path.join(args.output,convergence_name), 'w')
fconv.write('epoch,split,metric,value\n')
fconv.close()
### Main training loop
if args.resume:
epochs = range(ch['epoch']+1,args.nepochs+1)
else:
epochs = range(args.nepochs+1)
for epoch in epochs:
if args.optimizer == 'sgd':
utils.adjust_learning_rate(optimizer, epoch, args.lr_anneal, args.lr)
### Training logic
if epoch > 0:
loss = train(epoch, train_loader, model, criterion, optimizer)
else:
loss = np.nan
### Printing stats
fconv = open(os.path.join(args.output,convergence_name), 'a')
fconv.write('{},train,loss,{}\n'.format(epoch, loss))
fconv.close()
### Validation logic
# Evaluate on train data
train_probs = test(epoch, trainval_loader, model)
train_auc, train_ber, train_fpr, train_fnr = train_dset.errors(train_probs)
# Evaluate on validation set
val_probs = test(epoch, val_loader, model)
val_auc, val_ber, val_fpr, val_fnr = val_dset.errors(val_probs)
print('Epoch: [{}/{}]\tLoss: {:.6f}\tAUC: {:.4f}\t{:.4f}'.format(epoch, args.nepochs, loss, train_auc, val_auc))
fconv = open(os.path.join(args.output,convergence_name), 'a')
fconv.write('{},train,auc,{}\n'.format(epoch, train_auc))
fconv.write('{},train,ber,{}\n'.format(epoch, train_ber))
fconv.write('{},train,fpr,{}\n'.format(epoch, train_fpr))
fconv.write('{},train,fnr,{}\n'.format(epoch, train_fnr))
fconv.write('{},validation,auc,{}\n'.format(epoch, val_auc))
fconv.write('{},validation,ber,{}\n'.format(epoch, val_ber))
fconv.write('{},validation,fpr,{}\n'.format(epoch, val_fpr))
fconv.write('{},validation,fnr,{}\n'.format(epoch, val_fnr))
fconv.close()
### Create checkpoint dictionary
obj = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer' : optimizer.state_dict(),
'lr_scheduler' : lr_scheduler.lr_scheduler.state_dict(),
'early_stopping' : {'best_loss':early_stopping.best_loss,'counter':early_stopping.counter,'early_stop':early_stopping.early_stop,'min_delta': early_stopping.min_delta,'patience': early_stopping.patience},
'auc': val_auc,
}
### Save checkpoint
torch.save(obj, os.path.join(args.output,'checkpoint_split'+str(args.split_index)+'_run'+str(args.run)+'.pth'))
### Early stopping
if args.lr_scheduler:
lr_scheduler(-val_auc)
if args.early_stopping:
early_stopping(-val_auc)
if early_stopping.early_stop:
break
def test(epoch, loader, model):
# Set model in test mode
model.eval()
# Initialize probability vector
probs = torch.FloatTensor(len(loader.dataset)).cuda()
# Loop through batches
with torch.no_grad():
for i, (input,_) in enumerate(loader):
## Copy batch to GPU
input = input.cuda()
## Forward pass
y = model(input) #features, probabilities
p = F.softmax(y,dim=1)
## Clone output to output vector
probs[i*args.batch_size:i*args.batch_size+input.size(0)] = p.detach()[:,1].clone()
return probs.cpu().numpy()
def train(epoch, loader, model, criterion, optimizer):
# Set model in training mode
model.train()
# Initialize loss
running_loss = 0.
# Loop through batches
for i, (input,target) in enumerate(loader):
## Copy to GPU
input = input.cuda()
target_1hot = F.one_hot(target.long(),num_classes=2).cuda()
## Forward pass
y = model(input) #features, probabilities
## Calculate loss
loss = criterion(y, target_1hot.float())
## Optimization step
optimizer.zero_grad()
loss.backward()
optimizer.step()
## Store loss
running_loss += loss.item()*input.size(0)
return running_loss/len(loader.dataset)
if __name__ == '__main__':
main()
Datasets
Models
