estimator_5.py

def model_fn(features, labels, mode):
    logits = neural_net_model(features, mode)
    class_prediction = tf.argmax(logits, axis=-1)
    preds = class_prediction

    loss = None
    train_op = None
    eval_metric_ops = {}

    if mode in (tf.estimator.ModeKeys.EVAL, tf.estimator.ModeKeys.TRAIN):
        loss = tf.reduce_mean(
          tf.nn.sparse_softmax_cross_entropy_with_logits(labels=tf.cast(labels, dtype=tf.int32), logits=logits))

    if mode == tf.estimator.ModeKeys.TRAIN:
        train_op = tf.train.AdamOptimizer().minimize(loss, global_step=tf.train.get_global_step())

    if mode == tf.estimator.ModeKeys.EVAL:
        eval_metric_ops = {
            'accuracy': tf.metrics.accuracy(
                labels=labels,
                predictions=preds,
                name='accuracy')
        }

    return tf.estimator.EstimatorSpec(mode=mode,
                                      predictions=class_prediction,
                                      loss=loss,
                                      train_op=train_op,
                                      eval_metric_ops=eval_metric_ops)