def model_fn(features, labels, mode, params):
"""Model function for DenseNet classifier.
Args:
features: inputs.
labels: one hot encoded classes
mode: one of tf.estimator.ModeKeys.{TRAIN, INFER, EVAL}
params: a parameter dictionary with the following keys:
Returns:
ModelFnOps for Estimator API.
"""
number_classes = params.get('number_classes')
growth_rate = params.get('growth_rate')
dropout_keep_prob = params.get('dropout_keep_prob')
encoder_num_units = params.get('encoder_num_units')
decoder_num_units = params.get('decoder_num_units')
bottleneck_number_feature_maps = params.get(
'bottleneck_number_feature_maps')
densenet = DenseNet(growth_rate=growth_rate,
dropout_keep_prob=dropout_keep_prob,
number_classes=number_classes,
is_training=(mode == tf.estimator.ModeKeys.TRAIN))
densenet.encode(
features=features,
num_units=encoder_num_units,
bottleneck_number_feature_maps=bottleneck_number_feature_maps)
logits = densenet.decode(decoder_num_units).output
predictions = tf.argmax(logits, axis=1)
tf.summary.image('predictions', predictions)
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = {
'class_ids': predictions,
'probabilities': tf.nn.softmax(logits),
'logits': logits,
}
return tf.estimator.EstimatorSpec(mode, predictions=predictions)
# Add the loss.
# Calculate loss, which includes softmax cross entropy and L2 regularization.
# wraps the softmax_with_entropy fn. adds it to loss collection
tf.losses.softmax_cross_entropy(logits=logits, onehot_labels=labels)
# include the regulization losses in the loss collection.
loss = tf.losses.get_total_loss()
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode=mode,
eval_metric_ops={
"accuracy":
tf.metrics.accuracy(
labels, predictions)
})
assert mode == tf.estimator.ModeKeys.TRAIN
global_step = tf.train.get_global_step()
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
optimizer = tf.train.AdagradOptimizer(learning_rate=0.1)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss, global_step=global_step)
return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
def model_fn(features, labels, mode, params):
"""Model function for DenseNet classifier.
Args:
features: inputs.
labels: one hot encoded classes
mode: one of tf.estimator.ModeKeys.{TRAIN, INFER, EVAL}
params: a parameter dictionary with the following keys:
Returns:
ModelFnOps for Estimator API.
"""
number_classes = params.get('number_classes')
growth_rate = params.get('growth_rate')
dropout_keep_prob = params.get('dropout_keep_prob')
encoder_num_units = params.get('encoder_num_units')
decoder_num_units = params.get('decoder_num_units')
bottleneck_number_feature_maps = params.get(
'bottleneck_number_feature_maps')
tf.logging.info("features tensor {}".format(features))
features, image_ids, image_shapes = features['image'], features[
'image_id'], features['image_shape']
densenet = DenseNet(growth_rate=growth_rate,
dropout_keep_prob=dropout_keep_prob,
number_classes=number_classes,
is_training=(mode == tf.estimator.ModeKeys.TRAIN))
densenet.encode(
features=features,
num_units=encoder_num_units,
bottleneck_number_feature_maps=bottleneck_number_feature_maps)
logits = densenet.decode(decoder_num_units).output
probs = tf.nn.softmax(logits)
predictions = tf.argmax(probs, axis=-1)
if mode == tf.estimator.ModeKeys.PREDICT:
# resize the predictions back to original size.
# label_shape = tf.shape(features)[:2]
probs = tf.image.resize_bilinear(probs,
image_shapes,
name='resize_predictions')
predictions = tf.argmax(probs, axis=-1)
tf.logging.info("Starting to predict..")
predictions = {
'class_ids': predictions,
'probabilities': probs,
'logits': logits,
'image_ids': image_ids
}
tf.logging.info("prediction tensor {}".format(predictions))
return tf.estimator.EstimatorSpec(mode, predictions=predictions)
# Add the loss.
# Calculate loss, which includes softmax cross entropy and L2 regularization.
# wraps the softmax_with_entropy fn. adds it to loss collection
tf.losses.softmax_cross_entropy(logits=logits, onehot_labels=labels)
# include the regulization losses in the loss collection.
loss = tf.losses.get_total_loss()
if mode == tf.estimator.ModeKeys.EVAL:
tf.logging.info("Starting to evaluate..")
with tf.variable_scope('mean_iou_calc'):
prec = []
up_opts = []
for t in np.arange(0.5, 1.0, 0.05):
predicted_mask = tf.to_int32(probs > t)
score, up_opt = tf.metrics.mean_iou(labels, predicted_mask, 2)
up_opts.append(up_opt)
prec.append(score)
mean_iou = tf.reduce_mean(tf.stack(prec),
axis=0), tf.stack(up_opts)
eval_metrics = {'mean_iou': mean_iou}
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
eval_metric_ops=eval_metrics)
assert mode == tf.estimator.ModeKeys.TRAIN
tf.logging.info("Starting to train..")
global_step = tf.train.get_global_step()
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
optimizer = tf.train.AdagradOptimizer(learning_rate=1e-4)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss, global_step=global_step)
add_summaries(predictions, features, loss)
return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)