def get_model(fn_args):
"""
This function defines a Keras model and returns the model as a Keras object.
"""
base_model, layers, layer_names = create_base_model(
name=constants.BACKBONE_NAME,
weights=constants.PRETRAINED_WEIGHTS,
height=constants.HEIGHT,
width=constants.WIDTH,
include_top=False,
pooling=None)
model = DANet(n_classes=constants.N_MODEL_CLASSES,
base_model=base_model,
output_layers=layers,
backbone_trainable=constants.BACKBONE_TRAINABLE,
height=constants.HEIGHT,
width=constants.WIDTH).model()
opt = tf.keras.optimizers.SGD(learning_rate=0.2, momentum=0.9)
metrics = [IOUScore(threshold=0.5)]
categorical_focal_dice_loss = CategoricalFocalLoss(alpha=0.25,
gamma=2.0) + DiceLoss()
model.compile(
optimizer=opt,
loss=categorical_focal_dice_loss,
metrics=metrics,
)
return model
def __define_criterion(self,
class_weights,
delta_var,
delta_dist,
norm=2,
optimize_bg=False,
criterion='CE'):
assert criterion in ['CE', 'Dice', 'Multi', None]
smooth = 1.0
# Discriminative Loss
if self.use_instance_segmentation:
self.criterion_discriminative = DiscriminativeLoss(
delta_var, delta_dist, norm, self.usegpu)
if self.usegpu:
self.criterion_discriminative = \
self.criterion_discriminative.cuda()
# FG Segmentation Loss
if class_weights is not None:
class_weights = self.__define_variable(
torch.FloatTensor(class_weights))
if criterion in ['CE', 'Multi']:
self.criterion_ce = torch.nn.CrossEntropyLoss(class_weights)
if criterion in ['Dice', 'Multi']:
self.criterion_dice = DiceLoss(optimize_bg=optimize_bg,
weight=class_weights,
smooth=smooth)
else:
if criterion in ['CE', 'Multi']:
self.criterion_ce = torch.nn.CrossEntropyLoss()
if criterion in ['Dice', 'Multi']:
self.criterion_dice = DiceLoss(optimize_bg=optimize_bg,
smooth=smooth)
# MSE Loss
self.criterion_mse = torch.nn.MSELoss()
if self.usegpu:
if criterion in ['CE', 'Multi']:
self.criterion_ce = self.criterion_ce.cuda()
if criterion in ['Dice', 'Multi']:
self.criterion_dice = self.criterion_dice.cuda()
self.criterion_mse = self.criterion_mse.cuda()
def __init__(self,
model,
lr,
num_classes,
weight_ce,
weight_dice,
metrics=True):
super().__init__()
# model
self.model = model
# learning rate
self.lr = lr
# number of classes
self.num_classes = num_classes
# loss
self.register_buffer('weight_ce', weight_ce)
self.register_buffer('weight_dice', weight_dice)
self.dice_loss = DiceLoss(weight=self.weight_dice)
self.ce_loss = CrossEntropyLoss(weight=self.weight_ce)
# save hyperparameters
self.save_hyperparameters()
# metrics
self.metrics = metrics
if self.metrics:
self.f1_train = CustomMetric(metric=pl.metrics.functional.f1,
metric_name='F1',
num_classes=self.num_classes,
average='none')
self.f1_valid = CustomMetric(metric=pl.metrics.functional.f1,
metric_name='F1',
num_classes=self.num_classes,
average='none')
self.f1_test = CustomMetric(metric=pl.metrics.functional.f1,
metric_name='F1',
num_classes=self.num_classes,
average='none')
self.iou_train = CustomMetric(metric=pl.metrics.functional.iou,
metric_name='IoU',
num_classes=self.num_classes,
reduction='none')
self.iou_valid = CustomMetric(metric=pl.metrics.functional.iou,
metric_name='IoU',
num_classes=self.num_classes,
reduction='none')
self.iou_test = CustomMetric(metric=pl.metrics.functional.iou,
metric_name='IoU',
num_classes=self.num_classes,
reduction='none')
def select_loss(loss_function):
if loss_function == 'bce':
criterion = nn.BCELoss()
elif loss_function == 'bce_logit':
criterion = nn.BCEWithLogitsLoss()
elif loss_function == 'dice':
criterion = DiceLoss()
elif loss_function == 'mse':
criterion = nn.MSELoss()
elif loss_function == 'l1':
criterion = nn.L1Loss()
elif loss_function == 'kl' or loss_function == 'jsd':
criterion = nn.KLDivLoss()
elif loss_function == 'Cldice':
bce = nn.BCEWithLogitsLoss().cuda()
dice = DiceLoss().cuda()
criterion = ClDice(bce, dice, alpha=1, beta=1)
else:
raise ValueError('Not supported loss.')
return criterion.cuda()
def get_loss(mask_size):
if loss == 'bce-with-logits':
return WeightedBCEWithLogitsLoss(mask_size)
elif loss == 'dice':
return DiceLoss()
elif loss == 'cross-entropy':
l = torch.nn.CrossEntropyLoss()
l.__name__ = 'Cross Entropy Loss'
return l
else:
NotImplementedError("Unknown loss: {}".format(loss))
def loss_function_select(loss_function):
if loss_function == 'bce':
criterion = nn.BCEWithLogitsLoss(
pos_weight=torch.Tensor([args.bce_weight])).cuda()
elif loss_function == 'dice':
criterion = DiceLoss().cuda()
elif loss_function == 'mse':
criterion = nn.MSELoss().cuda()
elif loss_function == 'l1':
criterion = nn.L1Loss().cuda()
elif loss_function == 'kl' or loss_function == 'jsd':
criterion = nn.KLDivLoss().cuda()
elif loss_function == 'Cldice':
bce = nn.BCEWithLogitsLoss(
pos_weight=torch.Tensor([args.bce_weight])).cuda()
dice = DiceLoss().cuda()
criterion = ClDice(bce, dice, alpha=1, beta=1)
else:
raise ValueError('Not supported loss.')
return criterion
def case_test(self, model, device, test_loader, case_id):
self.model.eval()
testNum = len(test_loader.dataset)
diceloss = 0.
Criterion = DiceLoss()
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data).view_as(target)
diceloss += Criterion(output, target)
diceloss /= testNum
dice = 1 - diceloss
return dice.item()
def train(train_loader, student, teacher, optimizer, epoch):
global global_step
global best_dice
# set criterion
segmentation_criterion = DiceLoss(ignore=True)
consistency_criterion = nn.MSELoss()
#switch to train mode
student.train()
teacher.train()
for i, (data, label) in enumerate(train_loader):
# if we don't use cosine_annealing, adjust learning rate
if args.cosine_annealing == False:
adjust_learning_rate(optimizer, epoch, i, len(train_loader))
data = data.to(device)
label = label.to(device)
# get the result of the two models
student_pred = student(data)
teacher_pred = teacher(data)
# We don't want gradient descent in teacher model.
teacher_pred = torch.autograd.Variable(teacher_pred.detach().data,
requires_grad=False)
# calculate consistency criterion
consistency_weight = get_current_consistency_weight(epoch)
consistency_loss = consistency_weight * consistency_criterion(
student_pred, teacher_pred)
# calculate segmentation loss
segmentation_loss = segmentation_criterion(student_pred, label)
# combine them to get the final loss
loss = consistency_loss + segmentation_loss
# compute gradient and do optimization
optimizer.zero_grad()
loss.backward()
optimizer.step()
global_step += 1
update_ema_variables(student, teacher, args.ema_decay, global_step)
# print training information
if i % args.print_freq == 0:
LOG.info('Epoch: [{}][{}/{}]\tLoss: {:.3f}\t'.format(
epoch, i, len(train_loader), loss))
LOG.info('Conssistency Loss: {:.3f}\t'.format(consistency_loss))
def __init__(self, model, args, train_dataset, eval_dataset,
compute_metrics, loss_type, loss_gamma):
Trainer.__init__(self,
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
compute_metrics=compute_metrics)
if loss_type == 'DiceLoss':
self.loss_fct = DiceLoss()
elif loss_type == 'FocalLoss':
self.loss_fct = FocalLoss(gamma=loss_gamma)
elif loss_type == 'LabelSmoothingCrossEntropy':
self.loss_fct = LabelSmoothingCrossEntropy()
elif loss_type == 'CrossEntropyLoss':
self.loss_fct = CrossEntropyLoss()
elif loss_type == 'CourageLoss':
self.loss_fct = CourageLoss(gamma=loss_gamma)
else:
raise ValueError("Doesn't support such loss type")
def main():
args = arg_parse()
mode = args.mode
lossf = args.lossf.lower()
n_epochs = args.epochs
opt = args.opt.lower()
lr = args.lr
hash_code = '_'.join(
list(map(str, [args.hashcode, args.opt, args.lr, args.lossf])))
device = 'cuda:' + args.gpu
withlen = args.withlen.lower() == 'true'
normed = args.normed.lower() == 'true'
mu = args.mu
alpha = args.alpha
beta = args.beta
modelpath = args.modelpath
DEFAULT_BATCHSIZE = args.batchsize
modeltype = args.model
if mode == 'test' and not modelpath:
raise ValueError("Need model path for testing!!")
if mode == 'train':
DRIVE_train = DRIVEDataset("train",
DRIVE_train_imgs_original,
DRIVE_train_groudTruth,
DRIVE_train_narrowBand,
patch_height,
patch_width,
N_subimgs,
val_size=0)
print(len(DRIVE_train))
# DRIVE_valid = DRIVE_train.get_validation_dataset()
DRIVE_train_load = \
torch.utils.data.DataLoader(dataset=DRIVE_train,
batch_size=DEFAULT_BATCHSIZE, shuffle=True)
# DRIVE_val_load = \
# torch.utils.data.DataLoader(dataset=DRIVE_valid,
# batch_size=128, shuffle=True)
DRIVE_valid = DRIVEDataset("test", DRIVE_test_imgs_original,
DRIVE_test_groundTruth,
DRIVE_test_narrowBand, patch_height,
patch_width, None)
DRIVE_val_load = \
torch.utils.data.DataLoader(dataset=DRIVE_valid,
batch_size=DEFAULT_BATCHSIZE, shuffle=False)
else:
DRIVE_test = DRIVEDataset("test", DRIVE_test_imgs_original,
DRIVE_test_groundTruth,
DRIVE_test_narrowBand, patch_height,
patch_width, None)
DRIVE_test_load = \
torch.utils.data.DataLoader(dataset=DRIVE_test,
batch_size=DEFAULT_BATCHSIZE, shuffle=False)
shape = (1, 48, 48)
if torch.cuda.is_available():
if modeltype.lower() == 'fcn':
model = FCN8s(n_class=1).to(device)
else:
model = UNet1024(shape).to(device)
else:
if modeltype.lower() == 'fcn':
model = FCN8s(n_class=1)
else:
model = UNet1024(shape)
if lossf == 'bce':
criterion = BinaryCrossEntropyLoss2d()
elif lossf == 'dice':
criterion = DiceLoss()
elif lossf == 'contour':
criterion = ContourLoss(device=device,
mu=mu,
alpha=alpha,
beta=beta,
normed=normed,
withlen=withlen)
elif lossf == 'contour-v3':
criterion = ContourLossV3(device=device,
mu=mu,
alpha=alpha,
beta=beta,
normed=normed,
withlen=withlen)
elif lossf == 'contour-v2':
criterion = ContourLossV2(device=device,
mu=mu,
alpha=alpha,
beta=beta,
normed=normed,
withlen=withlen)
elif lossf == 'contour-v4':
criterion = ContourLossV4(device=device,
mu=mu,
normed=normed,
withlen=withlen)
elif lossf == 'focal':
criterion = FocalLoss()
else:
raise ValueError('Undefined loss type')
optimizer = None
if opt == 'rmsprop':
optimizer = torch.optim.RMSprop(model.parameters(), lr=lr)
if opt == 'adam':
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
# Saving History to csv
header = [
'epoch', 'train loss', 'train auc', 'train accuracy', 'val loss',
'val auc', 'val accuracy', 'val jaccard', 'val sensitivity',
'val specitivity', 'val precision', 'val f1', 'val pr auc'
]
save_file_name = f"../history/{hash_code}/history.csv"
save_dir = f"../history/{hash_code}"
# Saving images and models directories
model_save_dir = f"../history/{hash_code}/saved_models"
image_save_path = f"../history/{hash_code}/result_images"
# Train
if mode == 'train':
print("Initializing Training!")
min_loss = 1e9
best_epoch = 0
early_stop_count = 0
max_count = 5
for i in range(n_epochs):
train_loss, train_acc = train_model(i, model, DRIVE_train_load,
criterion, optimizer, device)
print('Epoch', str(i + 1), 'Train loss:', train_loss, 'Train acc:',
train_acc)
if (i + 1) % 5 == 0:
# val_loss, val_acc = validate_model(model, DRIVE_val_load, criterion, device)
# print('Epoch', str(i+1),
# 'Val loss:', val_loss,
# 'Val acc:', val_acc,
# )
val_loss, val_auc, val_accuracy, val_jaccard, val_sensitivity,\
val_specitivity, val_precision, val_f1, val_pr_auc, val_iou = test_model(model, DRIVE_val_load, criterion,
test_border_masks,
f'{image_save_path}/{i+1}/', device)
print(
'Epoch',
str(i + 1),
'Val loss:',
val_loss,
'Val acc:',
val_accuracy,
)
values = [
i + 1, train_loss, train_acc, val_loss, val_auc,
val_accuracy, val_jaccard, val_sensitivity,
val_specitivity, val_precision, val_f1, val_pr_auc, val_iou
]
export_history(header, values, save_dir, save_file_name)
if val_loss < min_loss:
early_stop_count = 0
min_loss = val_loss
best_epoch = i
save_model(model, model_save_dir, i + 1)
else:
early_stop_count += 1
if early_stop_count > max_count:
print(
'Traning can not improve from epoch {}\tBest loss: {}'
.format(best_epoch, min_loss))
break
else:
print("Initializing Testing!")
model.load_state_dict(torch.load(modelpath))
model.eval()
test_auc, test_accuracy, test_jaccard, test_sensitivity, test_specitivity, \
test_precision, test_f1, test_pr_auc \
= test_model_img(model,
DRIVE_test_load,
test_border_masks,
f'{image_save_path}/',
device)
print('Test auc: ', test_auc)
print('Test accuracy: ', test_accuracy)
print('Test jaccard: ', test_jaccard)
print('Test sensitivity: ', test_sensitivity)
print('Test specitivity: ', test_specitivity)
print('Test precision: ', test_precision)
print('Test f1: ', test_f1)
print('Test PR-AUC: ', test_pr_auc)
print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
print('AUC mean: {} - std: {}'.format(*mean_std(test_auc)))
print('ACCURACY mean: {} - std: {}'.format(*mean_std(test_accuracy)))
print('JACCARD mean: {} - std: {}'.format(*mean_std(test_jaccard)))
print('SENS mean: {} - std: {}'.format(*mean_std(test_sensitivity)))
print('SPEC mean: {} - std: {}'.format(*mean_std(test_specitivity)))
print('PRECISION mean: {} - std: {}'.format(*mean_std(test_precision)))
print('F1 mean: {} - std: {}'.format(*mean_std(test_f1)))
print('PR AUC: {} - std: {}'.format(*mean_std(test_pr_auc)))
print("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
def train_net(net, options):
data_path = options.data_path + '240dataset/'
csv_file = options.data_path + 'new_train.csv'
origin_spacing_data_path = options.data_path + 'origin_spacing_croped/'
# z_size is the random crop size along z-axis, you can set it larger if have enough gpu memory
trainset = BrainDataset(csv_file,
data_path,
data_path,
mode='train',
z_size=40)
trainLoader = data.DataLoader(trainset,
batch_size=options.batch_size,
shuffle=True,
num_workers=0)
test_data_list, test_label_list = load_test_data(origin_spacing_data_path)
writer = SummaryWriter(options.log_path + options.unique_name)
optimizer = optim.SGD(net.parameters(),
lr=options.lr,
momentum=0.9,
weight_decay=0.0005)
org_weight = torch.FloatTensor(options.org_weight).unsqueeze(1).cuda()
criterion_fl = FocalLoss(10, alpha=org_weight)
criterion_dl = DiceLoss()
best_dice = 0
for epoch in range(options.epochs):
print('Starting epoch {}/{}'.format(epoch + 1, options.epochs))
epoch_loss = 0
multistep_scheduler = multistep_lr_scheduler_with_warmup(
optimizer,
init_lr=options.lr,
epoch=epoch,
warmup_epoch=5,
lr_decay_epoch=[200, 400],
max_epoch=options.epochs,
gamma=0.1)
print('current lr:', multistep_scheduler)
net.train()
for i, (img, label, weight) in enumerate(trainLoader, 0):
img = img.cuda()
label = label.cuda()
weight = weight.cuda()
end = time.time()
optimizer.zero_grad()
result = net(img)
if options.rlt > 0:
loss = criterion_fl(result, label,
weight) + options.rlt * criterion_dl(
result, label, weight)
else:
loss = criterion_dl(result, label, weight)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
batch_time = time.time() - end
print('batch loss: %.5f, batch_time:%.5f' %
(loss.item(), batch_time))
print('[epoch %d] epoch loss: %.5f' % (epoch + 1, epoch_loss /
(i + 1)))
writer.add_scalar('Train/Loss', epoch_loss / (i + 1), epoch + 1)
writer.add_scalar('LR', multistep_scheduler, epoch + 1)
if os.path.isdir('%s%s/' % (options.cp_path, options.unique_name)):
pass
else:
os.mkdir('%s%s/' % (options.cp_path, options.unique_name))
if (epoch + 1) % 10 == 0:
torch.save(
net.state_dict(), '%s%s/CP%d.pth' %
(options.cp_path, options.unique_name, epoch))
avg_dice, dice_list = validation(net, test_data_list, test_label_list)
writer.add_scalar('Test/AVG_Dice', avg_dice, epoch + 1)
for idx in range(9):
writer.add_scalar('Test/Dice%d' % (idx + 1), dice_list[idx],
epoch + 1)
if avg_dice >= best_dice:
best_dice = avg_dice
torch.save(
net.state_dict(),
'%s%s/best.pth' % (options.cp_path, options.unique_name))
print('save done')
print('dice: %.5f/best dice: %.5f' % (avg_dice, best_dice))
valid_idx = fold[1]
train_set = Dataset([folder_paths[i] for i in train_idx],
[folder_ids[i] for i in train_idx])
valid_set = Dataset([folder_paths[i] for i in valid_idx],
[folder_ids[i] for i in valid_idx])
train_loader = data.DataLoader(train_set, **params)
valid_loader = data.DataLoader(valid_set, **params)
# Model
model = Modified3DUNet(in_channels, n_classes, base_n_filter)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
# Loss and optimizer
#criterion = torch.nn.CrossEntropyLoss().to(device) # For cross entropy
criterion = DiceLoss() # For dice loss
optimizer = torch.optim.Adam(model.parameters())
# Load model and optimizer parameters if the training was interrupted and must be continued - need to also change epoch range in for loop
# checkpoint = torch.load("/content/drive/My Drive/Brats2019/Model_Saves_KFold/Fold_1_Epoch_30_Train_Loss_0.0140_Valid_Loss_0.0137.tar")
# model.load_state_dict(checkpoint['model_state_dict'])
model.to(device)
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
#model.train()
for epoch in range(1, max_epochs + 1):
start_time = time.time()
train_losses = []
for batch, labels in train_loader:
# Data Augment
#augmenter = DataAugment(batch,labels)
#batch,labels = augmenter.augment()
'lr': 5e-3,
'weight_decay': 0.00003
},
{
'params': model.encoder.parameters(),
'lr': 5e-3 / 10,
'weight_decay': 0.00003
},
])
optimizer = contrib.nn.Lookahead(base_optimizer)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.2,
patience=5)
diceloss = DiceLoss()
loss_list = []
dice_list = []
for epoch in tqdm(range(EPOCHES)):
###Train
train_loss = 0
for data in d_train:
optimizer.zero_grad()
img, mask = data
img = img.to(DEVICE)
mask = mask.to(DEVICE)
outputs = model(img)
dice = dice_coef(outputs, mask)
loss = diceloss(outputs, mask)
loss.backward()
def main():
work_dir = os.path.join(args.work_dir, args.exp)
if not os.path.exists(work_dir):
os.makedirs(work_dir)
# copy this file to work dir to keep training configuration
shutil.copy(__file__, os.path.join(work_dir, 'main.py'))
with open(os.path.join(work_dir, 'args.pkl'), 'wb') as f:
pickle.dump(args, f)
#train
trn_image_root = os.path.join(args.trn_root, 'images')
exam_ids = os.listdir(trn_image_root)
random.shuffle(exam_ids)
train_exam_ids = exam_ids
#train_exam_ids = exam_ids[:int(len(exam_ids)*0.8)]
#val_exam_ids = exam_ids[int(len(exam_ids) * 0.8):]
# train_dataset
trn_dataset = DatasetTrain(args.trn_root,
train_exam_ids,
options=args,
input_stats=[0.5, 0.5])
trn_loader = torch.utils.data.DataLoader(trn_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
# save input stats for later use
np.save(os.path.join(work_dir, 'input_stats.npy'), trn_dataset.input_stats)
#val
val_image_root = os.path.join(args.val_root, 'images')
val_exam = os.listdir(val_image_root)
random.shuffle(val_exam)
val_exam_ids = val_exam
# val_dataset
val_dataset = DatasetVal(args.val_root,
val_exam_ids,
options=args,
input_stats=trn_dataset.input_stats)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
# make logger
trn_logger = Logger(os.path.join(work_dir, 'train.log'))
trn_raw_logger = Logger(os.path.join(work_dir, 'train_raw.log'))
val_logger = Logger(os.path.join(work_dir, 'validation.log'))
# model_select
if args.model == 'unet':
net = UNet3D(1,
1,
f_maps=args.f_maps,
depth_stride=args.depth_stride,
conv_layer_order=args.conv_layer_order,
num_groups=args.num_groups)
else:
raise ValueError('Not supported network.')
# loss_select
if args.loss_function == 'bce':
criterion = nn.BCEWithLogitsLoss(
pos_weight=torch.Tensor([args.bce_weight])).cuda()
elif args.loss_function == 'dice':
criterion = DiceLoss().cuda()
elif args.loss_function == 'weight_bce':
criterion = nn.BCEWithLogitsLoss(
pos_weight=torch.FloatTensor([5])).cuda()
else:
raise ValueError('{} loss is not supported yet.'.format(
args.loss_function))
# optim_select
if args.optim == 'sgd':
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=False)
elif args.optim == 'adam':
optimizer = optim.Adam(net.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
else:
raise ValueError('{} optim is not supported yet.'.format(args.optim))
net = nn.DataParallel(net).cuda()
cudnn.benchmark = True
lr_schedule = args.lr_schedule
lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=lr_schedule[:-1],
gamma=0.1)
best_iou = 0
for epoch in range(lr_schedule[-1]):
train(trn_loader, net, criterion, optimizer, epoch, trn_logger,
trn_raw_logger)
iou = validate(val_loader, net, criterion, epoch, val_logger)
lr_scheduler.step()
# save model parameter
is_best = iou > best_iou
best_iou = max(iou, best_iou)
checkpoint_filename = 'model_checkpoint_{:0>3}.pth'.format(epoch + 1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict()
}, is_best, work_dir, checkpoint_filename)
# visualize curve
draw_curve(work_dir, trn_logger, val_logger)
if args.inplace_test:
# calc overall performance and save figures
print('Test mode ...')
main_test(model=net, args=args)
def main():
print(args.work_dir, args.exp)
work_dir = os.path.join(args.work_dir, args.exp)
if not os.path.exists(work_dir):
os.makedirs(work_dir)
# copy this file to work dir to keep training configuration
shutil.copy(__file__, os.path.join(work_dir, 'main.py'))
with open(os.path.join(work_dir, 'args.pkl'), 'wb') as f:
pickle.dump(args, f)
# 1.dataset
train_filename = args.trn_root
test_filename = args.test_root
trainset = Segmentation_2d_data(train_filename)
valiset = Segmentation_2d_data(test_filename)
train_loader = data.DataLoader(trainset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True)
valid_loader = data.DataLoader(valiset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True)
trn_logger = Logger(os.path.join(work_dir, 'train.log'))
trn_raw_logger = Logger(os.path.join(work_dir, 'train_raw.log'))
val_logger = Logger(os.path.join(work_dir, 'validation.log'))
if args.model == 'unet':
net = Unet2D(in_shape=(1, 512, 512), padding=args.padding_size, momentum=args.batchnorm_momentum)
elif args.model == 'unetcoord':
net = Unet2D_coordconv(in_shape=(1, 512, 512), padding=args.padding_size,
momentum=args.batchnorm_momentum, coordnumber=args.coordconv_no, radius=False)
elif args.model == 'unetmultiinput':
net = Unet2D_multiinput(in_shape=(1, 512, 512), padding=args.padding_size,
momentum=args.batchnorm_momentum)
elif args.model == 'scse_block':
net = Unet_sae(in_shape=(1, 512, 512), padding=args.padding_size, momentum=args.batchnorm_momentum)
else:
raise ValueError('Not supported network.')
# loss
if args.loss_function == 'bce':
criterion = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([args.bce_weight])).cuda()
elif args.loss_function == 'dice':
criterion = DiceLoss().cuda()
else:
raise ValueError('{} loss is not supported yet.'.format(args.loss_function))
# optim
if args.optim_function == 'sgd':
optimizer = torch.optim.SGD(net.parameters(), lr=args.initial_lr, momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optim_function == 'adam':
optimizer = torch.optim.Adam(net.parameters(), lr=args.initial_lr, weight_decay=args.weight_decay)
elif args.optim_function == 'radam':
optimizer = RAdam(net.parameters(), lr=args.initial_lr, weight_decay = args.weight_decay)
else:
raise ValueError('{} loss is not supported yet.'.format(args.optim_function))
net = nn.DataParallel(net).cuda()
cudnn.benchmark = True
lr_schedule = args.lr_schedule
lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=lr_schedule[:-1],
gamma=0.1)
best_iou = 0
for epoch in range(lr_schedule[-1]):
train(train_loader, net, criterion, optimizer, epoch, trn_logger, trn_raw_logger)
iou = validate(valid_loader, net, criterion, epoch, val_logger)
lr_scheduler.step()
is_best = iou > best_iou
best_iou = max(iou, best_iou)
checkpoint_filename = 'model_checkpoint_{:0>3}.pth'.format(epoch + 1)
save_checkpoint({'epoch': epoch + 1,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict()},
is_best,
work_dir,
checkpoint_filename)
draw_curve(work_dir, trn_logger, val_logger)
def get_lossfn():
bce = nn.BCEWithLogitsLoss()
dice = DiceLoss(mode='binary', log_loss=True, smooth=1e-7)
focal = losses.BinaryFocalLoss(alpha=0.25, reduced_threshold=0.5)
criterion = ImanipLoss(bce, seglossA=dice, seglossB=focal)
return criterion