def train_ffn(model_cls, **model_kwargs): with tf.Graph().as_default(): with tf.device( tf.train.replica_device_setter(FLAGS.ps_tasks, merge_devices=True)): # The constructor might define TF ops/placeholders, so it is important # that the FFN is instantiated within the current context. model = model_cls(with_membrane=FLAGS.with_membrane, is_training=not (FLAGS.validation_mode), adabn=FLAGS.adabn, grad_clip_val=FLAGS.cap_gradient, **model_kwargs) eval_shape_zyx = train_eval_size(model).tolist()[::-1] eval_tracker = EvalTracker(eval_shape_zyx) load_data_ops = define_data_input(model, queue_batch=1) prepare_ffn(model) merge_summaries_op = tf.summary.merge_all() if FLAGS.task == 0: save_flags() # Start supervisor. train_dir = FLAGS.train_dir summary_rate_secs = 999999 save_model_secs = 30 if FLAGS.topup_mode else 999999 sv = tf.train.Supervisor(logdir=train_dir, is_chief=(FLAGS.task == 0), saver=model.saver, summary_op=None, save_summaries_secs=summary_rate_secs, save_model_secs=save_model_secs, recovery_wait_secs=5) sess = sv.prepare_or_wait_for_session( FLAGS.master, config=tf.ConfigProto(log_device_placement=False, allow_soft_placement=True)) eval_tracker.sess = sess # TODO (jk): load from ckpt step = int(sess.run(model.global_step)) step_since_session_start = 0 if FLAGS.task > 0: # To avoid early instabilities when using multiple replicas, we use # a launch schedule where new replicas are brought online gradually. logging.info('Delaying replica start.') while True: if (int(sess.run(model.global_step)) >= FLAGS.replica_step_delay * FLAGS.task): break time.sleep(5.0) else: summary_writer = tf.summary.FileWriterCache.get( FLAGS.train_dir) summary_writer.add_session_log( tf.summary.SessionLog(status=tf.summary.SessionLog.START), step) fov_shifts = list(model.shifts) # x, y, z if FLAGS.shuffle_moves: random.shuffle(fov_shifts) policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts), 'max_pred_moves': max_pred_offsets } batch_it = get_batch(lambda: sess.run(load_data_ops), eval_tracker, model, FLAGS.batch_size, policy_map[FLAGS.fov_policy]) t_last = time.time() # TODO (jk): text log of learning curve. refresh file. max_steps = FLAGS.eval_steps step_since_session_start = 0 while step_since_session_start < max_steps: if (step % 20 == 0) & (step_since_session_start > 0): # TODO (jk): text log of learning curve. refresh file. logging.info( 'Step: ' + str(step) + ', prec: ' + str( np.round(1000 * eval_tracker.tp / (eval_tracker.tp + eval_tracker.fp + 0.000001))) + ', recll: ' + str( np.round(1000 * eval_tracker.tp / (eval_tracker.tp + eval_tracker.fn + 0.000001))) + ', acc: ' + str( np.round(1000 * (eval_tracker.tp + eval_tracker.tn) / (eval_tracker.tp + eval_tracker.tn + eval_tracker.fp + eval_tracker.fn + 0.000001))) + ', #patches: ' + str(eval_tracker.num_patches)) seed, patches, labels, weights = next(batch_it) updated_seed, step, summ, accuracy = run_training_step( sess, model, None, None, feed_dict={ model.loss_weights: weights, model.labels: labels, model.offset_label: 'off', model.input_patches: patches, model.input_seed: seed, }) step += 1 step_since_session_start += 1 mask.update_at(seed, (0, 0, 0), updated_seed) # RECORD RESULT eval_curve_txt = open(os.path.join(FLAGS.train_dir, 'eval.txt'), "a") eval_curve_txt.write( '\nStep: ' + str(step) + ', prec: ' + str((eval_tracker.tp / (eval_tracker.tp + eval_tracker.fp + 0.000001))) + ', recll: ' + str((eval_tracker.tp / (eval_tracker.tp + eval_tracker.fn + 0.000001))) + ', acc: ' + str(((eval_tracker.tp + eval_tracker.tn) / (eval_tracker.tp + eval_tracker.tn + eval_tracker.fp + eval_tracker.fn + 0.000001))) + ', #patches: ' + str(eval_tracker.num_patches)) eval_curve_txt.close() print( ' prec: ' + str( np.round(1000 * eval_tracker.tp / (eval_tracker.tp + eval_tracker.fp + 0.000001))) + ', recll: ' + str( np.round(1000 * eval_tracker.tp / (eval_tracker.tp + eval_tracker.fn + 0.000001))) + ', acc: ' + str( np.round(1000 * (eval_tracker.tp + eval_tracker.tn) / (eval_tracker.tp + eval_tracker.tn + eval_tracker.fp + eval_tracker.fn + 0.000001))))
def train_ffn(model_cls, **model_kwargs): with tf.Graph().as_default(): with tf.device( tf.train.replica_device_setter(FLAGS.ps_tasks, merge_devices=True)): # The constructor might define TF ops/placeholders, so it is important # that the FFN is instantiated within the current context. model = model_cls(**model_kwargs) eval_shape_zyx = train_eval_size(model).tolist()[::-1] eval_tracker = EvalTracker(eval_shape_zyx) load_data_ops = define_data_input(model, queue_batch=1) prepare_ffn(model) merge_summaries_op = tf.summary.merge_all() if FLAGS.task == 0: save_flags() summary_writer = None saver = tf.train.Saver(keep_checkpoint_every_n_hours=0.25) scaffold = tf.train.Scaffold(saver=saver) with tf.train.MonitoredTrainingSession( master=FLAGS.master, is_chief=(FLAGS.task == 0), save_summaries_steps=None, save_checkpoint_secs=300, config=tf.ConfigProto(log_device_placement=False, allow_soft_placement=True), checkpoint_dir=FLAGS.train_dir, scaffold=scaffold) as sess: eval_tracker.sess = sess step = int(sess.run(model.global_step)) if FLAGS.task > 0: # To avoid early instabilities when using multiple replicas, we use # a launch schedule where new replicas are brought online gradually. logging.info('Delaying replica start.') while step < FLAGS.replica_step_delay * FLAGS.task: time.sleep(5.0) step = int(sess.run(model.global_step)) else: summary_writer = tf.summary.FileWriterCache.get( FLAGS.train_dir) summary_writer.add_session_log( tf.summary.SessionLog( status=tf.summary.SessionLog.START), step) fov_shifts = list(model.shifts) # x, y, z if FLAGS.shuffle_moves: random.shuffle(fov_shifts) policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts), 'max_pred_moves': max_pred_offsets } batch_it = get_batch(lambda: sess.run(load_data_ops), eval_tracker, model, FLAGS.batch_size, policy_map[FLAGS.fov_policy]) t_last = time.time() while not sess.should_stop() and step < FLAGS.max_steps: # Run summaries periodically. t_curr = time.time() if t_curr - t_last > FLAGS.summary_rate_secs and FLAGS.task == 0: summ_op = merge_summaries_op t_last = t_curr else: summ_op = None seed, patches, labels, weights = next(batch_it) updated_seed, step, summ = run_training_step( sess, model, summ_op, feed_dict={ model.loss_weights: weights, model.labels: labels, model.input_patches: patches, model.input_seed: seed, }) # Save prediction results in the original seed array so that # they can be used in subsequent steps. mask.update_at(seed, (0, 0, 0), updated_seed) # Record summaries. if summ is not None: logging.info('Saving summaries.') summ = tf.Summary.FromString(summ) # Compute a loss over the whole training patch (i.e. more than a # single-step field of view of the network). This quantifies the # quality of the final object mask. summ.value.extend(eval_tracker.get_summaries()) eval_tracker.reset() assert summary_writer is not None summary_writer.add_summary(summ, step) if summary_writer is not None: summary_writer.flush()
def train_ffn(model_cls, **model_kwargs): with tf.Graph().as_default(): with tf.device( tf.train.replica_device_setter(FLAGS.ps_tasks, merge_devices=True)): # The constructor might define TF ops/placeholders, so it is important # that the FFN is instantiated within the current context. if not FLAGS.validation_mode: model = model_cls(with_membrane=FLAGS.with_membrane, is_training=True, grad_clip_val=5.0, **model_kwargs) else: model = model_cls(with_membrane=FLAGS.with_membrane, is_training=False, adabn=FLAGS.adabn, **model_kwargs) eval_shape_zyx = train_eval_size(model).tolist()[::-1] eval_tracker = EvalTracker(eval_shape_zyx) load_data_ops = define_data_input(model, queue_batch=1) prepare_ffn(model) merge_summaries_op = tf.summary.merge_all() if FLAGS.task == 0: save_flags() # Start supervisor. if not FLAGS.validation_mode: save_model_secs = 1800 summary_rate_secs = FLAGS.summary_rate_secs train_dir = FLAGS.train_dir sv = tf.train.Supervisor(logdir=train_dir, is_chief=(FLAGS.task == 0), saver=model.saver, summary_op=None, save_summaries_secs=summary_rate_secs, save_model_secs=save_model_secs, recovery_wait_secs=5) sess = sv.prepare_or_wait_for_session( FLAGS.master, config=tf.ConfigProto(log_device_placement=False, allow_soft_placement=True)) else: train_dir = FLAGS.train_dir + 'v' summary_rate_secs = 999999 save_model_secs = 999999 sv = tf.train.Supervisor(logdir=train_dir, is_chief=(FLAGS.task == 0), saver=None, summary_op=None, save_summaries_secs=summary_rate_secs, save_model_secs=save_model_secs, recovery_wait_secs=5) sess = sv.prepare_or_wait_for_session( FLAGS.master, config=tf.ConfigProto(log_device_placement=False, allow_soft_placement=True)) eval_tracker.sess = sess # TODO (jk): load from ckpt if FLAGS.load_from_ckpt != 'None': logging.info('>>>>>>>>>>>>>>>>>>>>> Loading checkpoint.') model.saver.restore(eval_tracker.sess, FLAGS.load_from_ckpt) logging.info('>>>>>>>>>>>>>>>>>>>>> Checkpoint loaded.') step = int(sess.run(model.global_step)) step_since_session_start = 0 if FLAGS.task > 0: # To avoid early instabilities when using multiple replicas, we use # a launch schedule where new replicas are brought online gradually. logging.info('Delaying replica start.') while True: if (int(sess.run(model.global_step)) >= FLAGS.replica_step_delay * FLAGS.task): break time.sleep(5.0) else: summary_writer = tf.summary.FileWriterCache.get( FLAGS.train_dir) summary_writer.add_session_log( tf.summary.SessionLog(status=tf.summary.SessionLog.START), step) fov_shifts = list(model.shifts) # x, y, z if FLAGS.shuffle_moves: random.shuffle(fov_shifts) policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts), 'max_pred_moves': max_pred_offsets } batch_it = get_batch(lambda: sess.run(load_data_ops), eval_tracker, model, FLAGS.batch_size, policy_map[FLAGS.fov_policy]) t_last = time.time() if not FLAGS.validation_mode: # TODO (jk): text log of learning curve. refresh file. learning_curve_txt = open( os.path.join(FLAGS.train_dir, 'lc.txt'), "w") learning_curve_txt.close() max_steps = FLAGS.max_steps else: max_steps = step + 50000 / FLAGS.batch_size # if FLAGS.adabn: # ####################### BASIC ADABN # sess.run(model.ada_initializer) mean1o = 0. mean2o = 0. var1o = 1. var2o = 1. while step < max_steps: if (step % 20 == 0) & (step_since_session_start > 0): # TODO (jk): text log of learning curve. refresh file. logging.info( 'Step: ' + str(step) + ', prec: ' + str( np.round(1000 * eval_tracker.tp / (eval_tracker.tp + eval_tracker.fp + 0.000001))) + ', recll: ' + str( np.round(1000 * eval_tracker.tp / (eval_tracker.tp + eval_tracker.fn + 0.000001))) + ', acc: ' + str( np.round(1000 * (eval_tracker.tp + eval_tracker.tn) / (eval_tracker.tp + eval_tracker.tn + eval_tracker.fp + eval_tracker.fn + 0.000001))) + ', #patches: ' + str(eval_tracker.num_patches)) # Run summaries periodically. t_curr = time.time() # TIME-BASED SUMMARY # if t_curr - t_last > FLAGS.summary_rate_secs and FLAGS.task == 0: # summ_op = merge_summaries_op # t_last = t_curr # else: # summ_op = None # TODO (jk): ITERATION-BASED SUMMARY if (step % 100 == 0) & (step > 0): summ_op = merge_summaries_op else: summ_op = None if FLAGS.validation_mode & ( (step_since_session_start % 500) == 0) & (step_since_session_start > 0): print('REFRESHING EVAL TRACKER...') eval_tracker.reset() # if ((step_since_session_start % 5) == 0) & (model.moment_list is not None): # mean1 = sess.run(model.moment_list[0].name) # mean2 = sess.run(model.moment_list[-2].name) # var1 = sess.run(model.moment_list[1].name) # var2 = sess.run(model.moment_list[-1].name) # diff = np.sum(np.square(mean1-mean1o)) + np.sum(np.square(mean2-mean2o)) + np.sum(np.square(var1-var1o)) + np.sum(np.square(var2-var2o)) # mean1o = mean1 # mean2o = mean2 # var1o = var1 # var2o = var2 # print('moment displacement = ' + str(diff)) # # print(mean1) # # print(var1) # if FLAGS.adabn: # ####################### REMOVE THIS TO TURN OFF INSTANCE NORMALIZATION # sess.run(model.ada_initializer) # if ((step_since_session_start % 50) == 0) & (step_since_session_start > 0): # sess.run(model.ada_initializer) # # sess.run(model.fgru_ada_initializer) # # sess.run(model.ext_ada_initializer) # eval_tracker.reset() # print('REFRESHING MOMENTS, eval tracker...') seed, patches, labels, weights = next(batch_it) updated_seed, step, summ, accuracy = run_training_step( sess, model, summ_op, None, feed_dict={ model.loss_weights: weights, model.labels: labels, model.offset_label: 'off', model.input_patches: patches, model.input_seed: seed, }) if FLAGS.validation_mode: step += 1 step_since_session_start += 1 # Save prediction results in the original seed array so that # they can be used in subsequent steps. mask.update_at(seed, (0, 0, 0), updated_seed) # Record summaries. if summ is not None: # TODO (jk): text log of learning curve learning_curve_txt = open( os.path.join(FLAGS.train_dir, 'lc.txt'), "a") precision = eval_tracker.tp / (eval_tracker.tp + eval_tracker.fp + 0.0001) recall = eval_tracker.tp / (eval_tracker.tp + eval_tracker.fn + 0.0001) accuracy = (eval_tracker.tp + eval_tracker.tn) / ( eval_tracker.tp + eval_tracker.tn + eval_tracker.fp + eval_tracker.fn + 0.0001) if not FLAGS.validation_mode: learning_curve_txt.write('step_' + str(step) + '_precision_' + str(precision) + '_recall_' + str(recall) + '_accuracy_' + str(accuracy)) learning_curve_txt.write("\n") learning_curve_txt.close() logging.info('Saving summaries.') summ = tf.Summary.FromString(summ) # Compute a loss over the whole training patch (i.e. more than a # single-step field of view of the network). This quantifies the # quality of the final object mask. summ.value.extend(eval_tracker.get_summaries()) eval_tracker.reset() assert summary_writer is not None summary_writer.add_summary(summ, step) if np.min([precision, recall]) > 0.97: logging.info( '>>>>>>>>>>>>>>>>>>>>> Target performance (both prec and recall >0.9) reached.' ) break if summary_writer is not None: summary_writer.flush()
def train_ffn(model_cls, **model_kwargs): hvd.init() logging.info('Rank: %d %d' % (hvd.rank(), rank)) with tf.Graph().as_default(): # The constructor might define TF ops/placeholders, so it is important # that the FFN is instantiated within the current context. model = model_cls(**model_kwargs) eval_shape_zyx = train_eval_size(model).tolist()[::-1] eval_tracker = EvalTracker(eval_shape_zyx) load_data_ops = h5_distributed_dataset(model, queue_batch=1) print(load_data_ops) prepare_ffn(model) merge_summaries_op = tf.compat.v1.summary.merge_all() if FLAGS.task == 0: save_flags() hooks = [ # Horovod: BroadcastGlobalVariablesHook broadcasts initial variable states # from rank 0 to all other processes. This is necessary to ensure consistent # initialization of all workers when training is started with random weights # or restored from a checkpoint. hvd.BroadcastGlobalVariablesHook(0), # Horovod: adjust number of steps based on number of GPUs. tf.estimator.StopAtStepHook(last_step=FLAGS.max_steps // hvd.size()), ] config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True config.gpu_options.visible_device_list = str(hvd.local_rank()) checkpoint_dir = FLAGS.train_dir if hvd.rank() == 0 else None summary_writer = None saver = tf.compat.v1.train.Saver(max_to_keep=None, keep_checkpoint_every_n_hours=24) scaffold = tf.compat.v1.train.Scaffold(saver=saver) # model.global_step = None # logging.warning('GLOBAL STEP %s %s', model.global_step, model.global_step.dtype.base_dtype.is_integer) # logging.warning('GLOBAL_STEP assert') # tf.compat.v1.train.training_utils.assert_global_step(model.global_step) # tf.train.training_utils.assert_global_step(model.global_step) # global_step = tf.Variable(0, name='global_step', trainable=False) # g = ops.get_default_graph() # logging.warning('GRAPH %s', g) # glst = g.get_collection(ops.GraphKeys.GLOBAL_STEP) # logging.warning('GRAPH GLOBAL%s', glst) with tf.compat.v1.train.MonitoredTrainingSession( master=FLAGS.master, is_chief=(FLAGS.task == 0), checkpoint_dir=checkpoint_dir, hooks=hooks, save_checkpoint_secs=300, save_summaries_steps=None, config=config, scaffold=scaffold) as sess: eval_tracker.sess = sess step = int(sess.run(model.global_step)) if FLAGS.task > 0: # To avoid early instabilities when using multiple replicas, we use # a launch schedule where new replicas are brought online gradually. logging.info('Delaying replica start.') while step < FLAGS.replica_step_delay * FLAGS.task: time.sleep(5.0) if rank == 0: summary_writer = tf.compat.v1.summary.FileWriterCache.get( FLAGS.train_dir) summary_writer.add_session_log( tf.compat.v1.summary.SessionLog( status=tf.compat.v1.summary.SessionLog.START), step) fov_shifts = list(model.shifts) # x, y, z if FLAGS.shuffle_moves: random.shuffle(fov_shifts) policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts), 'max_pred_moves': max_pred_offsets } batch_it = get_batch( lambda: sess.run(load_data_ops), eval_tracker, model, FLAGS.batch_size, # eval_tracker, model, 1, policy_map[FLAGS.fov_policy]) t_last = time.time() while not sess.should_stop() and step < FLAGS.max_steps: # Run summaries periodically. t_curr = time.time() if t_curr - t_last > FLAGS.summary_rate_secs and FLAGS.task == 0: summ_op = merge_summaries_op t_last = t_curr else: summ_op = None seed, patches, labels, weights = next(batch_it) updated_seed, step, summ = run_training_step( sess, model, summ_op, feed_dict={ model.loss_weights: weights, model.labels: labels, model.input_patches: patches, model.input_seed: seed, }) # Save prediction results in the original seed array so that # they can be used in subsequent steps. mask.update_at(seed, (0, 0, 0), updated_seed) # Record summaries. if hvd.rank() == 0 and summ is not None: logging.info('Saving summaries.') summ = tf.compat.v1.Summary.FromString(summ) # Compute a loss over the whole training patch (i.e. more than a # single-step field of view of the network). This quantifies the # quality of the final object mask. summ.value.extend(eval_tracker.get_summaries()) eval_tracker.reset() assert summary_writer is not None summary_writer.add_summary(summ, step) if summary_writer is not None: summary_writer.flush()
def train_ffn(model_cls, **model_kwargs): with tf.Graph().as_default(): with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks, merge_devices=True)): # The constructor might define TF ops/placeholders, so it is important # that the FFN is instantiated within the current context. model = model_cls(**model_kwargs) eval_shape_zyx = train_eval_size(model).tolist()[::-1] eval_tracker = EvalTracker(eval_shape_zyx) load_data_ops = define_data_input(model, queue_batch=1) prepare_ffn(model) merge_summaries_op = tf.summary.merge_all() if FLAGS.task == 0: save_flags() # Start supervisor. sv = tf.train.Supervisor( logdir=FLAGS.train_dir, is_chief=(FLAGS.task == 0), saver=model.saver, summary_op=None, save_summaries_secs=FLAGS.summary_rate_secs, save_model_secs=300, recovery_wait_secs=5) sess = sv.prepare_or_wait_for_session( FLAGS.master, config=tf.ConfigProto( log_device_placement=False, allow_soft_placement=True)) eval_tracker.sess = sess options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) run_metadata = tf.RunMetadata() if FLAGS.task > 0: # To avoid early instabilities when using multiple replicas, we use # a launch schedule where new replicas are brought online gradually. logging.info('Delaying replica start.') while True: if (int(sess.run(model.global_step,options=options, run_metadata=run_metadata)) >= FLAGS.replica_step_delay * FLAGS.task): break time.sleep(5.0) fov_shifts = list(model.shifts) # x, y, z if FLAGS.shuffle_moves: random.shuffle(fov_shifts) policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts), 'max_pred_moves': max_pred_offsets } batch_it = get_batch(lambda: sess.run(load_data_ops,options=options, run_metadata=run_metadata), eval_tracker, model, FLAGS.batch_size, policy_map[FLAGS.fov_policy]) fetched_timeline = timeline.Timeline(run_metadata.step_stats) chrome_trace = fetched_timeline.generate_chrome_trace_format() with open('timeline_01.json', 'w') as f: f.write(chrome_trace) step = 0 t_last = time.time() while step < FLAGS.max_steps: # Run summaries periodically. logging.info('Iteration' + str(step)) t_curr = time.time() if t_curr - t_last > FLAGS.summary_rate_secs and FLAGS.task == 0: summ_op = merge_summaries_op t_last = t_curr else: summ_op = None seed, patches, labels, weights = next(batch_it) summaries = [] updated_seed, step, summ = run_training_step( sess, model, summ_op, feed_dict={ model.loss_weights: weights, model.labels: labels, model.offset_label: 'off', model.input_patches: patches, model.input_seed: seed, }) # Save prediction results in the original seed array so that # they can be used in subsequent steps. mask.update_at(seed, (0, 0, 0), updated_seed) if summ is not None: summaries.append(tf.Summary.FromString(summ)) # Record summaries. if FLAGS.task == 0 and summ_op is not None: # Compute a loss over the whole training patch (i.e. more than a # single-step field of view of the network). This quantifies the # quality of the final object mask. logging.info('Saving summaries.') summ = tf.Summary() summ.value.extend(eval_tracker.get_summaries()) eval_tracker.reset() for s in summaries: summ.value.extend(s.value) sv.summary_computed(sess, summ, step)
def train_ffn(model_cls, **model_kwargs): with tf.Graph().as_default(): # The constructor might define TF ops/placeholders, so it is important # that the FFN is instantiated within the current context. model = model_cls(**model_kwargs) eval_shape_zyx = train_eval_size(model).tolist()[::-1] eval_tracker = EvalTracker(eval_shape_zyx) load_data_ops = define_data_input(model, queue_batch=1) prepare_ffn(model) merge_summaries_op = tf.summary.merge_all() if hvd.rank() == 0: save_flags() summary_writer = None saver = tf.train.Saver(keep_checkpoint_every_n_hours=0.25,max_to_keep=20) scaffold = tf.train.Scaffold(saver=saver) if horovodworks: hooks = [hvd.BroadcastGlobalVariablesHook(0), tf.train.StopAtStepHook(last_step=FLAGS.max_steps),] else: hooks = [tf.train.StopAtStepHook(last_step=FLAGS.max_steps),] config=tf.ConfigProto(log_device_placement=False, allow_soft_placement=True, intra_op_parallelism_threads = FLAGS.num_intra_threads, inter_op_parallelism_threads = FLAGS.num_inter_threads) config.gpu_options.allow_growth = True config.gpu_options.visible_device_list = str(hvd.local_rank()) # Horovod: save checkpoints only on worker 0 to prevent other workers from # corrupting them. checkpoint_dir = FLAGS.train_dir if hvd.rank() == 0 else None with tf.train.MonitoredTrainingSession( master=FLAGS.master, save_summaries_steps=None, save_checkpoint_secs=FLAGS.summary_rate_secs, config=config, checkpoint_dir=checkpoint_dir, hooks=hooks, scaffold=scaffold) as sess: eval_tracker.sess = sess step = int(sess.run(model.global_step)) if hvd.rank() == 0: summary_writer = tf.summary.FileWriterCache.get(FLAGS.train_dir) summary_writer.add_session_log( tf.summary.SessionLog(status=tf.summary.SessionLog.START), step) fov_shifts = list(model.shifts) # x, y, z if FLAGS.shuffle_moves: random.shuffle(fov_shifts) policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts), 'max_pred_moves': max_pred_offsets } batch_it = get_batch(lambda: sess.run(load_data_ops), eval_tracker, model, FLAGS.batch_size, policy_map[FLAGS.fov_policy]) t_last = time.time() while not sess.should_stop() and step < FLAGS.max_steps: # Run summaries periodically. t_curr = time.time() if t_curr - t_last > FLAGS.summary_rate_secs and hvd.rank() == 0: summ_op = merge_summaries_op t_last = t_curr else: summ_op = None seed, patches, labels, weights = next(batch_it) target_lr = get_learning_rate(step,FLAGS.batch_size) updated_seed, step, summ = run_training_step( sess, model, summ_op, feed_dict={ model.loss_weights: weights, model.labels: labels, model.input_patches: patches, model.input_seed: seed, model.learning_rate: target_lr # Wushi: update learning rate }) # Save prediction results in the original seed array so that # they can be used in subsequent steps. mask.update_at(seed, (0, 0, 0), updated_seed) # Record summaries. if summ is not None and hvd.rank() == 0: logging.info('Saving summaries.') summ = tf.Summary.FromString(summ) # Compute a loss over the whole training patch (i.e. more than a # single-step field of view of the network). This quantifies the # quality of the final object mask. summ.value.extend(eval_tracker.get_summaries()) eval_tracker.reset() assert summary_writer is not None summary_writer.add_summary(summ, step) if summary_writer is not None: summary_writer.flush()
def train_ffn(model_cls, **model_kwargs): with tf.Graph().as_default(): model = model_cls(**model_kwargs) # initialize the model eval_shape_zyx = train_eval_size(model).tolist( )[::-1] # size of the subvolume (within which the FOV moves) eval_tracker = EvalTracker( eval_shape_zyx) # computes summary statistics inside EFOV load_data_ops = define_data_input( model, queue_batch=1) # this creates a batch of training subvolumes prepare_ffn(model) # here the tf graph is defined merge_summaries_op = tf.summary.merge_all( ) # merges all summaries defined in the graph. if hvd.rank() == 0: save_flags() var_to_reduce = tf.placeholder(tf.float32) bcast_op = hvd.broadcast_global_variables(0) avg_op = hvd.allreduce(var_to_reduce, average=True) # Start supervisor. sv = tf.train.Supervisor( logdir=(FLAGS.train_dir if hvd.rank() == 0 else None), is_chief=True, saver=(tf.train.Saver(max_to_keep=FLAGS.max_to_keep, keep_checkpoint_every_n_hours=1) if hvd.rank() == 0 else None), save_model_secs=(FLAGS.save_model_secs if hvd.rank() == 0 else 0), summary_op=None, save_summaries_secs=0, # will perform custom summaries instead ) sess = sv.prepare_or_wait_for_session( FLAGS.master, config=tf.ConfigProto( log_device_placement=False, allow_soft_placement=True, intra_op_parallelism_threads=FLAGS.num_intra_threads, inter_op_parallelism_threads=FLAGS.num_inter_threads)) # broadcast initial weights. This ensures that all horovod ranks # start at the same point in parameter space if hvd.rank() == 0: print("broadcasting initial weights") sess.run(bcast_op) eval_tracker.sess = sess #--connect the eval tracker to the session eval_tracker.avg_op = avg_op fov_shifts = list(model.shifts) # x, y, z if FLAGS.shuffle_moves: random.shuffle( fov_shifts ) #--this will shuffle the FOV positions that make up an extended FOV (EFOV) policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts), 'max_pred_moves': max_pred_offsets } # batch iterator for getting the next batch batch_it = get_batch(lambda: sess.run(load_data_ops), eval_tracker, model, FLAGS.batch_size, policy_map[FLAGS.fov_policy]) step = 0 t_last = time.time() if hvd.rank() == 0: timing = [] # list of times for benchmarking steps_since_last_summary = 0 if hvd.rank() == 0: print("starting training") while step < FLAGS.max_steps: time_step_start = time.time() if steps_since_last_summary == FLAGS.summary_every_steps: summ_op = merge_summaries_op steps_since_last_summary = 1 if hvd.rank() == 0: print("step ", step, "is a summary step") else: summ_op = None steps_since_last_summary += 1 # get the next batch - this is reading the data from disk. seed, patches, labels, weights = next(batch_it) summaries = [] scaled_lr = FLAGS.learning_rate if (FLAGS.scaling_rule > 0): #--scale the learning rate linearly (1) or sqrt (2) if FLAGS.scaled_lr == 1: scaled_lr *= hvd.size() elif FLAGS.scaled_lr == 2: scaled_lr *= np.sqrt(hvd.size()) if step < FLAGS.warmup_steps: scaled_lr = FLAGS.learning_rate + (step / float( FLAGS.warmup_steps)) * (scaled_lr - FLAGS.learning_rate) if ( FLAGS.decay_learning_rate_fraction > 0 ): # constantly decay the learning rate using exponential decay scaled_lr *= (FLAGS.decay_learning_rate_fraction)**( step / FLAGS.decay_learning_rate_steps) updated_seed, step, summ, my_loss = run_training_step( # run training step on a SINGLE FOV sess, model, summ_op, feed_dict={ model.loss_weights: weights, model.labels: labels, model.offset_label: 'off', model.input_patches: patches, model.input_seed: seed, model.learning_rate: scaled_lr }) # compute average loss avg_loss = sess.run(avg_op, feed_dict={var_to_reduce: my_loss}) # Save prediction results in the original seed array so that # they can be used in subsequent steps. mask.update_at( seed, (0, 0, 0), updated_seed) # updates the mask inside the subvolume batches if hvd.rank() == 0: this_time = time.time( ) - time_step_start # how long did this step take timing.append(this_time) print("step %i took %.2f seconds" % (step - 1, this_time)) if summ is not None: summaries.append(tf.Summary.FromString( summ)) # this adds the summaries from the single FOV # Compute a loss over the whole training patch (i.e. more than a # single-step field of view of the network). This quantifies the # quality of the final object mask. tp, fp, tn, fn, num_patches = eval_tracker.get_summaries_scalar( ) tp_sum = hvd.size() * sess.run( avg_op, feed_dict={var_to_reduce: float(tp)}) fp_sum = hvd.size() * sess.run( avg_op, feed_dict={var_to_reduce: float(fp)}) tn_sum = hvd.size() * sess.run( avg_op, feed_dict={var_to_reduce: float(tn)}) fn_sum = hvd.size() * sess.run( avg_op, feed_dict={var_to_reduce: float(fn)}) avg_num_patches = sess.run( avg_op, feed_dict={var_to_reduce: float(num_patches)}) accuracy = (tp_sum + tn_sum) / (tp_sum + fp_sum + tn_sum + fn_sum) precision = (tp_sum) / (tp_sum + fp_sum) recall = (tp_sum) / (tp_sum + fn_sum) f1 = 2.0 * precision * recall / (precision + recall) eval_tracker_summaries = ([ tf.Summary.Value(tag='eval/patches', simple_value=avg_num_patches), tf.Summary.Value(tag='eval/accuracy', simple_value=accuracy), tf.Summary.Value(tag='eval/precision', simple_value=precision), tf.Summary.Value(tag='eval/recall', simple_value=recall), tf.Summary.Value(tag='eval/f1', simple_value=f1) ]) if hvd.rank() == 0: logging.info('Saving summaries.') summ = tf.Summary() #initialize tensorflow summary summ.value.extend( eval_tracker_summaries) # add EFOV metrics summ.value.extend(eval_tracker.get_summaries_images() ) # add image summaries for s in summaries: summ.value.extend(s.value) # add FOV metrics # other custom summary items: summ.value.extend([ tf.Summary.Value(tag='avg_pixel_loss', simple_value=avg_loss) ]) #avg pixel loss summ.value.extend([ tf.Summary.Value(tag='learning_rate', simple_value=scaled_lr) ]) #(scaled) learning rate summ.value.extend([ tf.Summary.Value(tag='avg_time_per_step', simple_value=np.mean(timing)) ]) #avg time per step print("avg time per step: ", np.mean(timing), np.std(timing)) print("avg throughput: ", FLAGS.batch_size / np.mean(timing)) timing = [] # reset the timing array sv.summary_computed(sess, summ, step) # reset eval tracker before the next training step. eval_tracker.reset() if hvd.rank() == 0: print("all steps done!") if (FLAGS.do_benchmark_test == 1): print("benchmark result: ") print("steps, ranks, threads, mean, sigma:") string = str(FLAGS.max_steps) + "," + str( FLAGS.batch_size) + "," + str( hvd.size()) + "," + str(FLAGS.nthreads) + "," + str( np.mean(timing)) + "," + str(np.std(timing)) + "\n" print(string) with open(FLAGS.timelog, "a") as myfile: myfile.write(string)
def main(argv): experiment_uid = uid_from_flags(FLAGS) train_dir = os.path.join(FLAGS.train_base_dir, experiment_uid) with tf.Graph().as_default(): with tf.device( tf.train.replica_device_setter(FLAGS.ps_tasks, merge_devices=True)): # The constructor might define TF ops/placeholders, so it is important # that the FFN is instantiated within the current context. # Note: all inputs to ffn.training.model.FFNModel are in format (x, y, z), # except the fov_size, for historical reasons. # If fov_size is specified at command line, it will be stored as a list of # strings; these need to be coerced to integers. model = FFNTracerModel(batch_size=FLAGS.batch_size, adv_args=FLAGS.adv_args, **json.loads(FLAGS.model_args)) eval_shape_zyx = train_eval_size(model).tolist()[::-1] eval_tracker = EvalTracker(eval_shape_zyx) load_data_ops = define_data_input(model, queue_batch=1) prepare_ffn(model) merge_summaries_op = tf.summary.merge_all() # if FLAGS.task == 0: # save_flags() summary_writer = None saver = tf.train.Saver(keep_checkpoint_every_n_hours=0.25) scaffold = tf.train.Scaffold(saver=saver) with tf.train.MonitoredTrainingSession( master=FLAGS.master, is_chief=(FLAGS.task == 0), save_summaries_steps=None, save_checkpoint_secs=300, config=tf.ConfigProto( log_device_placement=False, allow_soft_placement=True, gpu_options=tf.GPUOptions( allow_growth=True, # visible_device_list=FLAGS.visible_gpus )), checkpoint_dir=train_dir, scaffold=scaffold) as sess: eval_tracker.sess = sess step = int(sess.run(model.global_step)) if FLAGS.task > 0: # To avoid early instabilities when using multiple replicas, we use # a launch schedule where new replicas are brought online gradually. logging.info('Delaying replica start.') while step < FLAGS.replica_step_delay * FLAGS.task: time.sleep(5.0) step = int(sess.run(model.global_step)) else: summary_writer = tf.summary.FileWriterCache.get(train_dir) summary_writer.add_session_log( tf.summary.SessionLog( status=tf.summary.SessionLog.START), step) fov_shifts_xyz = list(model.shifts) if FLAGS.shuffle_moves: random.shuffle(fov_shifts_xyz) # Policy_map is a dict mapping a fov_policy to a callable that # generates offsets, given a model and a seed as inputs. Note that # 'fixed' is the policy used for training, and max_pred_moves is the # policy used for inference. policy_map = { 'fixed': partial(fixed_offsets, fov_shifts=fov_shifts_xyz), 'max_pred_moves': max_pred_offsets } # JG: batch_it contains (seed, image, label, weights), where each is of # shape [b, z, y, x, 1] batch_it = get_batch(lambda: sess.run(load_data_ops), eval_tracker, model, FLAGS.batch_size, policy_map[FLAGS.fov_policy]) t_last = time.time() while not sess.should_stop() and step < FLAGS.max_steps: # Run summaries periodically. t_curr = time.time() if t_curr - t_last > FLAGS.summary_rate_secs and FLAGS.task == 0: summ_op = merge_summaries_op t_last = t_curr else: summ_op = None seed, patches, labels, weights = next(batch_it) # JG: weights, labels, patches, and seed all have # shape [b, z, y, x, 1] at this stage. feed_dict = { model.loss_weights: weights, model.labels: labels, model.input_patches: patches, model.input_seed: seed, } if is_adversary_update_step(step, model): # Update the adversary step, summ = run_adversary_training_step( sess, model, summ_op, feed_dict=feed_dict) else: # Update the FFN model updated_seed, step, summ = run_training_step( sess, model, summ_op, feed_dict=feed_dict) # Save prediction results in the original seed array so that # they can be used in subsequent steps. mask.update_at(seed, (0, 0, 0), updated_seed) # Record summaries. if summ is not None: logging.info('Saving summaries.') summ = tf.Summary.FromString(summ) # Compute a loss over the whole training patch (i.e. more than a # single-step field of view of the network). This quantifies the # quality of the final object mask. summ.value.extend(eval_tracker.get_summaries()) eval_tracker.reset() assert summary_writer is not None summary_writer.add_summary(summ, step) if summary_writer is not None: summary_writer.flush()