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
}
