import sys
sys.path.append('architectures/deeplab_3D/')
sys.path.append('architectures/unet_3D/')
sys.path.append('architectures/hrnet_3D/')
sys.path.append('architectures/experiment_nets_3D/')
sys.path.append('utils/')
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import scipy.misc
import os
from tqdm import *
import random
from random import randint
from docopt import docopt
import deeplab_resnet_3D
import unet_3D
import highresnet_3D
import exp_net_3D
import lossF
import PP
import augmentations as AUG
import nibabel as nib
import evalF as EF
import evalFP as EFP
docstr = """Write something here
Usage:
train.py [options]
Options:
-h, --help Print this message
--archId=<int> Architecture to run, 0 is DeepLab 3D, 1 is U-net3D, 2 is HRNet [default: 2]
--trainMethod=<int> 0 is full image, 1 is by patches (random), 2 is by patches (center pixel) [default: 1]
--lossFunction=<str> Loss function name. 'dice' option available [default: dice]
--imgSize=<str> Image size [default: 200x200x100]
--mainFolderPath=<str> Main folder path [default: ../Data/MS2017b/]
--patchSize=<int> Size of the patch [default: 60]
--patchSizeStage0=<int> Size of the patch at stage 0 [default: 41]
--namePostfix=<str> Postfix of flair. i.e. to use FLAIR_s postfix is _s. This also determines the train file [default: _200x200x100orig]
--modelPath=<str> Path of model to continue training on [default: none]
--NoLabels=<int> The number of different labels in training data, including background [default: 2]
--maxIter=<int> Maximum number of iterations [default: 20000]
--maxIterStage0=<int> Maximum number of iterations for stage 0 training [default: -1]
-i, --iterSize=<int> Num iters to accumulate gradients over [default: 1]
--lr=<float> Learning Rate [default: 0.0001]
--gpu0=<int> GPU number [default: 0]
--useGPU=<int> Use GPU or not [default: 0]
--experiment=<str> Specify experiment instead to run. e.g. 1x1x1x1x1x1_1_0 means 1 dilations all 6 blocks, with priv, no ASPP [default: None]
"""
args = docopt(docstr, version='v0.1')
print(args)
arch_id = int(args['--archId'])
train_method = int(args['--trainMethod'])
loss_name = args['--lossFunction']
img_dims = np.array(args['--imgSize'].split('x'), dtype=np.int64)
main_folder_path = args['--mainFolderPath']
patch_size = int(args['--patchSize'])
postfix = args['--namePostfix']
model_path = args['--modelPath']
num_labels = int(args['--NoLabels'])
max_iter = int(args['--maxIter'])
iter_size = int(args['--iterSize'])
base_lr = float(args['--lr'])
experiment = str(args['--experiment'])
gpu0 = int(args['--gpu0'])
useGPU = int(args['--useGPU'])
batch_size = 1
#img_dims = [197, 233, 189]
list_path = main_folder_path + 'train' + postfix + '.txt'
print('READING from ', list_path)
img_type_path = 'pre/FLAIR' + postfix + '.nii.gz'
gt_type_path = 'wmh' + postfix + '.nii.gz'
patch_size_stage0 = int(args['--patchSizeStage0'])
max_iter_stage0 = int(args['--maxIterStage0'])
iter_low = 1
iter_high = max_iter + 1
if model_path != 'none':
iter_low = int(model_path.split('iter_')[-1].replace('.pth','')) + 1
if iter_low >= iter_high:
print('Model already at ' + str(iter_low) + ' iterations. Change max iter size')
sys.exit()
num_labels2 = 209
#change to 0 to enable stage 0 patch learning
if num_labels == 2:
onlyLesions = True
else:
onlyLesions = False
if useGPU:
cudnn.enabled = True
else:
cudnn.enabled = False
if experiment != 'None':
snapshot_prefix = 'EXP3D' + '_' + experiment + '_' + loss_name + '_' + str(train_method)
else:
if arch_id == 0:
snapshot_prefix = 'DL3D_' + loss_name + '_' + str(train_method) + '_' + PP.getTime()
elif arch_id == 1:
snapshot_prefix = 'UNET3D_' + loss_name + '_' + str(train_method) + '_' + PP.getTime()
elif arch_id == 2:
snapshot_prefix = 'HR3D' + loss_name + '_' + str(train_method) + '_' + PP.getTime()
to_center_pixel = False
center_pixel_folder_path, locs_lesion, locs_other = (None, None, None)
if train_method == 2:
to_center_pixel = True
if not os.path.exists(os.path.join(main_folder_path, 'centerPixelPatches' + postfix + '_' + str(patch_size))):
print('Pixel patch folder does not exist')
sys.exit()
#load few files
img_list = PP.read_file(list_path)
results_folder = 'train_results/'
log_file_path = os.path.join(results_folder, 'logs', snapshot_prefix + '_log.txt')
model_file_path = os.path.join(results_folder, 'models', snapshot_prefix + '_best.pth')
logfile = open(log_file_path, 'w+')
info_run = "arch ID: {:d} | max iters: {:10d} | max iters stage 0 : {:10d} | train method : {} | lr : {}".format(arch_id, max_iter, max_iter_stage0, train_method, base_lr)
logfile.write(info_run + '\n')
logfile.flush()
def lr_poly(base_lr, iter,max_iter,power):
return base_lr*((1-float(iter)/max_iter)**(power))
def modelInit():
isPriv = False
if arch_id > 10:
isPriv = True
if experiment != 'None':
dilation_arr, isPriv, withASPP = PP.getExperimentInfo(experiment)
model = exp_net_3D.getExpNet(num_labels, dilation_arr, isPriv, NoLabels2 = num_labels2, withASPP = withASPP)
elif arch_id == 0:
model = deeplab_resnet_3D.Res_Deeplab(num_labels)
elif arch_id == 1:
model = unet_3D.UNet3D(1, num_labels)
elif arch_id == 2:
model = highresnet_3D.getHRNet(num_labels)
if model_path != 'none':
if useGPU:
#loading on GPU when model was saved on GPU
saved_state_dict = torch.load(model_path)
else:
#loading on CPU when model was saved on GPU
saved_state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
model.load_state_dict(saved_state_dict)
model.float()
model.eval() # use_global_stats = True
return model, isPriv
def trainModel(model):
if useGPU:
model.cuda(gpu0)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = base_lr)
optimizer.zero_grad()
print(model)
curr_val = 0
best_val = 0
val_change = False
loss_arr = np.zeros([iter_size])
loss_arr_i = 0
stage = 0
print('---------------')
print('STAGE ' + str(stage))
print('---------------')
for iter in range(iter_low, iter_high):
if iter > max_iter_stage0 and stage != 1:
print('---------------')
print('Stage 1')
print('---------------')
stage = 1
if train_method == 0:
img_b, label_b, _ = PP.extractImgBatch(batch_size, img_list, img_dims, onlyLesions,
main_folder_path = '../Data/MS2017b/')
elif train_method == 1 or train_method == 2:
if stage == 0:
batch_size = 1
img_b, label_b, _ = PP.extractPatchBatch(batch_size, patch_size_stage0, img_list, onlyLesions, center_pixel = to_center_pixel, main_folder_path = '../Data/MS2017b/', postfix=postfix)
else:
batch_size = 1
img_b, label_b, _ = PP.extractPatchBatch(batch_size, patch_size, img_list, onlyLesions, center_pixel = to_center_pixel, main_folder_path = '../Data/MS2017b/', postfix=postfix)
else:
print('Invalid training method format')
sys.exit()
if stage == 0:
img_b, label_b = AUG.augmentPatchLossLess([img_b, label_b])
img_b, label_b = AUG.augmentPatchLossy([img_b, label_b])
#img_b, label_b = AUG.augmentPatchLossless(img_b, label_b)
#img_b is of shape (batch_num) x 1 x dim1 x dim2 x dim3
#label_b is of shape (batch_num) x 1 x dim1 x dim2 x dim3
#batch_num should be 1 since too memory intensive
label_b = label_b.astype(np.int64)
#convert label from (batch_num x 1 x dim1 x dim2 x dim3)
# to ((batch_numxdim1*dim2*dim3) x 3) (one hot)
temp = label_b.reshape([-1])
label_b = np.zeros([temp.size, num_labels])
label_b[np.arange(temp.size),temp] = 1
label_b = torch.from_numpy(label_b).float()
imgs = torch.from_numpy(img_b).float()
if useGPU:
imgs, label_b = Variable(imgs).cuda(gpu0), Variable(label_b).cuda(gpu0)
else:
imgs, label_b = Variable(imgs), Variable(label_b)
#---------------------------------------------
#out size is (1, 3, dim1, dim2, dim3)
#---------------------------------------------
out = model(imgs)
out = out.permute(0,2,3,4,1).contiguous()
out = out.view(-1, num_labels)
#---------------------------------------------
#out size is (1 * dim1 * dim2 * dim3, 3)
#---------------------------------------------
#loss function
m = nn.Softmax()
loss = lossF.simple_dice_loss3D(m(out), label_b)
loss /= iter_size
loss.backward()
loss_val = loss.data.cpu().numpy()
loss_arr[loss_arr_i] = loss_val
loss_arr_i = (loss_arr_i + 1) % iter_size
if iter % 1 == 0:
if val_change:
print "iter = {:6d}/{:6d} Loss: {:1.6f} Val Score: {:1.6f} \r".format(iter-1, max_iter, float(loss_val)*iter_size, curr_val),
sys.stdout.flush()
print ""
val_change = False
print "iter = {:6d}/{:6d} Loss: {:1.6f} Val Score: {:1.6f} \r".format(iter, max_iter, float(loss_val)*iter_size, curr_val),
sys.stdout.flush()
if iter % 1000 == 0:
val_change = True
curr_val = EF.evalModelX(model, num_labels, postfix, main_folder_path, (train_method != 0), gpu0, useGPU, eval_metric = 'iou', patch_size = patch_size, extra_patch = 5)
if curr_val > best_val:
best_val = curr_val
print('\nSaving better model...')
torch.save(model.state_dict(), model_file_path)
logfile.write("iter = {:6d}/{:6d} Loss: {:1.6f} Val Score: {:1.6f} \n".format(iter, max_iter, np.sum(loss_arr), curr_val))
logfile.flush()
if iter % iter_size == 0:
optimizer.step()
optimizer.zero_grad()
del out, loss
def setupGIFVar(gif_b):
gif_b = gif_b.astype(np.int64)
gif_b = gif_b.reshape([-1])
gif_b = torch.from_numpy(gif_b).long()
if useGPU:
gif_b = Variable(gif_b).cuda(gpu0)
else:
gif_b = Variable(gif_b)
return gif_b
def trainModelPriv(model):
if useGPU:
model.cuda(gpu0)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = base_lr)
optimizer.zero_grad()
print(model)
curr_val1 = 0
curr_val2 = 0
best_val2 = 0
val_change = False
loss_arr1 = np.zeros([iter_size])
loss_arr2 = np.zeros([iter_size])
loss_arr_i = 0
stage = 0
print('---------------')
print('STAGE ' + str(stage))
print('---------------')
for iter in range(iter_low, iter_high):
if iter > max_iter_stage0 and stage != 1:
print('---------------')
print('Stage 1')
print('---------------')
stage = 1
if train_method == 0:
img_b, label_b, gif_b = PP.extractImgBatch(batch_size, img_list, img_dims, onlyLesions,
main_folder_path = '../Data/MS2017b/', with_priv = True)
elif train_method == 1 or train_method == 2:
if stage == 0:
batch_size = 5
img_b, label_b, gif_b = PP.extractPatchBatch(batch_size, patch_size_stage0, img_list, onlyLesions,
center_pixel = to_center_pixel,
main_folder_path = '../Data/MS2017b/',
postfix=postfix, with_priv= True)
else:
batch_size = 1
img_b, label_b, gif_b = PP.extractPatchBatch(batch_size, patch_size, img_list, onlyLesions,
center_pixel = to_center_pixel,
main_folder_path = '../Data/MS2017b/',
postfix=postfix, with_priv= True)
else:
print('Invalid training method format')
sys.exit()
img_b, label_b, gif_b = AUG.augmentPatchLossy([img_b, label_b, gif_b])
#img_b is of shape (batch_num) x 1 x dim1 x dim2 x dim3
#label_b is of shape (batch_num) x 1 x dim1 x dim2 x dim3
label_b = label_b.astype(np.int64)
#convert label from (batch_num x 1 x dim1 x dim2 x dim3)
# to ((batch_numxdim1*dim2*dim3) x 3) (one hot)
temp = label_b.reshape([-1])
label_b = np.zeros([temp.size, num_labels])
label_b[np.arange(temp.size),temp] = 1
label_b = torch.from_numpy(label_b).float()
imgs = torch.from_numpy(img_b).float()
if useGPU:
imgs, label_b = Variable(imgs).cuda(gpu0), Variable(label_b).cuda(gpu0)
else:
imgs, label_b = Variable(imgs), Variable(label_b)
gif_b = setupGIFVar(gif_b)
#---------------------------------------------
#out size is (1, 3, dim1, dim2, dim3)
#---------------------------------------------
#out1 is extra info
out1, out2 = model(imgs)
out1 = out1.permute(0,2,3,4,1).contiguous()
out1 = out1.view(-1, num_labels2)
out2 = out2.permute(0,2,3,4,1).contiguous()
out2 = out2.view(-1, num_labels)
#---------------------------------------------
#out size is (1 * dim1 * dim2 * dim3, 3)
#---------------------------------------------
m2 = nn.Softmax()
loss2 = lossF.simple_dice_loss3D(m2(out2), label_b)
m1 = nn.LogSoftmax()
loss1 = F.nll_loss(m1(out1), gif_b)
loss1 /= iter_size
loss2 /= iter_size
torch.autograd.backward([loss1, loss2])
loss_val1 = float(loss1.data.cpu().numpy())
loss_arr1[loss_arr_i] = loss_val1
loss_val2 = float(loss2.data.cpu().numpy())
loss_arr2[loss_arr_i] = loss_val2
loss_arr_i = (loss_arr_i + 1) % iter_size
if iter % 1 == 0:
if val_change:
print "iter = {:6d}/{:6d} Loss_main: {:1.6f} Loss_secondary: {:1.6f} Val Score: {:1.6f} Val Score secondary: {:1.6f} \r".format(iter-1, max_iter, loss_val2*iter_size, loss_val1*iter_size, curr_val2, curr_val1),
sys.stdout.flush()
print ""
val_change = False
print "iter = {:6d}/{:6d} Loss_main: {:1.6f} Loss_secondary: {:1.6f} Val Score main: {:1.6f} Val Score secondary: {:1.6f} \r".format(iter, max_iter, loss_val2*iter_size, loss_val1*iter_size, curr_val2, curr_val1),
sys.stdout.flush()
if iter % 2000 == 0:
val_change = True
curr_val1, curr_val2 = EFP.evalModelX(model, num_labels, num_labels2, postfix, main_folder_path, (train_method != 0), gpu0, useGPU, eval_metric = 'iou', patch_size = patch_size, extra_patch = 5, priv_eval = True)
if curr_val2 > best_val2:
best_val2 = curr_val2
torch.save(model.state_dict(), model_file_path)
print('\nSaving better model...')
logfile.write("iter = {:6d}/{:6d} Loss_main: {:1.6f} Loss_secondary: {:1.6f} Val Score main: {:1.6f} Val Score secondary: {:1.6f} \n".format(iter, max_iter, np.sum(loss_arr2), np.sum(loss_arr1), curr_val2, curr_val1))
logfile.flush()
if iter % iter_size == 0:
optimizer.step()
optimizer.zero_grad()
del out1, out2, loss1, loss2
if __name__ == "__main__":
model, with_priv = modelInit()
if with_priv:
trainModelPriv(model)
else:
trainModel(model)
logfile.close()
