In [1]:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from sklearn import metrics
import tensorflow as tf

In [2]:
layers = tf.contrib.layers
learn = tf.contrib.learn

In [3]:
def max_pool_2x2(tensor_in):
    return tf.nn.max_pool(
        tensor_in, ksize=[1,2,2,1], strides=[1,2,2,1], padding='SAME')

In [5]:
def conv_model(feature, target, mode):
    """2-layer convolution model."""
    target = tf.one_hot(tf.cast(target, tf.int32), 10, 1, 0)
    feature = tf.reshape(feature, [-1, 28, 28, 1])
    with tf.variabale_score('conv_layer1'):
        h_conv1 = layers.convolution2d(
            feature, 32, kernel_size[5,5], activation_fn=tf.nn.relu)
        h_pool1 = max_pool_2x2(h_conv1)
    
    with tf.variable_scope('conv_layer2'):
        h_conv2 = layers.convolution2d(
            h_pool1, 64, kernel_size = [5,5], activation_fn=tf.nn.relu)
        h_pool2 = max_pool_2x2(h_conv2)
        h_pool2_flat = tf.reshape(h_pool2, [-1, 7*7*64])
    
    h_fc1 = layers.dropout(
        layers.fully_connected(
            h_pool2_flat, 1024, activation_fn=tf.nn.relu),
        keep_prob=0.5,
        is_training = (mode==tf.contrib.learn.ModeKeys.TRAIN)
        )
    
    logits = layers.fully_connected(h_fc1, 10, activation_fn=None)
    loss = tf.losses.softmax_cross_entropy(target, logits)
    
    train_op = layers.optimize_loss(
        loss,
        tf.contrib.framework.get_global_step(),
        optimizer='SGD',
        learning_rate=0.001)
    
    return tf.argmax(logits, 1), loss, train_op

In [ ]:
def main(unused_args):
    # Download and load MNIST dataset
    mnist = learn.datasets.load_dataset('mnist')
    
    # Linear classifier
    feature_columns =