iris_estimator

iris_estimator.py

import pandas as pd
import tensorflow as tf
from sklearn import datasets
from tensorflow.contrib.estimator import InMemoryEvaluatorHook


def input_fn(df, batch_size, train=True):
    dataset = tf.data.Dataset.from_tensor_slices((df.as_matrix(), df.target.as_matrix().astype("int")))

    if train:
        return dataset.shuffle(1000).repeat().batch(batch_size)
    else:
        return dataset.batch(batch_size)


def model_fn(features, labels, mode, params):
    net = features
    for i, units in enumerate(params['hidden_units']):
        net = tf.layers.dense(net, units=units, activation=tf.nn.relu)

    logits = tf.layers.dense(net, units=params["n_classes"], activation=None)

    predicted_classes = tf.argmax(logits, 1)
    optimizer = tf.train.AdamOptimizer()
    acc = tf.metrics.accuracy(predictions=predicted_classes, labels=labels)
    loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits))
    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

    metrics = {"accuracy": acc}

    if mode == tf.estimator.ModeKeys.PREDICT:
        result = {
            'class_ids': predicted_classes,
            'probabilities': tf.nn.softmax(logits),
            'logits': logits
        }
        return tf.estimator.EstimatorSpec(mode, predictions=result)
    elif mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)
    else:
        return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

def main():
    tf.logging.set_verbosity(tf.logging.INFO)

    iris = datasets.load_iris()
    df = pd.DataFrame(iris.data, columns=iris.feature_names)
    df['target'] = iris.target

    columns = ["sepal_length", "sepal_width", "petal_length", "petal_width", "target"]
    df.columns = columns

    feature_columns = []
    for key in columns:
        feature_columns.append(tf.feature_column.numeric_column(key=key))

    estimator = tf.estimator.Estimator(model_fn=model_fn, params={"hidden_units":[15, 22], 'n_classes':3,
                                                                  'feature_columns':feature_columns})

    evaluator = InMemoryEvaluatorHook(
        estimator, lambda: input_fn(df, 15, False), every_n_iter=10)
    estimator.train(input_fn=lambda: input_fn(df, 15), hooks=[evaluator], steps=100)


if __name__ == "__main__":
    main()