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 __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 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 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 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 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 get_loss(self, label, pred, task): if self.config.loss_function == 'cross_entropy': loss = nn.CrossEntropyLoss(weight=self.label_weights[task] if self. label_weights is not None else None) elif self.config.loss_function == 'focal_loss': if self.label_weights is None: raise RuntimeError("标签权重为空") loss = FocalLoss.MultiFocalLoss(len(self.label_weights[task]), self.label_weights[task]) elif self.config.loss_function == 'dice_loss': loss = DiceLoss.MultiDiceLoss() else: raise RuntimeError("没有实现的损失函数") return self.uw.get_loss(loss( pred, label), task) if self.uw is not None else loss(pred, label)
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 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
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()
class Model(object): def __init__(self, dataset, model_name, n_classes, max_n_objects, use_instance_segmentation=False, use_coords=False, load_model_path='', usegpu=True): self.dataset = dataset self.model_name = model_name self.n_classes = n_classes self.max_n_objects = max_n_objects self.use_instance_segmentation = use_instance_segmentation self.use_coords = use_coords self.load_model_path = load_model_path self.usegpu = usegpu assert self.dataset in [ 'CVPPP', ] assert self.model_name in ['ReSeg', 'StackedRecurrentHourglass'] if self.dataset == 'CVPPP': if self.model_name == 'ReSeg': self.model = ReSeg(self.n_classes, self.use_instance_segmentation, pretrained=True, use_coordinates=self.use_coords, usegpu=self.usegpu) elif self.model_name == 'StackedRecurrentHourglass': self.model = SRecHg(self.n_classes, self.use_instance_segmentation, self.use_coords, pretrained=True, usegpu=self.usegpu) self.__load_weights() if self.usegpu: cudnn.benchmark = True self.model.cuda() # self.model = torch.nn.DataParallel(self.model, # device_ids=range(self.ngpus)) print self.model self.vis = visdom.Visdom() self.training_metric_vis, self.test_metric_vis = None, None if self.use_instance_segmentation: self.instance_seg_vis = None def __load_weights(self): if self.load_model_path != '': assert os.path.isfile(self.load_model_path), 'Model : {} does not \ exists!'.format(self.load_model_path) print 'Loading model from {}'.format(self.load_model_path) model_state_dict = self.model.state_dict() if self.usegpu: pretrained_state_dict = torch.load(self.load_model_path) else: pretrained_state_dict = torch.load( self.load_model_path, map_location=lambda storage, loc: storage) model_state_dict.update(pretrained_state_dict) self.model.load_state_dict(model_state_dict) def __define_variable(self, tensor, volatile=False): if volatile: with torch.no_grad(): return Variable(tensor) return Variable(tensor) def __define_input_variables(self, features, fg_labels, ins_labels, n_objects, mode): volatile = True if mode == 'training': volatile = False features_var = self.__define_variable(features, volatile=volatile) fg_labels_var = self.__define_variable(fg_labels, volatile=volatile) ins_labels_var = self.__define_variable(ins_labels, volatile=volatile) n_objects_var = self.__define_variable(n_objects, volatile=volatile) return features_var, fg_labels_var, ins_labels_var, n_objects_var 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 __define_optimizer(self, learning_rate, weight_decay, lr_drop_factor, lr_drop_patience, optimizer='Adam'): assert optimizer in ['RMSprop', 'Adam', 'Adadelta', 'SGD'] parameters = ifilter(lambda p: p.requires_grad, self.model.parameters()) if optimizer == 'RMSprop': self.optimizer = optim.RMSprop(parameters, lr=learning_rate, weight_decay=weight_decay) elif optimizer == 'Adadelta': self.optimizer = optim.Adadelta(parameters, lr=learning_rate, weight_decay=weight_decay) elif optimizer == 'Adam': self.optimizer = optim.Adam(parameters, lr=learning_rate, weight_decay=weight_decay) elif optimizer == 'SGD': self.optimizer = optim.SGD(parameters, lr=learning_rate, momentum=0.9, weight_decay=weight_decay) self.lr_scheduler = ReduceLROnPlateau(self.optimizer, mode='min', factor=lr_drop_factor, patience=lr_drop_patience, verbose=True) @staticmethod def __get_loss_averager(): return averager() def __minibatch(self, train_test_iter, clip_grad_norm, criterion_type, train_cnn=True, mode='training', debug=False): assert mode in ['training', 'test'], 'Mode must be either "training" or "test"' if mode == 'training': for param in self.model.parameters(): param.requires_grad = True if not train_cnn: for param in self.model.cnn.parameters(): param.requires_grad = False self.model.train() else: for param in self.model.parameters(): param.requires_grad = False self.model.eval() cpu_images, cpu_sem_seg_annotations, \ cpu_ins_seg_annotations, cpu_n_objects = train_test_iter.next() cpu_images = cpu_images.contiguous() cpu_sem_seg_annotations = cpu_sem_seg_annotations.contiguous() cpu_ins_seg_annotations = cpu_ins_seg_annotations.contiguous() cpu_n_objects = cpu_n_objects.contiguous() if self.usegpu: gpu_images = cpu_images.cuda(async=True) gpu_sem_seg_annotations = cpu_sem_seg_annotations.cuda(async=True) gpu_ins_seg_annotations = cpu_ins_seg_annotations.cuda(async=True) gpu_n_objects = cpu_n_objects.cuda(async=True) else: gpu_images = cpu_images gpu_sem_seg_annotations = cpu_sem_seg_annotations gpu_ins_seg_annotations = cpu_ins_seg_annotations gpu_n_objects = cpu_n_objects gpu_images, gpu_sem_seg_annotations, \ gpu_ins_seg_annotations, gpu_n_objects = \ self.__define_input_variables(gpu_images, gpu_sem_seg_annotations, gpu_ins_seg_annotations, gpu_n_objects, mode) gpu_n_objects = gpu_n_objects.unsqueeze(dim=1) gpu_n_objects_normalized = gpu_n_objects.float() / self.max_n_objects sem_seg_predictions, ins_seg_predictions, \ n_objects_predictions = self.model(gpu_images) if mode == 'test': if debug: _vis_prob = np.random.rand() if _vis_prob > 0.7: if self.use_instance_segmentation: sem_seg_preds = np.argmax( sem_seg_predictions.data.cpu().numpy(), axis=1) seg_preds = ins_seg_predictions.data.cpu().numpy() _bs, _n_feats = seg_preds.shape[:2] _sample_idx = np.random.randint(_bs) _sem_seg_preds_sample = sem_seg_preds[_sample_idx] _seg_preds_sample = seg_preds[_sample_idx] fg_ins_embeddings = np.stack([ _seg_preds_sample[i][np.where( _sem_seg_preds_sample == 1)] for i in range(_n_feats) ], axis=1) _n_fg_samples = fg_ins_embeddings.shape[0] if _n_fg_samples > 0: fg_ins_embeddings = \ fg_ins_embeddings[np.random.choice( range(_n_fg_samples), size=400)] tsne = TSNE(n_components=2, random_state=0) fg_ins_embeddings_vis = tsne.fit_transform( fg_ins_embeddings) if self.instance_seg_vis: self.vis.scatter(X=fg_ins_embeddings_vis, win=self.instance_seg_vis, opts={ 'title': 'Predicted Embeddings \ for Foreground \ Predictions', 'markersize': 2 }) else: self.instance_seg_vis =\ self.vis.scatter(X=fg_ins_embeddings_vis, opts={'title': 'Predicted \ Embeddings for \ Foreground \ Predictions' , 'markersize': 2}) cost = 0.0 out_metrics = dict() if self.use_instance_segmentation: disc_cost = self.criterion_discriminative( ins_seg_predictions, gpu_ins_seg_annotations.float(), cpu_n_objects, self.max_n_objects) cost += disc_cost out_metrics['Discriminative Cost'] = disc_cost.data if criterion_type in ['CE', 'Multi']: _, gpu_sem_seg_annotations_criterion_ce = \ gpu_sem_seg_annotations.max(1) ce_cost = self.criterion_ce( sem_seg_predictions.permute(0, 2, 3, 1).contiguous().view( -1, self.n_classes), gpu_sem_seg_annotations_criterion_ce.view(-1)) cost += ce_cost out_metrics['CE Cost'] = ce_cost.data if criterion_type in ['Dice', 'Multi']: dice_cost = self.criterion_dice(sem_seg_predictions, gpu_sem_seg_annotations) cost += dice_cost out_metrics['Dice Cost'] = dice_cost.data mse_cost = self.criterion_mse(n_objects_predictions, gpu_n_objects_normalized) cost += mse_cost out_metrics['MSE Cost'] = mse_cost.data if mode == 'training': self.model.zero_grad() cost.backward() if clip_grad_norm != 0: torch.nn.utils.clip_grad_norm_(self.model.parameters(), clip_grad_norm) self.optimizer.step() return out_metrics def __test(self, test_loader, criterion_type, epoch, debug): n_minibatches = len(test_loader) test_iter = iter(test_loader) out_metrics = dict() for minibatch_index in range(n_minibatches): mb_out_metrics = self.__minibatch(test_iter, 0.0, criterion_type, train_cnn=False, mode='test', debug=debug) for mk, mv in mb_out_metrics.iteritems(): if mk not in out_metrics: out_metrics[mk] = [] out_metrics[mk].append(mv) test_metric_vis_data, test_metric_vis_legend = [], [] metrics_as_str = 'Testing: [METRIC]' for mk, mv in out_metrics.iteritems(): out_metrics[mk] = torch.stack(mv, dim=0).mean() metrics_as_str += ' {} : {} |'.format(mk, out_metrics[mk]) test_metric_vis_data.append(out_metrics[mk]) test_metric_vis_legend.append(mk) print metrics_as_str test_metric_vis_data = np.expand_dims(np.array(test_metric_vis_data), 0) if self.test_metric_vis: self.vis.line(X=np.array([epoch]), Y=test_metric_vis_data, win=self.test_metric_vis, update='append') else: self.test_metric_vis = self.vis.line(X=np.array([epoch]), Y=test_metric_vis_data, opts={ 'legend': test_metric_vis_legend, 'title': 'Test Metrics', 'showlegend': True, 'xlabel': 'Epoch', 'ylabel': 'Metric' }) return out_metrics def fit(self, criterion_type, delta_var, delta_dist, norm, learning_rate, weight_decay, clip_grad_norm, lr_drop_factor, lr_drop_patience, optimize_bg, optimizer, train_cnn, n_epochs, class_weights, train_loader, test_loader, model_save_path, debug): assert criterion_type in ['CE', 'Dice', 'Multi'] training_log_file = open(os.path.join(model_save_path, 'training.log'), 'w') validation_log_file = open( os.path.join(model_save_path, 'validation.log'), 'w') training_log_file.write('Epoch,Cost\n') validation_log_file.write('Epoch,Cost\n') self.__define_criterion(class_weights, delta_var, delta_dist, norm=norm, optimize_bg=optimize_bg, criterion=criterion_type) self.__define_optimizer(learning_rate, weight_decay, lr_drop_factor, lr_drop_patience, optimizer=optimizer) self.__test(test_loader, criterion_type, -1.0, debug) best_val_cost = np.Inf for epoch in range(n_epochs): epoch_start = time.time() train_iter = iter(train_loader) n_minibatches = len(train_loader) train_out_metrics = dict() minibatch_index = 0 while minibatch_index < n_minibatches: mb_out_metrics = self.__minibatch(train_iter, clip_grad_norm, criterion_type, train_cnn=train_cnn, mode='training', debug=debug) for mk, mv in mb_out_metrics.iteritems(): if mk not in train_out_metrics: train_out_metrics[mk] = [] train_out_metrics[mk].append(mv) minibatch_index += 1 epoch_end = time.time() epoch_duration = epoch_end - epoch_start training_metric_vis_data, training_metric_vis_legend = [], [] print 'Epoch : [{}/{}] - [{}]'.format(epoch, n_epochs, epoch_duration) metrics_as_str = 'Training: [METRIC]' for mk, mv in train_out_metrics.iteritems(): train_out_metrics[mk] = torch.stack(mv, dim=0).mean() metrics_as_str += ' {} : {} |'.format(mk, train_out_metrics[mk]) training_metric_vis_data.append(train_out_metrics[mk]) training_metric_vis_legend.append(mk) print metrics_as_str training_metric_vis_data = np.expand_dims( np.array(training_metric_vis_data), 0) if self.training_metric_vis: self.vis.line(X=np.array([epoch]), Y=training_metric_vis_data, win=self.training_metric_vis, update='append') else: self.training_metric_vis = self.vis.line( X=np.array([epoch]), Y=training_metric_vis_data, opts={ 'legend': training_metric_vis_legend, 'title': 'Training Metrics', 'showlegend': True, 'xlabel': 'Epoch', 'ylabel': 'Metric' }) val_out_metrics = self.__test(test_loader, criterion_type, epoch, debug) if self.use_instance_segmentation: val_cost = val_out_metrics['Discriminative Cost'] train_cost = train_out_metrics['Discriminative Cost'] elif criterion_type in ['Dice', 'Multi']: val_cost = val_out_metrics['Dice Cost'] train_cost = train_out_metrics['Dice Cost'] else: val_cost = val_out_metrics['CE Cost'] train_cost = train_out_metrics['CE Cost'] self.lr_scheduler.step(val_cost) is_best_model = val_cost <= best_val_cost if is_best_model: best_val_cost = val_cost torch.save( self.model.state_dict(), os.path.join(model_save_path, 'model_{}_{}.pth'.format(epoch, val_cost))) training_log_file.write('{},{}\n'.format(epoch, train_cost)) validation_log_file.write('{},{}\n'.format(epoch, val_cost)) training_log_file.flush() validation_log_file.flush() training_log_file.close() validation_log_file.close() def predict(self, images): assert len(images.size()) == 4 # b, c, h, w for param in self.model.parameters(): param.requires_grad = False self.model.eval() images = images.contiguous() if self.usegpu: images = images.cuda(async=True) images = self.__define_variable(images, volatile=True) sem_seg_predictions, ins_seg_predictions, n_objects_predictions = \ self.model(images) sem_seg_predictions = torch.nn.functional.softmax(sem_seg_predictions, dim=1) n_objects_predictions = n_objects_predictions * self.max_n_objects n_objects_predictions = torch.round(n_objects_predictions).int() sem_seg_predictions = sem_seg_predictions.data.cpu() ins_seg_predictions = ins_seg_predictions.data.cpu() n_objects_predictions = n_objects_predictions.data.cpu() return sem_seg_predictions, ins_seg_predictions, n_objects_predictions
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 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)
class Model(object): def __init__( self, dataset, n_classes, max_n_objects, use_instance_segmentation=False, use_coords=False, load_model_path='', usegpu=True): self.dataset = dataset self.n_classes = n_classes self.max_n_objects = max_n_objects self.use_instance_segmentation = use_instance_segmentation self.use_coords = use_coords self.load_model_path = load_model_path self.usegpu = usegpu assert self.dataset in ['CVPPP', 'cityscapes'] if self.dataset == 'CVPPP': self.model = CVPPPArchitecture( self.n_classes, self.use_instance_segmentation, self.use_coords, usegpu=self.usegpu) elif self.dataset == 'cityscapes': self.model = CityscapesArchitecture( self.n_classes, self.use_instance_segmentation, self.use_coords, usegpu=self.usegpu) self.__load_weights() if self.usegpu: cudnn.benchmark = True self.model.cuda() # self.model = torch.nn.DataParallel(self.model, # device_ids=range(self.ngpus)) print self.model self.vis = visdom.Visdom() self.training_metric_vis, self.test_metric_vis = None, None if self.use_instance_segmentation: self.instance_seg_vis = None def __load_weights(self): if self.load_model_path != '': assert os.path.isfile(self.load_model_path), 'Model : {} does not \ exists!'.format(self.load_model_path) print 'Loading model from {}'.format(self.load_model_path) model_state_dict = self.model.state_dict() if self.usegpu: pretrained_state_dict = torch.load(self.load_model_path) else: pretrained_state_dict = torch.load( self.load_model_path, map_location=lambda storage, loc: storage) model_state_dict.update(pretrained_state_dict) self.model.load_state_dict(model_state_dict) def __define_variable(self, tensor, volatile=False): return Variable(tensor, volatile=volatile) def __define_input_variables( self, features, fg_labels, ins_labels, n_objects, mode): volatile = True if mode == 'training': volatile = False features_var = self.__define_variable(features, volatile=volatile) fg_labels_var = self.__define_variable(fg_labels, volatile=volatile) ins_labels_var = self.__define_variable(ins_labels, volatile=volatile) n_objects_var = self.__define_variable(n_objects, volatile=volatile) return features_var, fg_labels_var, ins_labels_var, n_objects_var 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 __define_optimizer(self, learning_rate, weight_decay, lr_drop_factor, lr_drop_patience, optimizer='Adam'): assert optimizer in ['RMSprop', 'Adam', 'Adadelta', 'SGD'] parameters = ifilter(lambda p: p.requires_grad, self.model.parameters()) if optimizer == 'RMSprop': self.optimizer = optim.RMSprop( parameters, lr=learning_rate, weight_decay=weight_decay) elif optimizer == 'Adadelta': self.optimizer = optim.Adadelta( parameters, lr=learning_rate, weight_decay=weight_decay) elif optimizer == 'Adam': self.optimizer = optim.Adam( parameters, lr=learning_rate, weight_decay=weight_decay) elif optimizer == 'SGD': self.optimizer = optim.SGD( parameters, lr=learning_rate, momentum=0.9, weight_decay=weight_decay) self.lr_scheduler = ReduceLROnPlateau( self.optimizer, mode='min', factor=lr_drop_factor, patience=lr_drop_patience, verbose=True) @staticmethod def __get_loss_averager(): return averager() def __minibatch(self, train_test_iter, clip_grad_norm, criterion_type, train_cnn=True, mode='training', debug=False): assert mode in ['training', 'test'], 'Mode must be either "training" or "test"' if mode == 'training': for param in self.model.parameters(): param.requires_grad = True if not train_cnn: for param in self.model.cnn.parameters(): param.requires_grad = False self.model.train() else: for param in self.model.parameters(): param.requires_grad = False self.model.eval() cpu_images, cpu_sem_seg_annotations, \ cpu_ins_seg_annotations, cpu_n_objects = train_test_iter.next() cpu_images = cpu_images.contiguous() cpu_sem_seg_annotations = cpu_sem_seg_annotations.contiguous() cpu_ins_seg_annotations = cpu_ins_seg_annotations.contiguous() cpu_n_objects = cpu_n_objects.contiguous() if self.usegpu: gpu_images = cpu_images.cuda(async=True) gpu_sem_seg_annotations = cpu_sem_seg_annotations.cuda(async=True) gpu_ins_seg_annotations = cpu_ins_seg_annotations.cuda(async=True) gpu_n_objects = cpu_n_objects.cuda(async=True) else: gpu_images = cpu_images gpu_sem_seg_annotations = cpu_sem_seg_annotations gpu_ins_seg_annotations = cpu_ins_seg_annotations gpu_n_objects = cpu_n_objects gpu_images, gpu_sem_seg_annotations, \ gpu_ins_seg_annotations, gpu_n_objects = \ self.__define_input_variables(gpu_images, gpu_sem_seg_annotations, gpu_ins_seg_annotations, gpu_n_objects, mode) gpu_n_objects_normalized = gpu_n_objects.float() / self.max_n_objects sem_seg_predictions, ins_seg_predictions, \ n_objects_predictions = self.model(gpu_images) if mode == 'test': if debug: _vis_prob = np.random.rand() if _vis_prob > 0.7: if self.use_instance_segmentation: sem_seg_preds = np.argmax( sem_seg_predictions.data.cpu().numpy(), axis=1) seg_preds = ins_seg_predictions.data.cpu().numpy() _bs, _n_feats = seg_preds.shape[:2] _sample_idx = np.random.randint(_bs) _sem_seg_preds_sample = sem_seg_preds[_sample_idx] _seg_preds_sample = seg_preds[_sample_idx] fg_ins_embeddings = np.stack( [_seg_preds_sample[i][np.where( _sem_seg_preds_sample == 1)] for i in range(_n_feats)], axis=1) _n_fg_samples = fg_ins_embeddings.shape[0] if _n_fg_samples > 0: fg_ins_embeddings = \ fg_ins_embeddings[np.random.choice( range(_n_fg_samples), size=400)] tsne = TSNE(n_components=2, random_state=0) fg_ins_embeddings_vis = tsne.fit_transform( fg_ins_embeddings) if self.instance_seg_vis: self.vis.scatter(X=fg_ins_embeddings_vis, win=self.instance_seg_vis, opts={'title': 'Predicted Embeddings \ for Foreground \ Predictions', 'markersize': 2}) else: self.instance_seg_vis =\ self.vis.scatter(X=fg_ins_embeddings_vis, opts={'title': 'Predicted \ Embeddings for \ Foreground \ Predictions', 'markersize': 2}) cost = 0.0 out_metrics = dict() if self.use_instance_segmentation: disc_cost = self.criterion_discriminative( ins_seg_predictions, gpu_ins_seg_annotations.float(), cpu_n_objects, self.max_n_objects) cost += disc_cost out_metrics['Discriminative Cost'] = disc_cost.data if criterion_type in ['CE', 'Multi']: _, gpu_sem_seg_annotations_criterion_ce = \ gpu_sem_seg_annotations.max(1) ce_cost = self.criterion_ce( sem_seg_predictions.permute(0, 2, 3, 1).contiguous().view( -1, self.n_classes), gpu_sem_seg_annotations_criterion_ce.view(-1)) cost += ce_cost out_metrics['CE Cost'] = ce_cost.data if criterion_type in ['Dice', 'Multi']: dice_cost = self.criterion_dice( sem_seg_predictions, gpu_sem_seg_annotations) cost += dice_cost out_metrics['Dice Cost'] = dice_cost.data mse_cost = self.criterion_mse( n_objects_predictions, gpu_n_objects_normalized) cost += mse_cost out_metrics['MSE Cost'] = mse_cost.data if mode == 'training': self.model.zero_grad() cost.backward() if clip_grad_norm != 0: torch.nn.utils.clip_grad_norm( self.model.parameters(), clip_grad_norm) self.optimizer.step() return out_metrics def __test(self, test_loader, criterion_type, epoch, debug): n_minibatches = len(test_loader) test_iter = iter(test_loader) out_metrics = dict() for minibatch_index in range(n_minibatches): mb_out_metrics = self.__minibatch( test_iter, 0.0, criterion_type, train_cnn=False, mode='test', debug=debug) for mk, mv in mb_out_metrics.iteritems(): if mk not in out_metrics: out_metrics[mk] = [] out_metrics[mk].append(mv) test_metric_vis_data, test_metric_vis_legend = [], [] metrics_as_str = 'Testing: [METRIC]' for mk, mv in out_metrics.iteritems(): out_metrics[mk] = torch.stack(mv, dim=0).mean() metrics_as_str += ' {} : {} |'.format(mk, out_metrics[mk]) test_metric_vis_data.append(out_metrics[mk]) test_metric_vis_legend.append(mk) print metrics_as_str test_metric_vis_data = np.expand_dims( np.array(test_metric_vis_data), 0) if self.test_metric_vis: self.vis.line(X=np.array([epoch]), Y=test_metric_vis_data, win=self.test_metric_vis, update='append') else: self.test_metric_vis = self.vis.line(X=np.array([epoch]), Y=test_metric_vis_data, opts={'legend': test_metric_vis_legend, 'title': 'Test Metrics', 'showlegend': True, 'xlabel': 'Epoch', 'ylabel': 'Metric'}) return out_metrics def fit(self, criterion_type, delta_var, delta_dist, norm, learning_rate, weight_decay, clip_grad_norm, lr_drop_factor, lr_drop_patience, optimize_bg, optimizer, train_cnn, n_epochs, class_weights, train_loader, test_loader, model_save_path, debug): assert criterion_type in ['CE', 'Dice', 'Multi'] training_log_file = open(os.path.join( model_save_path, 'training.log'), 'w') validation_log_file = open(os.path.join( model_save_path, 'validation.log'), 'w') training_log_file.write('Epoch,Cost\n') validation_log_file.write('Epoch,Cost\n') self.__define_criterion(class_weights, delta_var, delta_dist, norm=norm, optimize_bg=optimize_bg, criterion=criterion_type) self.__define_optimizer(learning_rate, weight_decay, lr_drop_factor, lr_drop_patience, optimizer=optimizer) self.__test(test_loader, criterion_type, -1.0, debug) best_val_cost = np.Inf for epoch in range(n_epochs): epoch_start = time.time() train_iter = iter(train_loader) n_minibatches = len(train_loader) train_out_metrics = dict() minibatch_index = 0 while minibatch_index < n_minibatches: mb_out_metrics = self.__minibatch(train_iter, clip_grad_norm, criterion_type, train_cnn=train_cnn, mode='training', debug=debug) for mk, mv in mb_out_metrics.iteritems(): if mk not in train_out_metrics: train_out_metrics[mk] = [] train_out_metrics[mk].append(mv) minibatch_index += 1 epoch_end = time.time() epoch_duration = epoch_end - epoch_start training_metric_vis_data, training_metric_vis_legend = [], [] print 'Epoch : [{}/{}] - [{}]'.format(epoch, n_epochs, epoch_duration) metrics_as_str = 'Training: [METRIC]' for mk, mv in train_out_metrics.iteritems(): train_out_metrics[mk] = torch.stack(mv, dim=0).mean() metrics_as_str += ' {} : {} |'.format(mk, train_out_metrics[mk]) training_metric_vis_data.append(train_out_metrics[mk]) training_metric_vis_legend.append(mk) print metrics_as_str training_metric_vis_data = np.expand_dims( np.array(training_metric_vis_data), 0) if self.training_metric_vis: self.vis.line(X=np.array([epoch]), Y=training_metric_vis_data, win=self.training_metric_vis, update='append') else: self.training_metric_vis = self.vis.line( X=np.array([epoch]), Y=training_metric_vis_data, opts={'legend': training_metric_vis_legend, 'title': 'Training Metrics', 'showlegend': True, 'xlabel': 'Epoch', 'ylabel': 'Metric'}) val_out_metrics = self.__test( test_loader, criterion_type, epoch, debug) if self.use_instance_segmentation: val_cost = val_out_metrics['Discriminative Cost'] train_cost = train_out_metrics['Discriminative Cost'] elif criterion_type in ['Dice', 'Multi']: val_cost = val_out_metrics['Dice Cost'] train_cost = train_out_metrics['Dice Cost'] else: val_cost = val_out_metrics['CE Cost'] train_cost = train_out_metrics['CE Cost'] self.lr_scheduler.step(val_cost) is_best_model = val_cost <= best_val_cost if is_best_model: best_val_cost = val_cost torch.save(self.model.state_dict(), os.path.join( model_save_path, 'model_{}_{}.pth'.format(epoch, val_cost))) training_log_file.write('{},{}\n'.format(epoch, train_cost)) validation_log_file.write('{},{}\n'.format(epoch, val_cost)) training_log_file.flush() validation_log_file.flush() training_log_file.close() validation_log_file.close() def predict(self, images): assert len(images.size()) == 4 # b, c, h, w for param in self.model.parameters(): param.requires_grad = False self.model.eval() images = images.contiguous() if self.usegpu: images = images.cuda(async=True) images = self.__define_variable(images, volatile=True) sem_seg_predictions, ins_seg_predictions, n_objects_predictions = \ self.model(images) sem_seg_predictions = torch.nn.functional.softmax( sem_seg_predictions, dim=1) n_objects_predictions = n_objects_predictions * self.max_n_objects n_objects_predictions = torch.round(n_objects_predictions).int() sem_seg_predictions = sem_seg_predictions.data.cpu() ins_seg_predictions = ins_seg_predictions.data.cpu() n_objects_predictions = n_objects_predictions.data.cpu() return sem_seg_predictions, ins_seg_predictions, n_objects_predictions
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 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))
'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()