mnist_model_fn_simples.py

# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import numpy as np
 
print ('Versão do TensorFlow:', tf.__version__)

def model_fn(features, labels, mode, params):
  # Deve conter os seguintes passos:
  
  # 1. Definir o modelo via TF
  # nosso modelo consiste de apenas uma dense layer que produz
  # uma saída com 10 unidades (já que temos 10 dígitos possíveis)
  output = tf.layers.dense(inputs=features['x'], units=10)
  
  # 2. Definir como gerar predições
  predicoes = {
      'classes': tf.argmax(input=output, axis=1),
      'probabilidades': tf.nn.softmax(output, name='softmax_tensor')
  }
  
  # Se estamos realizando predição precisamos apenas especificar estas operações ;)
  if mode == tf.estimator.ModeKeys.PREDICT:
    # Return an EstimatorSpec object
    return tf.estimator.EstimatorSpec(mode=mode, predictions=predicoes)
  
  
  # 3. Definir a loss function para treino e avaliação
  # loss => erro, como calcular nosso erro?
  loss =  tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=output)


  # Se estamos treinando precisamos definir como otimizar nosso modelo
  # Para o Estimator esta definição é chamada train_op
  if mode == tf.estimator.ModeKeys.TRAIN:
    
    # 4. Definir como otimizar o modelo (optimizer)
    # como otimizar nosso modelo utilizando o erro que calculamos?
    # Neste exemplo estamos usando um algoritmo de otimizacao chamado Adam
    train_op = tf.contrib.layers.optimize_loss(
      loss=loss,
      global_step=tf.train.get_global_step(),
      learning_rate=params['learning_rate'],
      optimizer=params['optimizer'])
    
    return tf.estimator.EstimatorSpec(mode=mode, predictions=predicoes,
                                      loss=loss, train_op=train_op)
  
  
  # 5. Definir métricas para avaliação
  eval_metric_ops = {
      'acuracia': tf.metrics.accuracy(
          tf.argmax(input=output, axis=1),
          tf.argmax(input=labels, axis=1))
  }
  
  # 6. Retornar um EstimatorSpec definindo os passos acima
  return tf.estimator.EstimatorSpec(mode=mode, predictions=predicoes,
                                      loss=loss, eval_metric_ops=eval_metric_ops)