def main(_): if not FLAGS.dataset_dir: raise ValueError('You must supply the dataset directory with --dataset_dir') tf.logging.set_verbosity(tf.logging.DEBUG) with tf.Graph().as_default(): # Config model_deploy. Keep TF Slim Models structure. # Useful if want to need multiple GPUs and/or servers in the future. deploy_config = model_deploy.DeploymentConfig( num_clones=FLAGS.num_clones, clone_on_cpu=FLAGS.clone_on_cpu, replica_id=0, num_replicas=1, num_ps_tasks=0) # Create global_step. with tf.device(deploy_config.variables_device()): global_step = slim.create_global_step() # Select the dataset. dataset = dataset_factory.get_dataset( FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.dataset_dir) # Get the SSD network and its anchors. ssd_class = nets_factory.get_network(FLAGS.model_name) ssd_params = ssd_class.default_params._replace(num_classes=FLAGS.num_classes) ssd_net = ssd_class(ssd_params) ssd_shape = ssd_net.params.img_shape ssd_anchors = ssd_net.anchors(ssd_shape) # Select the preprocessing function. preprocessing_name = FLAGS.preprocessing_name or FLAGS.model_name image_preprocessing_fn = preprocessing_factory.get_preprocessing( preprocessing_name, is_training=True) tf_utils.print_configuration(FLAGS.__flags, ssd_params, dataset.data_sources, FLAGS.train_dir) # =================================================================== # # Create a dataset provider and batches. # =================================================================== # with tf.device(deploy_config.inputs_device()): with tf.name_scope(FLAGS.dataset_name + '_data_provider'): provider = slim.dataset_data_provider.DatasetDataProvider( dataset, num_readers=FLAGS.num_readers, common_queue_capacity=20 * FLAGS.batch_size, common_queue_min=10 * FLAGS.batch_size, shuffle=True) # Get for SSD network: image, labels, bboxes. [image, shape, glabels, gbboxes] = provider.get(['image', 'shape', 'object/label', 'object/bbox']) # Pre-processing image, labels and bboxes. image, glabels, gbboxes = \ image_preprocessing_fn(image, glabels, gbboxes, out_shape=ssd_shape, data_format=DATA_FORMAT) # Encode groundtruth labels and bboxes. gclasses, glocalisations, gscores = \ ssd_net.bboxes_encode(glabels, gbboxes, ssd_anchors) batch_shape = [1] + [len(ssd_anchors)] * 3 # Training batches and queue. r = tf.train.batch( tf_utils.reshape_list([image, gclasses, glocalisations, gscores]), batch_size=FLAGS.batch_size, num_threads=FLAGS.num_preprocessing_threads, capacity=5 * FLAGS.batch_size) b_image, b_gclasses, b_glocalisations, b_gscores = \ tf_utils.reshape_list(r, batch_shape) # Intermediate queueing: unique batch computation pipeline for all # GPUs running the training. batch_queue = slim.prefetch_queue.prefetch_queue( tf_utils.reshape_list([b_image, b_gclasses, b_glocalisations, b_gscores]), capacity=2 * deploy_config.num_clones) # =================================================================== # # Define the model running on every GPU. # =================================================================== # def clone_fn(batch_queue): """Allows data parallelism by creating multiple clones of network_fn.""" # Dequeue batch. b_image, b_gclasses, b_glocalisations, b_gscores = \ tf_utils.reshape_list(batch_queue.dequeue(), batch_shape) # Construct SSD network. arg_scope = ssd_net.arg_scope(weight_decay=FLAGS.weight_decay, data_format=DATA_FORMAT) with slim.arg_scope(arg_scope): predictions, localisations, logits, end_points = \ ssd_net.net(b_image, is_training=True) # Add loss function. ssd_net.losses(logits, localisations, b_gclasses, b_glocalisations, b_gscores, match_threshold=FLAGS.match_threshold, negative_ratio=FLAGS.negative_ratio, alpha=FLAGS.loss_alpha, label_smoothing=FLAGS.label_smoothing) return end_points # Gather initial summaries. summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES)) # =================================================================== # # Add summaries from first clone. # =================================================================== # clones = model_deploy.create_clones(deploy_config, clone_fn, [batch_queue]) first_clone_scope = deploy_config.clone_scope(0) # Gather update_ops from the first clone. These contain, for example, # the updates for the batch_norm variables created by network_fn. update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, first_clone_scope) # Add summaries for end_points. end_points = clones[0].outputs for end_point in end_points: x = end_points[end_point] summaries.add(tf.summary.histogram('activations/' + end_point, x)) summaries.add(tf.summary.scalar('sparsity/' + end_point, tf.nn.zero_fraction(x))) # Add summaries for losses and extra losses. for loss in tf.get_collection(tf.GraphKeys.LOSSES, first_clone_scope): summaries.add(tf.summary.scalar(loss.op.name, loss)) for loss in tf.get_collection('EXTRA_LOSSES', first_clone_scope): summaries.add(tf.summary.scalar(loss.op.name, loss)) # Add summaries for variables. for variable in slim.get_model_variables(): summaries.add(tf.summary.histogram(variable.op.name, variable)) # =================================================================== # # Configure the moving averages. # =================================================================== # if FLAGS.moving_average_decay: moving_average_variables = slim.get_model_variables() variable_averages = tf.train.ExponentialMovingAverage( FLAGS.moving_average_decay, global_step) else: moving_average_variables, variable_averages = None, None # =================================================================== # # Configure the optimization procedure. # =================================================================== # with tf.device(deploy_config.optimizer_device()): learning_rate = tf_utils.configure_learning_rate(FLAGS, dataset.num_samples, global_step) optimizer = tf_utils.configure_optimizer(FLAGS, learning_rate) summaries.add(tf.summary.scalar('learning_rate', learning_rate)) if FLAGS.moving_average_decay: # Update ops executed locally by trainer. update_ops.append(variable_averages.apply(moving_average_variables)) # Variables to train. variables_to_train = tf_utils.get_variables_to_train(FLAGS) # and returns a train_tensor and summary_op total_loss, clones_gradients = model_deploy.optimize_clones( clones, optimizer, var_list=variables_to_train) # Add total_loss to summary. summaries.add(tf.summary.scalar('total_loss', total_loss)) # Create gradient updates. grad_updates = optimizer.apply_gradients(clones_gradients, global_step=global_step) update_ops.append(grad_updates) update_op = tf.group(*update_ops) train_tensor = control_flow_ops.with_dependencies([update_op], total_loss, name='train_op') # Add the summaries from the first clone. These contain the summaries summaries |= set(tf.get_collection(tf.GraphKeys.SUMMARIES, first_clone_scope)) # Merge all summaries together. summary_op = tf.summary.merge(list(summaries), name='summary_op') # =================================================================== # # Kicks off the training. # =================================================================== # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction) config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options) saver = tf.train.Saver(max_to_keep=5, keep_checkpoint_every_n_hours=1.0, write_version=2, pad_step_number=False) slim.learning.train( train_tensor, logdir=FLAGS.train_dir, master='', is_chief=True, init_fn=tf_utils.get_init_fn(FLAGS), summary_op=summary_op, number_of_steps=FLAGS.max_number_of_steps, log_every_n_steps=FLAGS.log_every_n_steps, save_summaries_secs=FLAGS.save_summaries_secs, saver=saver, save_interval_secs=FLAGS.save_interval_secs, session_config=config, sync_optimizer=None)
def main(_): if FLAGS.train_on_cpu: os.environ["CUDA_VISIBLE_DEVICES"] = "-1" else: os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu_device if not FLAGS.dataset_dir: raise ValueError( "You must supply the dataset directory with --dataset-dir.") tf.logging.set_verbosity(tf.logging.DEBUG) g = tf.Graph() with g.as_default(): # select the dataset dataset = dataset_factory.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.dataset_dir) # create global step, used for optimizer moving average decay with tf.device("/cpu:0"): global_step = tf.train.create_global_step() # pdb.set_trace() # get the ssd network and its anchors ssd_cls = ssd.SSDnet ssd_params = ssd_cls.default_params._replace( num_classes=FLAGS.num_classes) ssd_net = ssd_cls(ssd_params) image_size = ssd_net.params.img_shape ssd_anchors = ssd_net.anchors(img_shape=image_size) # select the preprocessing function preprocessing_name = FLAGS.preprocessing_name image_preprocessing_fn = preprocessing_factory.get_preprocessing( preprocessing_name, is_training=True) tf_utils.print_configuration(FLAGS.__flags, ssd_params, dataset.data_sources, FLAGS.train_dir) # create a dataset provider and batches. with tf.device("/cpu:0"): with tf.name_scope(FLAGS.dataset_name + "_data_provider"): provider = slim.dataset_data_provider.DatasetDataProvider( dataset, num_readers=FLAGS.num_readers, common_queue_capacity=20 * FLAGS.batch_size, common_queue_min=10 * FLAGS.batch_size, shuffle=True) # get for ssd network: image,labels,bboxes [image, shape, glabels, gbboxes] = provider.get( ["image", "shape", "object/label", "object/bbox"]) # pdb.set_trace() # preprocessing image,glabels,gbboxes = \ image_preprocessing_fn(image, glabels,gbboxes, out_shape=image_size, data_format="NHWC") # encode groundtruth labels and bboxes gclasses,glocalisations,gscores= \ ssd_net.bboxes_encode(glabels,gbboxes,ssd_anchors) batch_shape = [1] + [len(ssd_anchors)] * 3 # training batches and queue r = tf.train.batch(tf_utils.reshape_list( [image, gclasses, glocalisations, gscores]), batch_size=FLAGS.batch_size, num_threads=FLAGS.num_preprocessing_threads, capacity=5 * FLAGS.batch_size) b_image,b_gclasses,b_glocalisations,b_gscores = \ tf_utils.reshape_list(r,batch_shape) # prefetch queue batch_queue = slim.prefetch_queue.prefetch_queue( tf_utils.reshape_list( [b_image, b_gclasses, b_glocalisations, b_gscores]), capacity=8) # dequeue batch b_image, b_gclasses, b_glocalisations, b_gscores = \ tf_utils.reshape_list(batch_queue.dequeue(), batch_shape) # gather initial summaries summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES)) arg_scope = ssd_net.arg_scope(weight_decay=FLAGS.weight_decay) with slim.arg_scope(arg_scope): predictions,localisations,logits,end_points,mobilenet_var_list = \ ssd_net.net(b_image,is_training=True) # add loss function ssd_net.losses(logits, localisations, b_gclasses, b_glocalisations, b_gscores, match_threshold=FLAGS.match_threshold, negative_ratio=FLAGS.negative_ratio, alpha=FLAGS.loss_alpha, label_smoothing=FLAGS.label_smoothing) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) # add summaries for end_points for end_point in end_points: x = end_points[end_point] summaries.add(tf.summary.histogram("activations/" + end_point, x)) summaries.add( tf.summary.scalar("sparsity/" + end_point, tf.nn.zero_fraction(x))) # add summaries for losses and extra losses for loss in tf.get_collection(tf.GraphKeys.LOSSES): summaries.add(tf.summary.scalar(loss.op.name, loss)) for loss in tf.get_collection("EXTRA_LOSSES"): summaries.add(tf.summary.scalar(loss.op.name, loss)) # add summaries for variables for var in slim.get_model_variables(): summaries.add(tf.summary.histogram(var.op.name, var)) # configure the moving averages if FLAGS.moving_average_decay: # use moving average decay on weights variables moving_average_variables = slim.get_model_variables() variable_averages = tf.train.ExponentialMovingAverage( FLAGS.moving_average_decay, global_step) else: moving_average_variables, variable_averages = None, None # configure the optimization procedure with tf.device("/cpu:0"): learning_rate = tf_utils.configure_learning_rate( FLAGS, dataset.num_samples, global_step) optimizer = tf_utils.configure_optimizer(FLAGS, learning_rate) summaries.add(tf.summary.scalar("learning_rate", learning_rate)) if FLAGS.moving_average_decay: # update ops executed by trainer update_ops.append( variable_averages.apply(moving_average_variables)) # get variables to train variables_to_train = tf_utils.get_variables_to_train(FLAGS) # return a train tensor and summary op total_losses = tf.get_collection(tf.GraphKeys.LOSSES) total_loss = tf.add_n(total_losses, name="total_loss") summaries.add(tf.summary.scalar("total_loss", total_loss)) # create gradient updates grads = optimizer.compute_gradients(total_loss, var_list=variables_to_train) grad_updates = optimizer.apply_gradients(grads, global_step=global_step) update_ops.append(grad_updates) # create train op update_op = tf.group(*update_ops) train_tensor = control_flow_ops.with_dependencies([update_op], total_loss, name="train_op") # merge all summaries together summary_op = tf.summary.merge(list(summaries), name="summary_op") # start training gpu_options = tf.GPUOptions( per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction, allow_growth=FLAGS.allow_growth) config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options) saver = tf.train.Saver(max_to_keep=2, keep_checkpoint_every_n_hours=1.0, write_version=2, pad_step_number=False) # create initial assignment op init_assign_op, init_feed_dict = slim.assign_from_checkpoint( FLAGS.checkpoint_path, mobilenet_var_list, ignore_missing_vars=FLAGS.ignore_missing_vars) # create an initial assignment function for k, v in init_feed_dict.items(): if "global_step" in k.name: g_step = k init_feed_dict[g_step] = 0 # change the global_step to zero. init_fn = lambda sess: sess.run(init_assign_op, init_feed_dict) # run training slim.learning.train( train_tensor, logdir=FLAGS.train_dir, init_fn=init_fn, summary_op=summary_op, number_of_steps=FLAGS.max_number_of_steps, save_summaries_secs=FLAGS.save_summaries_secs, save_interval_secs=FLAGS.save_interval_secs, session_config=config, saver=saver, )
def main(_): if not FLAGS.dataset_dir: raise ValueError( 'You must supply the dataset directory with --dataset_dir') tf.logging.set_verbosity(tf.logging.INFO) with tf.Graph().as_default(): tf_global_step = slim.get_or_create_global_step() # =================================================================== # # Dataset + SSD model + Pre-processing # =================================================================== # dataset = dataset_factory.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.dataset_dir) # Get the SSD network and its anchors. ssd_class = nets_factory.get_network(FLAGS.model_name) ssd_params = ssd_class.default_params._replace( num_classes=FLAGS.num_classes, no_annotation_label=FLAGS.num_classes) ssd_net = ssd_class(ssd_params) # Evaluation shape and associated anchors: eval_image_size ssd_shape = ssd_net.params.img_shape ssd_anchors = ssd_net.anchors(ssd_shape) # Select the preprocessing function. preprocessing_name = FLAGS.preprocessing_name or FLAGS.model_name image_preprocessing_fn = preprocessing_factory.get_preprocessing( preprocessing_name, is_training=False) tf_utils.print_configuration(FLAGS.__flags, ssd_params, dataset.data_sources, FLAGS.eval_dir) # =================================================================== # # Create a dataset provider and batches. # =================================================================== # with tf.device('/cpu:0'): with tf.name_scope(FLAGS.dataset_name + '_data_provider'): provider = slim.dataset_data_provider.DatasetDataProvider( dataset, common_queue_capacity=2 * FLAGS.batch_size, common_queue_min=FLAGS.batch_size, shuffle=False) # Get for SSD network: image, labels, bboxes. [image, shape, glabels, gbboxes] = provider.get( ['image', 'shape', 'object/label', 'object/bbox']) if FLAGS.remove_difficult: [gdifficults] = provider.get(['object/difficult']) else: gdifficults = tf.zeros(tf.shape(glabels), dtype=tf.int64) # Pre-processing image, labels and bboxes. image, glabels, gbboxes, gbbox_img = \ image_preprocessing_fn(image, glabels, gbboxes, out_shape=ssd_shape, data_format=DATA_FORMAT, resize=FLAGS.eval_resize, difficults=None) # Encode groundtruth labels and bboxes. gclasses, glocalisations, gscores = \ ssd_net.bboxes_encode(glabels, gbboxes, ssd_anchors) batch_shape = [1] * 5 + [len(ssd_anchors)] * 3 # Evaluation batch. r = tf.train.batch(tf_utils.reshape_list([ image, glabels, gbboxes, gdifficults, gbbox_img, gclasses, glocalisations, gscores ]), batch_size=FLAGS.batch_size, num_threads=FLAGS.num_preprocessing_threads, capacity=5 * FLAGS.batch_size, dynamic_pad=True) (b_image, b_glabels, b_gbboxes, b_gdifficults, b_gbbox_img, b_gclasses, b_glocalisations, b_gscores) = tf_utils.reshape_list(r, batch_shape) # =================================================================== # # SSD Network + Ouputs decoding. # =================================================================== # dict_metrics = {} arg_scope = ssd_net.arg_scope(data_format=DATA_FORMAT) with slim.arg_scope(arg_scope): predictions, localisations, logits, end_points = \ ssd_net.net(b_image, is_training=False) # Add losses functions. ssd_net.losses(logits, localisations, b_gclasses, b_glocalisations, b_gscores) # Performing post-processing on CPU: loop-intensive, usually more efficient. with tf.device('/device:CPU:0'): # Detected objects from SSD output. localisations = ssd_net.bboxes_decode(localisations, ssd_anchors) rscores, rbboxes = \ ssd_net.detected_bboxes(predictions, localisations, select_threshold=FLAGS.select_threshold, nms_threshold=FLAGS.nms_threshold, clipping_bbox=None, top_k=FLAGS.select_top_k, keep_top_k=FLAGS.keep_top_k) # Compute TP and FP statistics. num_gbboxes, tp, fp, rscores = \ tfe.bboxes_matching_batch(rscores.keys(), rscores, rbboxes, b_glabels, b_gbboxes, b_gdifficults, matching_threshold=FLAGS.matching_threshold) # Variables to restore: moving avg. or normal weights. if FLAGS.moving_average_decay: variable_averages = tf.train.ExponentialMovingAverage( FLAGS.moving_average_decay, tf_global_step) variables_to_restore = variable_averages.variables_to_restore( slim.get_model_variables()) variables_to_restore[tf_global_step.op.name] = tf_global_step else: variables_to_restore = slim.get_variables_to_restore() # =================================================================== # # Evaluation metrics. # =================================================================== # with tf.device('/device:CPU:0'): dict_metrics = {} # First add all losses. for loss in tf.get_collection(tf.GraphKeys.LOSSES): dict_metrics[loss.op.name] = slim.metrics.streaming_mean(loss) # Extra losses as well. for loss in tf.get_collection('EXTRA_LOSSES'): dict_metrics[loss.op.name] = slim.metrics.streaming_mean(loss) # Add metrics to summaries and Print on screen. for name, metric in dict_metrics.items(): # summary_name = 'eval/%s' % name summary_name = name op = tf.summary.scalar(summary_name, metric[0], collections=[]) # op = tf.Print(op, [metric[0]], summary_name) tf.add_to_collection(tf.GraphKeys.SUMMARIES, op) # FP and TP metrics. tp_fp_metric = tfe.streaming_tp_fp_arrays(num_gbboxes, tp, fp, rscores) for c in tp_fp_metric[0].keys(): dict_metrics['tp_fp_%s' % c] = (tp_fp_metric[0][c], tp_fp_metric[1][c]) # Add to summaries precision/recall values. aps_voc07 = {} aps_voc12 = {} for c in tp_fp_metric[0].keys(): # Precison and recall values. prec, rec = tfe.precision_recall(*tp_fp_metric[0][c]) # Average precision VOC07. v = tfe.average_precision_voc07(prec, rec) summary_name = 'AP_VOC07/%s' % c op = tf.summary.scalar(summary_name, v, collections=[]) # op = tf.Print(op, [v], summary_name) tf.add_to_collection(tf.GraphKeys.SUMMARIES, op) aps_voc07[c] = v # Average precision VOC12. v = tfe.average_precision_voc12(prec, rec) summary_name = 'AP_VOC12/%s' % c op = tf.summary.scalar(summary_name, v, collections=[]) # op = tf.Print(op, [v], summary_name) tf.add_to_collection(tf.GraphKeys.SUMMARIES, op) aps_voc12[c] = v # Mean average precision VOC07. summary_name = 'AP_VOC07/mAP' mAP = tf.add_n(list(aps_voc07.values())) / len(aps_voc07) op = tf.summary.scalar(summary_name, mAP, collections=[]) op = tf.Print(op, [mAP], summary_name) tf.add_to_collection(tf.GraphKeys.SUMMARIES, op) # Mean average precision VOC12. summary_name = 'AP_VOC12/mAP' mAP = tf.add_n(list(aps_voc12.values())) / len(aps_voc12) op = tf.summary.scalar(summary_name, mAP, collections=[]) op = tf.Print(op, [mAP], summary_name) tf.add_to_collection(tf.GraphKeys.SUMMARIES, op) # for i, v in enumerate(l_precisions): # summary_name = 'eval/precision_at_recall_%.2f' % LIST_RECALLS[i] # op = tf.summary.scalar(summary_name, v, collections=[]) # op = tf.Print(op, [v], summary_name) # tf.add_to_collection(tf.GraphKeys.SUMMARIES, op) # Split into values and updates ops. names_to_values, names_to_updates = slim.metrics.aggregate_metric_map( dict_metrics) # =================================================================== # # Evaluation loop. # =================================================================== # gpu_options = tf.GPUOptions( per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction) config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options) # config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1 # Number of batches... if FLAGS.max_num_batches: num_batches = FLAGS.max_num_batches else: num_batches = math.ceil(dataset.num_samples / float(FLAGS.batch_size)) if not FLAGS.wait_for_checkpoints: if tf.gfile.IsDirectory(FLAGS.checkpoint_path): checkpoint_path = tf.train.latest_checkpoint( FLAGS.checkpoint_path) else: checkpoint_path = FLAGS.checkpoint_path tf.logging.info('Evaluating %s' % checkpoint_path) # Standard evaluation loop. start = time.time() slim.evaluation.evaluate_once( master=FLAGS.master, checkpoint_path=checkpoint_path, logdir=FLAGS.eval_dir, num_evals=num_batches, eval_op=list(names_to_updates.values()), variables_to_restore=variables_to_restore, session_config=config) # Log time spent. elapsed = time.time() elapsed = elapsed - start print('Time spent : %.3f seconds.' % elapsed) print('Time spent per BATCH: %.3f seconds.' % (elapsed / num_batches)) else: checkpoint_path = FLAGS.checkpoint_path tf.logging.info('Evaluating %s' % checkpoint_path) # Waiting loop. slim.evaluation.evaluation_loop( master=FLAGS.master, checkpoint_dir=checkpoint_path, logdir=FLAGS.eval_dir, num_evals=num_batches, eval_op=list(names_to_updates.values()), variables_to_restore=variables_to_restore, eval_interval_secs=60, max_number_of_evaluations=np.inf, session_config=config, timeout=None)
def main(_): if not FLAGS.dataset_dir: raise ValueError( 'You must supply the dataset directory with --dataset_dir') tf.logging.set_verbosity(tf.logging.DEBUG) with tf.Graph().as_default(): # Config model_deploy. Keep TF Slim Models structure. # Useful if want to need multiple GPUs and/or servers in the future. deploy_config = model_deploy.DeploymentConfig( num_clones=FLAGS.num_clones, clone_on_cpu=FLAGS.clone_on_cpu, replica_id=0, num_replicas=1, num_ps_tasks=0) # Create global_step. with tf.device(deploy_config.variables_device()): global_step = slim.create_global_step() # Select the dataset. dataset = dataset_factory.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.dataset_dir) # Get the SSD network and its anchors. ssd_class = nets_factory.get_network(FLAGS.model_name) ssd_params = ssd_class.default_params._replace( num_classes=FLAGS.num_classes) ssd_net = ssd_class(ssd_params) ssd_shape = ssd_net.params.img_shape ssd_anchors = ssd_net.anchors(ssd_shape) # Select the preprocessing function. preprocessing_name = FLAGS.preprocessing_name or FLAGS.model_name image_preprocessing_fn = preprocessing_factory.get_preprocessing( preprocessing_name, is_training=True) tf_utils.print_configuration(FLAGS.__flags, ssd_params, dataset.data_sources, FLAGS.train_dir) # =================================================================== # # Create a dataset provider and batches. # =================================================================== # with tf.device(deploy_config.inputs_device()): with tf.name_scope(FLAGS.dataset_name + '_data_provider'): provider = slim.dataset_data_provider.DatasetDataProvider( dataset, num_readers=FLAGS.num_readers, common_queue_capacity=20 * FLAGS.batch_size, common_queue_min=10 * FLAGS.batch_size, shuffle=True) # Get for SSD network: image, labels, bboxes. [image, shape, glabels, gbboxes] = provider.get( ['image', 'shape', 'object/label', 'object/bbox']) # Pre-processing image, labels and bboxes. image, glabels, gbboxes = \ image_preprocessing_fn(image, glabels, gbboxes, out_shape=ssd_shape, data_format=DATA_FORMAT) # Encode groundtruth labels and bboxes. gclasses, glocalisations, gscores = \ ssd_net.bboxes_encode(glabels, gbboxes, ssd_anchors) batch_shape = [1] + [len(ssd_anchors)] * 3 # Training batches and queue. r = tf.train.batch(tf_utils.reshape_list( [image, gclasses, glocalisations, gscores]), batch_size=FLAGS.batch_size, num_threads=FLAGS.num_preprocessing_threads, capacity=5 * FLAGS.batch_size) b_image, b_gclasses, b_glocalisations, b_gscores = \ tf_utils.reshape_list(r, batch_shape) # Intermediate queueing: unique batch computation pipeline for all # GPUs running the training. batch_queue = slim.prefetch_queue.prefetch_queue( tf_utils.reshape_list( [b_image, b_gclasses, b_glocalisations, b_gscores]), capacity=2 * deploy_config.num_clones) # =================================================================== # # Define the model running on every GPU. # =================================================================== # def clone_fn(batch_queue): """Allows data parallelism by creating multiple clones of network_fn.""" # Dequeue batch. b_image, b_gclasses, b_glocalisations, b_gscores = \ tf_utils.reshape_list(batch_queue.dequeue(), batch_shape) # Construct SSD network. arg_scope = ssd_net.arg_scope(weight_decay=FLAGS.weight_decay, data_format=DATA_FORMAT) with slim.arg_scope(arg_scope): predictions, localisations, logits, end_points = \ ssd_net.net(b_image, is_training=True) # Add loss function. ssd_net.losses(logits, localisations, b_gclasses, b_glocalisations, b_gscores, match_threshold=FLAGS.match_threshold, negative_ratio=FLAGS.negative_ratio, alpha=FLAGS.loss_alpha, label_smoothing=FLAGS.label_smoothing) return end_points # Gather initial summaries. summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES)) # =================================================================== # # Add summaries from first clone. # =================================================================== # clones = model_deploy.create_clones(deploy_config, clone_fn, [batch_queue]) first_clone_scope = deploy_config.clone_scope(0) # Gather update_ops from the first clone. These contain, for example, # the updates for the batch_norm variables created by network_fn. update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, first_clone_scope) # Add summaries for end_points. end_points = clones[0].outputs for end_point in end_points: x = end_points[end_point] summaries.add(tf.summary.histogram('activations/' + end_point, x)) summaries.add( tf.summary.scalar('sparsity/' + end_point, tf.nn.zero_fraction(x))) # Add summaries for losses and extra losses. for loss in tf.get_collection(tf.GraphKeys.LOSSES, first_clone_scope): summaries.add(tf.summary.scalar(loss.op.name, loss)) for loss in tf.get_collection('EXTRA_LOSSES', first_clone_scope): summaries.add(tf.summary.scalar(loss.op.name, loss)) # Add summaries for variables. for variable in slim.get_model_variables(): summaries.add(tf.summary.histogram(variable.op.name, variable)) # =================================================================== # # Configure the moving averages. # =================================================================== # if FLAGS.moving_average_decay: moving_average_variables = slim.get_model_variables() variable_averages = tf.train.ExponentialMovingAverage( FLAGS.moving_average_decay, global_step) else: moving_average_variables, variable_averages = None, None # =================================================================== # # Configure the optimization procedure. # =================================================================== # with tf.device(deploy_config.optimizer_device()): learning_rate = tf_utils.configure_learning_rate( FLAGS, dataset.num_samples, global_step) optimizer = tf_utils.configure_optimizer(FLAGS, learning_rate) summaries.add(tf.summary.scalar('learning_rate', learning_rate)) if FLAGS.moving_average_decay: # Update ops executed locally by trainer. update_ops.append( variable_averages.apply(moving_average_variables)) # Variables to train. variables_to_train = tf_utils.get_variables_to_train(FLAGS) # and returns a train_tensor and summary_op total_loss, clones_gradients = model_deploy.optimize_clones( clones, optimizer, var_list=variables_to_train) # Add total_loss to summary. summaries.add(tf.summary.scalar('total_loss', total_loss)) # Create gradient updates. grad_updates = optimizer.apply_gradients(clones_gradients, global_step=global_step) update_ops.append(grad_updates) update_op = tf.group(*update_ops) train_tensor = control_flow_ops.with_dependencies([update_op], total_loss, name='train_op') # Add the summaries from the first clone. These contain the summaries summaries |= set( tf.get_collection(tf.GraphKeys.SUMMARIES, first_clone_scope)) # Merge all summaries together. summary_op = tf.summary.merge(list(summaries), name='summary_op') # =================================================================== # # Kicks off the training. # =================================================================== # gpu_options = tf.GPUOptions( per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction) config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options) saver = tf.train.Saver(max_to_keep=5, keep_checkpoint_every_n_hours=1.0, write_version=2, pad_step_number=False) slim.learning.train(train_tensor, logdir=FLAGS.train_dir, master='', is_chief=True, init_fn=tf_utils.get_init_fn(FLAGS), summary_op=summary_op, number_of_steps=FLAGS.max_number_of_steps, log_every_n_steps=FLAGS.log_every_n_steps, save_summaries_secs=FLAGS.save_summaries_secs, saver=saver, save_interval_secs=FLAGS.save_interval_secs, session_config=config, sync_optimizer=None)
def train_input_fn(): # Select the dataset. dataset = dataset_factory.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.data_dir) tf_utils.print_configuration(FLAGS.__flags, ron_params, dataset.data_sources, FLAGS.model_dir) # =================================================================== # # Create a dataset provider and batches. # =================================================================== # with tf.name_scope(FLAGS.dataset_name + '_data_provider'): provider = slim.dataset_data_provider.DatasetDataProvider( dataset, num_readers=FLAGS.num_readers, common_queue_capacity=120 * FLAGS.batch_size, common_queue_min=80 * FLAGS.batch_size, shuffle=True) # Get for RON network: image, labels, bboxes. # (ymin, xmin, ymax, xmax) fro gbboxes [image, shape, glabels, gbboxes, isdifficult] = provider.get([ 'image', 'shape', 'object/label', 'object/bbox', 'object/difficult' ]) isdifficult_mask = tf.cond( tf.reduce_sum( tf.cast( tf.logical_not( tf.equal(tf.ones_like(isdifficult), isdifficult)), tf.float32)) < 1., lambda: tf.one_hot(0, tf.shape(isdifficult)[0], on_value=True, off_value=False, dtype=tf.bool), lambda: isdifficult < tf.ones_like(isdifficult)) glabels = tf.boolean_mask(glabels, isdifficult_mask) gbboxes = tf.boolean_mask(gbboxes, isdifficult_mask) # Select the preprocessing function. preprocessing_name = FLAGS.preprocessing_name or FLAGS.model_name image_preprocessing_fn = preprocessing_factory.get_preprocessing( preprocessing_name, is_training=True) # Pre-processing image, labels and bboxes. image, glabels, gbboxes = image_preprocessing_fn( image, glabels, gbboxes, out_shape=ron_shape, data_format=DATA_FORMAT) # Encode groundtruth labels and bboxes. # glocalisations is our regression object # gclasses is the ground_trutuh label # gscores is the the jaccard score with ground_truth gclasses, glocalisations, gscores = ron_net.bboxes_encode( glabels, gbboxes, ron_anchors, positive_threshold=FLAGS.match_threshold, ignore_threshold=FLAGS.neg_threshold) # each size of the batch elements # include one image, three others(gclasses, glocalisations, gscores) batch_shape = [1] + [len(ron_anchors)] * 3 # Training batches and queue. r = tf.train.batch(tf_utils.reshape_list( [image, gclasses, glocalisations, gscores]), batch_size=FLAGS.batch_size, num_threads=FLAGS.num_preprocessing_threads, capacity=120 * FLAGS.batch_size, shared_name=None) b_image, b_gclasses, b_glocalisations, b_gscores = tf_utils.reshape_list( r, batch_shape) return b_image, { 'b_gclasses': b_gclasses, 'b_glocalisations': b_glocalisations, 'b_gscores': b_gscores }