Пример #1
0
def run_training(continue_run):

    logging.info('EXPERIMENT NAME: %s' % exp_config.experiment_name)
    already_created_recursion = False
    print("ALready created recursion : " + str(already_created_recursion))
    init_step = 0
    # Load data
    base_data, recursion_data, recursion = acdc_data.load_and_maybe_process_scribbles(
        scribble_file=sys_config.project_root + exp_config.scribble_data,
        target_folder=log_dir,
        percent_full_sup=exp_config.percent_full_sup,
        scr_ratio=exp_config.length_ratio)
    #wrap everything from this point onwards in a try-except to catch keyboard interrupt so
    #can control h5py closing data
    try:
        loaded_previous_recursion = False
        start_epoch = 0
        if continue_run:
            logging.info(
                '!!!!!!!!!!!!!!!!!!!!!!!!!!!! Continuing previous run !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
            )
            try:
                try:
                    init_checkpoint_path = utils.get_latest_model_checkpoint_path(
                        log_dir, 'recursion_{}_model.ckpt'.format(recursion))

                except:
                    print("EXCEPTE GİRDİ")
                    init_checkpoint_path = utils.get_latest_model_checkpoint_path(
                        log_dir,
                        'recursion_{}_model.ckpt'.format(recursion - 1))
                    loaded_previous_recursion = True
                logging.info('Checkpoint path: %s' % init_checkpoint_path)
                init_step = int(
                    init_checkpoint_path.split('/')[-1].split('-')
                    [-1]) + 1  # plus 1 b/c otherwise starts with eval
                start_epoch = int(init_step /
                                  (len(base_data['images_train']) / 4))
                logging.info('Latest step was: %d' % init_step)
                logging.info('Continuing with epoch: %d' % start_epoch)
            except:
                logging.warning(
                    '!!! Did not find init checkpoint. Maybe first run failed. Disabling continue mode...'
                )
                continue_run = False
                init_step = 0
                start_epoch = 0

            logging.info(
                '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
            )

        if loaded_previous_recursion:
            logging.info(
                "Data file exists for recursion {} "
                "but checkpoints only present up to recursion {}".format(
                    recursion, recursion - 1))
            logging.info("Likely means postprocessing was terminated")

            if not already_created_recursion:

                recursion_data = acdc_data.load_different_recursion(
                    recursion_data, -1)
                recursion -= 1
            else:
                start_epoch = 0
                init_step = 0
        # load images and validation data
        images_train = np.array(base_data['images_train'])
        scribbles_train = np.array(base_data['scribbles_train'])
        images_val = np.array(base_data['images_test'])
        labels_val = np.array(base_data['masks_test'])

        # if exp_config.use_data_fraction:
        #     num_images = images_train.shape[0]
        #     new_last_index = int(float(num_images)*exp_config.use_data_fraction)
        #
        #     logging.warning('USING ONLY FRACTION OF DATA!')
        #     logging.warning(' - Number of imgs orig: %d, Number of imgs new: %d' % (num_images, new_last_index))
        #     images_train = images_train[0:new_last_index,...]
        #     labels_train = labels_train[0:new_last_index,...]

        logging.info('Data summary:')
        logging.info(' - Images:')
        logging.info(images_train.shape)
        logging.info(images_train.dtype)
        #logging.info(' - Labels:')
        #logging.info(labels_train.shape)
        #logging.info(labels_train.dtype)

        # Tell TensorFlow that the model will be built into the default Graph.
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        config.allow_soft_placement = True
        #        with tf.Graph().as_default():
        with tf.Session(config=config) as sess:
            # Generate placeholders for the images and labels.

            image_tensor_shape = [exp_config.batch_size] + list(
                exp_config.image_size) + [1]
            mask_tensor_shape = [exp_config.batch_size] + list(
                exp_config.image_size)

            images_placeholder = tf.placeholder(tf.float32,
                                                shape=image_tensor_shape,
                                                name='images')
            labels_placeholder = tf.placeholder(tf.uint8,
                                                shape=mask_tensor_shape,
                                                name='labels')

            learning_rate_placeholder = tf.placeholder(tf.float32, shape=[])
            crf_learning_rate_placeholder = tf.placeholder(tf.float32,
                                                           shape=[])
            training_time_placeholder = tf.placeholder(tf.bool, shape=[])

            tf.summary.scalar('learning_rate', learning_rate_placeholder)

            # Build a Graph that computes predictions from the inference model.
            logits = model.inference(images_placeholder,
                                     exp_config.model_handle,
                                     training=training_time_placeholder,
                                     nlabels=exp_config.nlabels)

            # Add to the Graph the Ops for loss calculation.
            [loss, _, weights_norm
             ] = model.loss(logits,
                            labels_placeholder,
                            nlabels=exp_config.nlabels,
                            loss_type=exp_config.loss_type,
                            weight_decay=exp_config.weight_decay
                            )  # second output is unregularised loss

            tf.summary.scalar('loss', loss)
            tf.summary.scalar('weights_norm_term', weights_norm)

            crf_variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                              scope='crf_scope')
            print(str(crf_variables))

            spatial_weight = []
            for v in crf_variables:
                if v.name[10:17] == 'spatial':
                    spatial_weight.append(v)
                    print("FOUNDDDDDDD")

#            crf_weights = []
#            crf_weights_name = []
#            for v in tf.all_variables():
#
#                if v.name[0:4]=='bila':
#                    print(str(v))
#                    crf_weights.append(v)
#                    crf_weights_name.append(v.name)
#                elif v.name[0:4] =='spat':
#                    print(str(v))
#                    crf_weights_name.append(v.name)
#                    crf_weights.append(v)
#                elif v.name[0:4] =='comp':
#                    print(str(v))
#                    crf_weights.append(v)
#                    crf_weights_name.append(v.name)
            restore_var = [
                v for v in tf.all_variables() if v.name not in crf_variables
            ]

            # Add to the Graph the Ops that calculate and apply gradients.

            global_step = tf.Variable(0, name='global_step', trainable=False)

            network_train_op = tf.train.AdamOptimizer(
                learning_rate=learning_rate_placeholder).minimize(
                    loss,
                    var_list=restore_var,
                    colocate_gradients_with_ops=True,
                    global_step=global_step)

            reg_loss = tf.nn.l2_loss(spatial_weight)
            crf_train_op = tf.train.AdamOptimizer(
                learning_rate=crf_learning_rate_placeholder).minimize(
                    loss + 0.0001 * reg_loss,
                    var_list=crf_variables,
                    colocate_gradients_with_ops=True,
                    global_step=global_step)

            # Add the Op to compare the logits to the labels during evaluation.
            #            eval_loss = model.evaluation(logits,
            #                                         labels_placeholder,
            #                                         images_placeholder,
            #                                         nlabels=exp_config.nlabels,
            #                                         loss_type=exp_config.loss_type,
            #                                         weak_supervision=True,
            #                                         cnn_threshold=exp_config.cnn_threshold,
            #                                         include_bg=True)
            eval_val_loss = model.evaluation(
                logits,
                labels_placeholder,
                images_placeholder,
                nlabels=exp_config.nlabels,
                loss_type=exp_config.loss_type,
                weak_supervision=True,
                cnn_threshold=exp_config.cnn_threshold,
                include_bg=False)

            # Build the summary Tensor based on the TF collection of Summaries.
            summary = tf.summary.merge_all()

            # Add the variable initializer Op.
            init = tf.global_variables_initializer()

            # Create a saver for writing training checkpoints.
            # Only keep two checkpoints, as checkpoints are kept for every recursion
            # and they can be 300MB +
            saver = tf.train.Saver(max_to_keep=2)
            saver_best_dice = tf.train.Saver(max_to_keep=2)
            saver_best_xent = tf.train.Saver(max_to_keep=2)

            # Create a session for running Ops on the Graph.
            sess = tf.Session()

            # Instantiate a SummaryWriter to output summaries and the Graph.
            summary_writer = tf.summary.FileWriter(log_dir, sess.graph)

            # with tf.name_scope('monitoring'):

            val_error_ = tf.placeholder(tf.float32, shape=[], name='val_error')
            val_error_summary = tf.summary.scalar('validation_loss',
                                                  val_error_)

            val_dice_ = tf.placeholder(tf.float32, shape=[], name='val_dice')
            val_dice_summary = tf.summary.scalar('validation_dice', val_dice_)

            val_summary = tf.summary.merge(
                [val_error_summary, val_dice_summary])

            train_error_ = tf.placeholder(tf.float32,
                                          shape=[],
                                          name='train_error')
            train_error_summary = tf.summary.scalar('training_loss',
                                                    train_error_)

            train_dice_ = tf.placeholder(tf.float32,
                                         shape=[],
                                         name='train_dice')
            train_dice_summary = tf.summary.scalar('training_dice',
                                                   train_dice_)

            train_summary = tf.summary.merge(
                [train_error_summary, train_dice_summary])

            # Run the Op to initialize the variables.
            sess.run(init)

            #            crf_training_variables =tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='crf_training_op')
            #            print(str(crf_training_variables))
            #            all_crf_variables = crf_training_variables+crf_variables
            #            print(str(all_crf_variables))

            network_variables = tf.get_collection(
                tf.GraphKeys.GLOBAL_VARIABLES, scope='network_scope')
            load_variables = []
            for v in network_variables:
                if 'Adam' in v.name:
                    continue
                else:
                    load_variables.append(v)
                # Restore session
#            crf_weights = []
#            for v in tf.all_variables():
#
#                if v.name[0:4]=='bila':
#                    print(str(v))
#                    crf_weights.append(v.name)
#                elif v.name[0:4] =='spat':
#                    print(str(v))
#                    crf_weights.append(v.name)
#                elif v.name[0:4] =='comp':
#                    print(str(v))
#                    crf_weights.append(v.name)
#            restore_var = [v for v in tf.all_variables() if v.name not in all_crf_variables]

#/////////////////////

#            load_saver = tf.train.Saver(var_list=load_variables)
#            load_saver.restore(sess, '/scratch_net/biwirender02/cany/basil/logdir/unet2D_ws_spot_blur/recursion_0_model.ckpt-7499')

#///////////////////////////

#            if continue_run:
#                # Restore session
#                saver.restore(sess, init_checkpoint_path)
            saver.restore(
                sess,
                '/scratch_net/biwirender02/cany/scribble/logdir/twisted_full_network/recursion_2_model_best_dice.ckpt-43499'
            )
            step = init_step
            curr_lr = exp_config.learning_rate
            crf_curr_lr = 1e-06
            no_improvement_counter = 0
            best_val = np.inf
            last_train = np.inf
            loss_history = []
            loss_gradient = np.inf
            best_dice = 0
            logging.info('RECURSION {0}'.format(recursion))

            # random walk - if it already has been random walked it won't redo
            if start_epoch == 0:
                recursion_data = acdc_data.random_walk_epoch(
                    recursion_data, exp_config.rw_beta,
                    exp_config.rw_threshold, exp_config.random_walk)
            print("Random walku geçti")
            #get ground truths
            labels_train = np.array(recursion_data['random_walked'])
            print("Start epoch : " + str(start_epoch) + " : max epochs : " +
                  str(exp_config.epochs_per_recursion))
            for epoch in range(start_epoch, exp_config.max_epochs):
                if (epoch % exp_config.epochs_per_recursion == 0 and epoch != 0) \
                        or loaded_previous_recursion:

                    loaded_previous_recursion = False
                    #Have reached end of recursion
                    recursion_data = predict_next_gt(
                        data=recursion_data,
                        images_train=images_train,
                        images_placeholder=images_placeholder,
                        training_time_placeholder=training_time_placeholder,
                        logits=logits,
                        sess=sess)

                    #                        recursion_data = postprocess_gt(data=recursion_data,
                    #                                                        images_train=images_train,
                    #                                                        scribbles_train=scribbles_train)
                    recursion += 1
                    #                        # random walk - if it already has been random walked it won't redo
                    #                        recursion_data = acdc_data.random_walk_epoch(recursion_data,
                    #                                                                     exp_config.rw_beta,
                    #                                                                     exp_config.rw_threshold,
                    #                                                                     exp_config.random_walk)
                    #get ground truths
                    labels_train = np.array(recursion_data['predicted'])

                    #reinitialise savers - otherwise, no checkpoints will be saved for each recursion
                    saver = tf.train.Saver(max_to_keep=2)
                    saver_best_dice = tf.train.Saver(max_to_keep=2)
                    saver_best_xent = tf.train.Saver(max_to_keep=2)
                logging.info(
                    'Epoch {0} ({1} of {2} epochs for recursion {3})'.format(
                        epoch, 1 + epoch % exp_config.epochs_per_recursion,
                        exp_config.epochs_per_recursion, recursion))
                # for batch in iterate_minibatches(images_train,
                #                                  labels_train,
                #                                  batch_size=exp_config.batch_size,
                #                                  augment_batch=exp_config.augment_batch):

                # You can run this loop with the BACKGROUND GENERATOR, which will lead to some improvements in the
                # training speed. However, be aware that currently an exception inside this loop may not be caught.
                # The batch generator may just continue running silently without warning even though the code has
                # crashed.

                for batch in BackgroundGenerator(
                        iterate_minibatches(
                            images_train,
                            labels_train,
                            batch_size=exp_config.batch_size,
                            augment_batch=exp_config.augment_batch)):

                    if exp_config.warmup_training:
                        if step < 50:
                            curr_lr = exp_config.learning_rate / 10.0
                        elif step == 50:
                            curr_lr = exp_config.learning_rate
                    if ((step % 5000 == 0) & (step > 0)):
                        curr_lr = curr_lr * 0.94
                    start_time = time.time()

                    # batch = bgn_train.retrieve()
                    x, y = batch

                    # TEMPORARY HACK (to avoid incomplete batches
                    if y.shape[0] < exp_config.batch_size:
                        step += 1
                        continue

                    network_feed_dict = {
                        images_placeholder: x,
                        labels_placeholder: y,
                        learning_rate_placeholder: curr_lr,
                        training_time_placeholder: True
                    }

                    crf_feed_dict = {
                        images_placeholder: x,
                        labels_placeholder: y,
                        crf_learning_rate_placeholder: crf_curr_lr,
                        training_time_placeholder: True
                    }
                    if (step % 1000 == 0):
                        print("CRF variables : " +
                              str(sess.run(crf_variables)))
                    if (step % 10 == 0):
                        _, loss_value = sess.run([crf_train_op, loss],
                                                 feed_dict=crf_feed_dict)
                    _, loss_value = sess.run([network_train_op, loss],
                                             feed_dict=network_feed_dict)
                    duration = time.time() - start_time

                    # Write the summaries and print an overview fairly often.
                    if step % 10 == 0:
                        # Print status to stdout.
                        logging.info('Step %d: loss = %.6f (%.3f sec)' %
                                     (step, loss_value, duration))
                        # Update the events file.


#                        summary_str = sess.run(summary, feed_dict=feed_dict)
#                        summary_writer.add_summary(summary_str, step)
#                        summary_writer.flush()

#                    if (step + 1) % exp_config.train_eval_frequency == 0:
#
#                        logging.info('Training Data Eval:')
#                        [train_loss, train_dice] = do_eval(sess,
#                                                           eval_loss,
#                                                           images_placeholder,
#                                                           labels_placeholder,
#                                                           training_time_placeholder,
#                                                           images_train,
#                                                           labels_train,
#                                                           exp_config.batch_size)
#
#                        train_summary_msg = sess.run(train_summary, feed_dict={train_error_: train_loss,
#                                                                               train_dice_: train_dice}
#                                                     )
#                        summary_writer.add_summary(train_summary_msg, step)
#
#                        loss_history.append(train_loss)
#                        if len(loss_history) > 5:
#                            loss_history.pop(0)
#                            loss_gradient = (loss_history[-5] - loss_history[-1]) / 2
#
#                        logging.info('loss gradient is currently %f' % loss_gradient)
#
#                        if exp_config.schedule_lr and loss_gradient < exp_config.schedule_gradient_threshold:
#                            logging.warning('Reducing learning rate!')
#                            curr_lr /= 10.0
#                            logging.info('Learning rate changed to: %f' % curr_lr)
#
#                            # reset loss history to give the optimisation some time to start decreasing again
#                            loss_gradient = np.inf
#                            loss_history = []
#
#                        if train_loss <= last_train:  # best_train:
#                            logging.info('Decrease in training error!')
#                        else:
#                            logging.info('No improvement in training error for %d steps' % no_improvement_counter)
#
#                        last_train = train_loss

# Save a checkpoint and evaluate the model periodically.
                    if (step + 1) % exp_config.val_eval_frequency == 0:

                        checkpoint_file = os.path.join(
                            log_dir,
                            'recursion_{}_model.ckpt'.format(recursion))
                        saver.save(sess, checkpoint_file, global_step=step)
                        # Evaluate against the training set.

                        # Evaluate against the validation set.
                        logging.info('Validation Data Eval:')
                        [val_loss, val_dice
                         ] = do_eval(sess, eval_val_loss, images_placeholder,
                                     labels_placeholder,
                                     training_time_placeholder, images_val,
                                     labels_val, exp_config.batch_size)

                        val_summary_msg = sess.run(val_summary,
                                                   feed_dict={
                                                       val_error_: val_loss,
                                                       val_dice_: val_dice
                                                   })
                        summary_writer.add_summary(val_summary_msg, step)

                        if val_dice > best_dice:
                            best_dice = val_dice
                            best_file = os.path.join(
                                log_dir,
                                'recursion_{}_model_best_dice.ckpt'.format(
                                    recursion))
                            saver_best_dice.save(sess,
                                                 best_file,
                                                 global_step=step)
                            logging.info(
                                'Found new best dice on validation set! - {} - '
                                'Saving recursion_{}_model_best_dice.ckpt'.
                                format(val_dice, recursion))
                            text_file = open('val_results.txt', "a")
                            text_file.write("\nVal dice " + str(step) + " : " +
                                            str(val_dice))
                            text_file.close()
                        if val_loss < best_val:
                            best_val = val_loss
                            best_file = os.path.join(
                                log_dir,
                                'recursion_{}_model_best_xent.ckpt'.format(
                                    recursion))
                            saver_best_xent.save(sess,
                                                 best_file,
                                                 global_step=step)
                            logging.info(
                                'Found new best crossentropy on validation set! - {} - '
                                'Saving recursion_{}_model_best_xent.ckpt'.
                                format(val_loss, recursion))

                    step += 1

    except Exception:
        raise
def run_training(continue_run):

    logging.info('EXPERIMENT NAME: %s' % exp_config.experiment_name)

    init_step = 0
    # Load data
    base_data, recursion_data, recursion = acdc_data.load_and_maybe_process_scribbles(scribble_file=sys_config.project_root + exp_config.scribble_data,
                                                                                      target_folder=log_dir,
                                                                                      percent_full_sup=exp_config.percent_full_sup,
                                                                                      scr_ratio=exp_config.length_ratio
                                                                                      )
    #wrap everything from this point onwards in a try-except to catch keyboard interrupt so
    #can control h5py closing data
    try:
        loaded_previous_recursion = False
        start_epoch = 0
        if continue_run:
            logging.info('!!!!!!!!!!!!!!!!!!!!!!!!!!!! Continuing previous run !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
            try:
                try:
                    init_checkpoint_path = utils.get_latest_model_checkpoint_path(log_dir, 'recursion_{}_model.ckpt'.format(recursion))
                except:
                    init_checkpoint_path = utils.get_latest_model_checkpoint_path(log_dir, 'recursion_{}_model.ckpt'.format(recursion - 1))
                    loaded_previous_recursion = True
                logging.info('Checkpoint path: %s' % init_checkpoint_path)
                init_step = int(init_checkpoint_path.split('/')[-1].split('-')[-1]) + 1  # plus 1 b/c otherwise starts with eval
                start_epoch = int(init_step/(len(base_data['images_train'])/4))
                logging.info('Latest step was: %d' % init_step)
                logging.info('Continuing with epoch: %d' % start_epoch)
            except:
                logging.warning('!!! Did not find init checkpoint. Maybe first run failed. Disabling continue mode...')
                continue_run = False
                init_step = 0
                start_epoch = 0

            logging.info('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')


        if loaded_previous_recursion:
            logging.info("Data file exists for recursion {} "
                         "but checkpoints only present up to recursion {}".format(recursion, recursion - 1))
            logging.info("Likely means postprocessing was terminated")
            recursion_data = acdc_data.load_different_recursion(recursion_data, -1)
            recursion-=1

        # load images and validation data
        images_train = np.array(base_data['images_train'])
        scribbles_train = np.array(base_data['scribbles_train'])
        images_val = np.array(base_data['images_test'])
        labels_val = np.array(base_data['masks_test'])


        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        config.allow_soft_placement = True
#        with tf.Graph().as_default():
        with tf.Session(config = config) as sess:
            # Generate placeholders for the images and labels.

            image_tensor_shape = [exp_config.batch_size] + list(exp_config.image_size) + [1]
            mask_tensor_shape = [exp_config.batch_size] + list(exp_config.image_size)

            images_placeholder = tf.placeholder(tf.float32, shape=image_tensor_shape, name='images')
            labels_placeholder = tf.placeholder(tf.uint8, shape=mask_tensor_shape, name='labels')

            learning_rate_placeholder = tf.placeholder(tf.float32, shape=[])
            training_time_placeholder = tf.placeholder(tf.bool, shape=[])

            tf.summary.scalar('learning_rate', learning_rate_placeholder)

            # Build a Graph that computes predictions from the inference model.
            logits = model.inference(images_placeholder,
                                     exp_config.model_handle,
                                     training=training_time_placeholder,
                                     nlabels=exp_config.nlabels)


            summary = tf.summary.merge_all()

            # Add the variable initializer Op.
            init = tf.global_variables_initializer()



            # Run the Op to initialize the variables.
            sess.run(init)
            
            saver = tf.train.Saver(max_to_keep=2)

            if continue_run:
                # Restore session
                saver.restore(sess, init_checkpoint_path)

            step = init_step
            curr_lr = exp_config.learning_rate

            no_improvement_counter = 0
            best_val = np.inf
            last_train = np.inf
            loss_history = []
            loss_gradient = np.inf
            best_dice = 0
            logging.info('RECURSION {0}'.format(recursion))

#            # random walk - if it already has been random walked it won't redo
#            recursion_data = acdc_data.random_walk_epoch(recursion_data, exp_config.rw_beta, exp_config.rw_threshold, exp_config.random_walk)


                    #Have reached end of recursion
            recursion_data = predict_next_gt(data=recursion_data,
                                             images_train=images_train,
                                             images_placeholder=images_placeholder,
                                             training_time_placeholder=training_time_placeholder,
                                             logits=logits,
                                             sess=sess)

            recursion_data = postprocess_gt(data=recursion_data,
                                            images_train=images_train,
                                            scribbles_train=scribbles_train)
            recursion += 1
            # random walk - if it already has been random walked it won't redo
            recursion_data = acdc_data.random_walk_epoch(recursion_data,
                                                         exp_config.rw_beta,
                                                         exp_config.rw_threshold,
                                                         exp_config.random_walk)
    except Exception:
        raise