def test_random_adjust_brightness(self):
preprocessor_text_proto = """
random_adjust_brightness {
max_delta: 0.15
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.random_adjust_brightness)
self.assert_dictionary_close(args, {
'max_delta': 0.15
})
def test_random_rgb_to_gray(self):
preprocessor_text_proto = """
random_rgb_to_gray {
probability: 0.15
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.random_rgb_to_gray)
self.assert_dictionary_close(args, {
'probability': 0.15
})
def test_build_resize_image_random_method(self):
preprocessor_text_proto = """
resize_image_random_method {
target_height: 384
target_width: 384
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.resize_image_random_method)
self.assert_dictionary_close(args, {
'target_size': [384, 384]
})
def test_random_pixel_value_scale(self):
preprocessor_text_proto = """
random_pixel_value_scale {
minval: 0.85
maxval: 1.25
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.random_pixel_value_scale)
self.assert_dictionary_close(args, {
'minval': 0.85,
'maxval': 1.25
})
def test_random_adjust_contrast(self):
preprocessor_text_proto = """
random_adjust_contrast {
min_delta: 0.7
max_delta: 1.3
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.random_adjust_contrast)
self.assert_dictionary_close(args, {
'min_delta': 0.7,
'max_delta': 1.3
})
def test_build_resize_image(self):
preprocessor_text_proto = """
resize_image {
target_height: 384
target_width: 384
method: BICUBIC
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.resize_image)
self.assertEqual(args, {
'target_size': [384, 384],
'method': tf.image.ResizeMethod.BICUBIC
})
def test_normalize_image(self):
preprocessor_text_proto = """
normalize_image {
original_minval: 0.0
original_maxval: 255.0
target_minval: 0.0
target_maxval: 1.0
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.normalize_image)
self.assert_dictionary_close(args, {
'original_minval': 0.0,
'original_maxval': 255.0,
'target_minval': 0.0,
'target_maxval': 1.0
})
def test_string_filtering_2(self):
preprocessor_text_proto = """
string_filtering {
lower_case: false
include_charset {
built_in_set: ALLCASES
}
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.string_filtering)
self.assert_dictionary_close(args, {
'lower_case': False,
'include_charset': "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
})
test_input_strings = [t.encode('utf-8') for t in ['abc', 'abcde', '!=ABC DE~']]
expected_output_string = [t.encode('utf-8') for t in ['abc', 'abcde', 'ABCDE']]
test_processed_strings = [function(t, **args) for t in test_input_strings]
with self.test_session() as sess:
outputs = sess.run(test_processed_strings)
self.assertAllEqual(outputs, expected_output_string)
def test_string_filtering(self):
preprocessor_text_proto = """
string_filtering {
lower_case: true
include_charset {
text_string: "abc"
}
}
"""
preprocessor_proto = preprocessor_pb2.PreprocessingStep()
text_format.Merge(preprocessor_text_proto, preprocessor_proto)
function, args = preprocessor_builder.build(preprocessor_proto)
self.assertEqual(function, preprocessor.string_filtering)
self.assert_dictionary_close(args, {
'lower_case': True,
'include_charset': "abc"
})
test_input_strings = [t.encode('utf-8') for t in ['abc', 'abcde', 'ABCDE']]
expected_output_string = [t.encode('utf-8') for t in ['abc', 'abc', 'abc']]
test_processed_strings = [function(t, **args) for t in test_input_strings]
with self.test_session() as sess:
outputs = sess.run(test_processed_strings)
self.assertAllEqual(outputs, expected_output_string)
def train(create_tensor_dict_fn_list, create_model_fn, train_config, master,
task, num_clones, worker_replicas, clone_on_cpu, ps_tasks,
worker_job_name, is_chief, train_dir):
"""Training function for models.
Args:
create_tensor_dict_fn: a function to create a tensor input dictionary.
create_model_fn: a function that creates a DetectionModel and generates
losses.
train_config: a train_pb2.TrainConfig protobuf.
master: BNS name of the TensorFlow master to use.
task: The task id of this training instance.
num_clones: The number of clones to run per machine.
worker_replicas: The number of work replicas to train with.
clone_on_cpu: True if clones should be forced to run on CPU.
ps_tasks: Number of parameter server tasks.
worker_job_name: Name of the worker job.
is_chief: Whether this replica is the chief replica.
train_dir: Directory to write checkpoints and training summaries to.
"""
data_augmentation_options = [
preprocessor_builder.build(step)
for step in train_config.data_augmentation_options
]
with tf.Graph().as_default():
# Build a configuration specifying multi-GPU and multi-replicas.
deploy_config = model_deploy.DeploymentConfig(
num_clones=num_clones,
clone_on_cpu=clone_on_cpu,
replica_id=task,
num_replicas=worker_replicas,
num_ps_tasks=ps_tasks,
worker_job_name=worker_job_name)
# Place the global step on the device storing the variables.
with tf.device(deploy_config.variables_device()):
global_step = tf.train.create_global_step()
with tf.device(deploy_config.inputs_device()), \
tf.name_scope('Input'):
input_queue_list = []
for i, create_tensor_dict_fn in enumerate(
create_tensor_dict_fn_list):
input_queue_list.append(
_create_input_queue(
train_config.batch_size[i] // num_clones,
create_tensor_dict_fn,
train_config.batch_queue_capacity,
train_config.num_batch_queue_threads,
train_config.prefetch_queue_capacity,
data_augmentation_options))
# Gather initial summaries.
summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES))
global_summaries = set([])
model_fn = functools.partial(_create_losses,
create_model_fn=create_model_fn)
clones = model_deploy.create_clones(deploy_config, model_fn,
[input_queue_list])
first_clone_scope = clones[0].scope
# Gather update_ops from the first clone. These contain, for example,
# the updates for the batch_norm variables created by model_fn.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS,
first_clone_scope)
with tf.device(deploy_config.optimizer_device()), \
tf.name_scope('Optimizer'):
training_optimizer = optimizer_builder.build(
train_config.optimizer, global_summaries)
sync_optimizer = None
if train_config.sync_replicas:
training_optimizer = tf.train.SyncReplicasOptimizer(
training_optimizer,
replicas_to_aggregate=train_config.replicas_to_aggregate,
total_num_replicas=train_config.worker_replicas)
sync_optimizer = training_optimizer
# Create ops required to initialize the model from a given checkpoint.
init_fn = None
if train_config.fine_tune_checkpoint:
var_map = detection_model.restore_map(
from_detection_checkpoint=train_config.
from_detection_checkpoint)
available_var_map = variables_helper.get_variables_available_in_checkpoint(
var_map, train_config.fine_tune_checkpoint)
init_saver = tf.train.Saver(available_var_map)
def initializer_fn(sess):
init_saver.restore(sess, train_config.fine_tune_checkpoint)
init_fn = initializer_fn
with tf.device(deploy_config.optimizer_device()), \
tf.variable_scope('OptimizeClones'):
total_loss, grads_and_vars = model_deploy.optimize_clones(
clones, training_optimizer, regularization_losses=None)
total_loss = tf.check_numerics(total_loss,
'LossTensor is inf or nan.')
# Optionally multiply bias gradients by train_config.bias_grad_multiplier.
if train_config.bias_grad_multiplier:
biases_regex_list = [r'.*bias(?:es)?', r'.*beta']
grads_and_vars = variables_helper.multiply_gradients_matching_regex(
grads_and_vars,
biases_regex_list,
multiplier=train_config.bias_grad_multiplier)
# Optionally freeze some layers by setting their gradients to be zero.
if train_config.freeze_variables:
grads_and_vars = variables_helper.freeze_gradients_matching_regex(
grads_and_vars, train_config.freeze_variables)
# Optionally clip gradients
if train_config.gradient_clipping_by_norm > 0:
with tf.name_scope('clip_grads'):
grads_and_vars = tf.contrib.training.clip_gradient_norms(
grads_and_vars, train_config.gradient_clipping_by_norm)
# Create gradient updates.
grad_updates = training_optimizer.apply_gradients(
grads_and_vars, global_step=global_step)
update_ops.append(grad_updates)
update_op = tf.group(*update_ops)
with tf.control_dependencies([update_op]):
train_tensor = tf.identity(total_loss, name='train_op')
# Add summaries.
for (grad, var) in grads_and_vars:
var_name = var.op.name
grad_name = 'grad/' + var_name
global_summaries.add(tf.summary.histogram(grad_name, grad))
global_summaries.add(tf.summary.histogram(var_name, var))
# for model_var in tf.contrib.framework.get_model_variables():
# global_summaries.add(tf.summary.histogram(model_var.op.name, model_var))
for loss_tensor in tf.losses.get_losses():
global_summaries.add(
tf.summary.scalar(loss_tensor.op.name, loss_tensor))
global_summaries.add(
tf.summary.scalar('TotalLoss', tf.losses.get_total_loss()))
# Add the summaries from the first clone. These contain the summaries
# created by model_fn and either optimize_clones() or _gather_clone_loss().
summaries |= set(
tf.get_collection(tf.GraphKeys.SUMMARIES, first_clone_scope))
summaries |= global_summaries
# Merge all summaries together.
summary_op = tf.summary.merge(list(summaries), name='summary_op')
# Soft placement allows placing on CPU ops without GPU implementation.
session_config = tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False)
# Save checkpoints regularly.
keep_checkpoint_every_n_hours = train_config.keep_checkpoint_every_n_hours
saver = tf.train.Saver(
keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours)
scaffold = tf.train.Scaffold(init_fn=init_fn,
summary_op=summary_op,
saver=saver)
stop_hook = tf.train.StopAtStepHook(
num_steps=(train_config.num_steps
if train_config.num_steps else None), )
profile_hook = profile_session_run_hooks.ProfileAtStepHook(
at_step=200, checkpoint_dir=train_dir)
tf.contrib.training.train(
train_tensor,
train_dir,
master=master,
is_chief=is_chief,
scaffold=scaffold,
hooks=[stop_hook, profile_hook],
chief_only_hooks=None,
save_checkpoint_secs=train_config.save_checkpoint_secs,
save_summaries_steps=train_config.save_summaries_steps,
config=session_config)
def evaluate(create_input_dict_fn, create_model_fn, eval_config,
checkpoint_dir, eval_dir,
repeat_evaluation=True):
model = create_model_fn()
data_preprocessing_steps = [
preprocessor_builder.build(step)
for step in eval_config.data_preprocessing_steps]
tensor_dict = _extract_prediction_tensors(
model=model,
create_input_dict_fn=create_input_dict_fn,
data_preprocessing_steps=data_preprocessing_steps,
ignore_groundtruth=eval_config.ignore_groundtruth,
evaluate_with_lexicon=eval_config.eval_with_lexicon)
summary_writer = tf.summary.FileWriter(eval_dir)
def _process_batch(tensor_dict, sess, batch_index, counters, update_op):
if batch_index >= eval_config.num_visualizations:
if 'original_image' in tensor_dict:
tensor_dict = {k: v for (k, v) in tensor_dict.items()
if k != 'original_image'}
try:
(result_dict, _) = sess.run([tensor_dict, update_op])
counters['success'] += 1
except tf.errors.InvalidArgumentError:
logging.info('Skipping image')
counters['skipped'] += 1
return {}
global_step = tf.train.global_step(sess, tf.train.get_global_step())
if batch_index < eval_config.num_visualizations:
eval_util.visualize_recognition_results(
result_dict,
'Recognition_{}'.format(batch_index),
global_step,
summary_dir=eval_dir,
export_dir=os.path.join(eval_dir, 'vis'),
summary_writer=summary_writer,
only_visualize_incorrect=eval_config.only_visualize_incorrect)
return result_dict
def _process_aggregated_results(result_lists):
eval_metric_fn_key = eval_config.metrics_set
if eval_metric_fn_key not in EVAL_METRICS_FN_DICT:
raise ValueError('Metric not found: {}'.format(eval_metric_fn_key))
return EVAL_METRICS_FN_DICT[eval_metric_fn_key](result_lists)
variables_to_restore = tf.global_variables()
global_step = tf.train.get_or_create_global_step()
variables_to_restore.append(global_step)
if eval_config.use_moving_averages:
variable_averages = tf.train.ExponentialMovingAverage(0.0)
variables_to_restore = variable_averages.variables_to_restore()
saver = tf.train.Saver(variables_to_restore)
def _restore_latest_checkpoint(sess):
latest_checkpoint = tf.train.latest_checkpoint(checkpoint_dir)
saver.restore(sess, latest_checkpoint)
eval_util.repeated_checkpoint_run(
tensor_dict=tensor_dict,
update_op=tf.no_op(),
summary_dir=eval_dir,
aggregated_result_processor=_process_aggregated_results,
batch_processor=_process_batch,
checkpoint_dirs=[checkpoint_dir],
variables_to_restore=None,
restore_fn=_restore_latest_checkpoint,
num_batches=eval_config.num_examples,
eval_interval_secs=eval_config.eval_interval_secs,
max_number_of_evaluations=(
1 if eval_config.ignore_groundtruth else
eval_config.max_evals if eval_config.max_evals else
None if repeat_evaluation else 1),
master=eval_config.eval_master,
save_graph=eval_config.save_graph,
save_graph_dir=(eval_dir if eval_config.save_graph else ''))
summary_writer.close()