#!/usr/bin/env python
# coding: utf-8

# nbviewer.jupyter.org/format/slides/gist/leechanwoo/

# In[1]:


from IPython.display import Image


# # Deep Learning Day 7.

# ## High Level APIs
#  * Estimator
#  * Experiment

# ### Estimator
#  * sklearn의 estimator와 유사하게 설계된 framework
#  * train, evaluate, predict 등을 쉽게 구현 할 수 있음.
#  * Session을 사용자가 직접 구현 할 필요 없음
#  * input pipeline만 작성하면 model 설계에 집중 할 수 있다.
#  * model export를 제공하며 serving시 필요한 protocol buffer file을 생성해줌

# #### input function
# 
#  * model에 투입되는 dataset을 load하는 pipeline을 구성하는 함수 정의
#  * 원하는 pipeline을 설계하고, 함수의 리턴은 train, labels로만 해주면 사용 가능
#  * Estimator에서 이 input function을 함수형으로 받기 때문에 입력 argument 필요시 아래 코드와 같이 lambda 식 활용

# In[1]:


import tensorflow as tf


# In[2]:


def parser(serialized_example):
    features = {
        'age': tf.FixedLenFeature([1], tf.int64),
        'img': tf.FixedLenFeature([61*49], tf.int64)
    }
    
    parsed_feature = tf.parse_single_example(serialized_example, features)
    age = parsed_feature['age']
    img = parsed_feature['img']
    return age, img


# In[3]:


tf.reset_default_graph()

def input_fn(dataset_dir, batch_size):
    dataset = tf.contrib.data.TFRecordDataset(dataset_dir).map(parser)
    dataset = dataset.batch(batch_size).shuffle(7777)
    
    itr = dataset.make_one_shot_iterator()
    
    age, img = itr.get_next()
    
    img = tf.reshape(img, [-1, 61, 49, 1])
    img = tf.cast(img, tf.float32)
    
    age = tf.reshape(age, [-1])
    age_onehot = tf.one_hot(age, depth=3, axis=-1)
    
    return {'img': img}, {'age': age, 'age_onehot': age_onehot}
    

def get_train_input_fn():
    return lambda: input_fn('./cnn_dataset/face_train.tfrecord', 10)

def get_test_input_fn():
    return lambda: input_fn('./cnn_dataset/face_test.tfrecord', 10)


# #### model function
# 
#  * 실제 모델이 들어가는 함수로써 입력인자(features, labels, mode) signature를 지켜주어야 한다.
#  * 이 model function도 마찬가지로 사용자의 hyper parameters를 외부에서 제어 할 수 있도록 lambda식을 사용하여 입력인자로 넘겨 받을 수 있다.
#  * 반드시 EstimatorSpec를 리턴시켜주어야 하며, TRAIN, EVAL, PREDICT 모드 별 트레인 모델을 분기문으로 나눠준다.

# In[4]:


def model(features, labels, mode, **params):
    conv1 = tf.layers.conv2d(features['img'], filters=10, kernel_size=3, 
                             padding='SAME', activation=params['activation'])
    pool1 = tf.layers.max_pooling2d(conv1, pool_size=2, strides=2)
    conv2 = tf.layers.conv2d(pool1, filters=10, kernel_size=3, 
                             padding='SAME', activation=params['activation'])
    pool2 = tf.layers.max_pooling2d(conv2, pool_size=2, strides=2)
    conv3 = tf.layers.conv2d(pool2, filters=10, kernel_size=3, 
                             padding='SAME', activation=params['activation'])
    pool3 = tf.layers.max_pooling2d(conv3, pool_size=2, strides=2)

    flat_size = int(pool3.shape[1]) * int(pool3.shape[2]) * int(pool3.shape[3])
    flat = tf.reshape(pool3, [-1, flat_size])
    
    if mode == tf.estimator.ModeKeys.TRAIN:
        dropout_prob = params['dropout_prob']
    else:
        dropout_prob = 1
    

    dropout1 = tf.layers.dropout(flat, dropout_prob)
    fc1 = tf.layers.dense(dropout1, 1000)

    dropout2 = tf.layers.dropout(fc1, dropout_prob)
    out = tf.layers.dense(dropout2, 3) 
    
    if mode == tf.estimator.ModeKeys.TRAIN:
        global_step = tf.train.get_global_step()
        loss = tf.losses.softmax_cross_entropy(labels['age_onehot'], out)
        train_op = tf.train.GradientDescentOptimizer(1e-6).minimize(loss, global_step)
        estimator_spec = tf.estimator.EstimatorSpec(train_op=train_op, loss=loss, mode=mode)

        
    elif mode == tf.estimator.ModeKeys.EVAL:
        loss = tf.losses.softmax_cross_entropy(labels['age_onehot'], out)
        
        pred = tf.argmax(tf.nn.softmax(out), axis=1)
        accuracy = tf.metrics.accuracy(labels['age'], pred)
        eval_metric_ops = {"acc": accuracy} 
        estimator_spec = tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
        
    else:
        estimator_spec = None
        
    return estimator_spec

def get_model_fn(activation, dropout_prob):
    return lambda features, labels, mode: model(features, labels, mode, 
                                                activation=activation, dropout_prob=0.7)


# #### Estimator
# 
#  * 아래와 같이 model 함수를 이용하여 Estimator를 생성해준다.
#  * 생성한 Estimator 이용하여 train, eval, predict 등을 수행 할 수 있다.

# In[ ]:


est_config = tf.estimator.RunConfig()
est = tf.estimator.Estimator(model_fn=get_model_fn(activation=tf.nn.relu, dropout_prob=0.7), 
                             model_dir='./logs/estimator',
                             config=config)

for epoch in range(3):
    est.train(get_train_input_fn())
    est.evaluate(get_test_input_fn())
    


# ### Experiment
#  * estimator를 기가막히게 돌려주는 framework
#  * local distributed setting만 되어 있다면 분산러닝 자동 지원
#  * training, inference, predict 모델을 따로 저장하기 때문에 serving 시 Experiment를 통한 export 권장
#  * 내부에 학습간 필요한 기능들을 다수 내장

# In[50]:


import json
import os
def get_experiment(output_dir):
    os.environ['TF_CONFIG'] = json.dumps({'environment': 'local'})
    config = tf.contrib.learn.RunConfig()
    return tf.contrib.learn.Experiment(estimator=tf.estimator.Estimator(model_fn=get_model_fn(activation=tf.nn.relu,
                                                                                              dropout_prob=0.7), 
                                                                        model_dir=output_dir,
                                                                        config=config),
                                       train_input_fn=get_train_input_fn(),
                                       eval_input_fn=get_test_input_fn())

tf.contrib.learn.learn_runner.run(experiment_fn=get_experiment, output_dir='./logs/estimator')