def model_fn(features, labels, mode, params):
"""Defines the model compatible with tf.estimator."""
del labels, params
if mode == tf.estimator.ModeKeys.TRAIN:
_build_deeplab(
features, model.get_output_to_num_classes(FLAGS),
model_input.dataset_descriptors[FLAGS.dataset].ignore_label)
# Print out the objective loss and regularization loss independently to
# track NaN loss issue
objective_losses = tf.losses.get_losses()
objective_losses = tf.Print(objective_losses, [objective_losses],
message='Objective Losses: ',
summarize=100)
objective_loss = tf.reduce_sum(objective_losses)
tf.summary.scalar('objective_loss', objective_loss)
reg_losses = tf.losses.get_regularization_losses()
reg_losses = tf.Print(reg_losses, [reg_losses],
message='Reg Losses: ',
summarize=100)
reg_loss = tf.reduce_sum(reg_losses)
tf.summary.scalar('regularization_loss', reg_loss)
loss = objective_loss + reg_loss
learning_rate = train_utils.get_model_learning_rate(
FLAGS.learning_policy, FLAGS.base_learning_rate,
FLAGS.learning_rate_decay_step, FLAGS.learning_rate_decay_factor,
FLAGS.training_number_of_steps, FLAGS.learning_power,
FLAGS.slow_start_step, FLAGS.slow_start_learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate)
tf.summary.scalar('learning_rate', learning_rate)
grads_and_vars = optimizer.compute_gradients(loss)
grad_updates = optimizer.apply_gradients(grads_and_vars,
tf.train.get_global_step())
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
update_ops.append(grad_updates)
update_op = tf.group(*update_ops)
with tf.control_dependencies([update_op]):
train_op = tf.identity(loss, name='train_op')
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
)
def model_fn(features, labels, mode, params):
"""Defines the model compatible with tf.estimator."""
del labels, params
if mode == tf.estimator.ModeKeys.TRAIN:
_build_deeplab(
features, {
'semantic':
model_input.dataset_descriptors[FLAGS.dataset].num_classes
}, model_input.dataset_descriptors[FLAGS.dataset].ignore_label)
loss = tf.losses.get_total_loss()
learning_rate = train_utils.get_model_learning_rate(
FLAGS.learning_policy, FLAGS.base_learning_rate,
FLAGS.learning_rate_decay_step, FLAGS.learning_rate_decay_factor,
FLAGS.training_number_of_steps, FLAGS.learning_power,
FLAGS.slow_start_step, FLAGS.slow_start_learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate)
tf.summary.scalar('learning_rate', learning_rate)
grads_and_vars = optimizer.compute_gradients(loss)
last_layers = model.get_extra_layer_scopes()
grad_mult = train_utils.get_model_gradient_multipliers(
last_layers, FLAGS.last_layer_gradient_multiplier)
if grad_mult:
grads_and_vars = tf.contrib.slim.learning.multiply_gradients(
grads_and_vars, grad_mult)
grad_updates = optimizer.apply_gradients(grads_and_vars,
tf.train.get_global_step())
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
update_ops.append(grad_updates)
update_op = tf.group(*update_ops)
with tf.control_dependencies([update_op]):
train_op = tf.identity(loss, name='train_op')
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
)