def get_dataset(dataset_name, split_name, dataset_dir, reader=None):
"""Given a dataset dataset_name and a split_name returns a Dataset.
Args:
dataset_name: String, the dataset_name of the dataset.
split_name: A train/test split dataset_name.
dataset_dir: The directory where the dataset files are stored.
reader: The subclass of tf.ReaderBase. If left as `None`, then the default
reader defined by each dataset is used.
Returns:
A `Dataset` class.
Raises:
ValueError: If the dataset `dataset_name` is unknown.
"""
if dataset_name not in datasets_map:
raise ValueError('Name of dataset unknown %s' % dataset_name)
dataset_config = datasets_map[dataset_name]
file_pattern = dataset_config.file_pattern
num_samples = dataset_config.split_sizes[split_name]
return dataset_utils.get_split(split_name, dataset_dir,file_pattern, num_samples, reader)
def get_dataset(dataset_name, split_name, dataset_dir, reader=None):
"""Given a dataset dataset_name and a split_name returns a Dataset.
Args:
dataset_name: String, the dataset_name of the dataset.
split_name: A train/test split dataset_name.
dataset_dir: The directory where the dataset files are stored.
reader: The subclass of tf.ReaderBase. If left as `None`, then the default
reader defined by each dataset is used.
Returns:
A `Dataset` class.
Raises:
ValueError: If the dataset `dataset_name` is unknown.
"""
if dataset_name not in datasets_map:
raise ValueError('Name of dataset unknown %s' % dataset_name)
dataset_config = datasets_map[dataset_name];
file_pattern = dataset_config.file_pattern
num_samples = dataset_config.split_sizes[split_name]
return dataset_utils.get_split(split_name, dataset_dir,file_pattern, num_samples, reader)
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():
global_step = slim.create_global_step() # create the global step
######################
# select the dataset #
######################
dataset = dataset_utils.get_split(FLAGS.dataset_name,
FLAGS.dataset_split_name,
FLAGS.dataset_dir)
######################
# create the network #
######################
# parse the options from a yaml file
model, options = models_factory.get_model(FLAGS.model_config)
####################################################
# create a dataset provider that loads the dataset #
####################################################
# dataset 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)
[image] = provider.get(['image'])
image_clip = preprocessing_image(image,
model.training_image_size,
model.training_image_size,
model.content_size,
is_training=True)
image_clip_batch = tf.train.batch(
[image_clip],
batch_size=FLAGS.batch_size,
num_threads=FLAGS.num_preprocessing_threads,
capacity=5 * FLAGS.batch_size)
# feque queue the inputs
batch_queue = slim.prefetch_queue.prefetch_queue([image_clip_batch])
###########################################
# build the models based on the given data #
###########################################
images = batch_queue.dequeue()
total_loss = model.build_train_graph(images)
####################################################
# gather the operations for training and summaries #
####################################################
summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES))
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
# configurate 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
# gather the optimizer operations
learning_rate = _configure_learning_rate(dataset.num_samples,
global_step)
optimizer = _configure_optimizer(learning_rate)
summaries.add(tf.summary.scalar('learning_rate', learning_rate))
if FLAGS.moving_average_decay:
update_ops.append(
variable_averages.apply(moving_average_variables))
# training operations
train_op = model.get_training_operations(
optimizer, global_step, _get_variables_to_train(options))
update_ops.append(train_op)
# gather the training summaries
summaries |= set(model.summaries)
# gather the update operation
update_op = tf.group(*update_ops)
watched_loss = control_flow_ops.with_dependencies([update_op],
total_loss,
name='train_op')
# merge the summaries
summaries |= set(tf.get_collection(tf.GraphKeys.SUMMARIES))
summary_op = tf.summary.merge(list(summaries), name='summary_op')
##############################
# start the training process #
##############################
slim.learning.train(watched_loss,
logdir=FLAGS.train_dir,
init_fn=_get_init_fn(options),
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,
save_interval_secs=FLAGS.save_interval_secs)
