#
# Copyright 2018 Analytics Zoo Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import argparse

import tensorflow as tf
import tensorflow_datasets as tfds
from zoo.orca import init_orca_context, stop_orca_context
from zoo.orca.learn.tf.estimator import Estimator


def preprocess(data):
    data['image'] = tf.cast(data["image"], tf.float32) / 255.
    return data['image'], data['label']


def main(max_epoch, dataset_dir):

    mnist_train = tfds.load(name="mnist", split="train", data_dir=dataset_dir)
    mnist_test = tfds.load(name="mnist", split="test", data_dir=dataset_dir)

    mnist_train = mnist_train.map(preprocess)
    mnist_test = mnist_test.map(preprocess)

    model = tf.keras.Sequential(
        [tf.keras.layers.Conv2D(20, kernel_size=(5, 5), strides=(1, 1), activation='tanh',
                                input_shape=(28, 28, 1), padding='valid'),
         tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'),
         tf.keras.layers.Conv2D(50, kernel_size=(5, 5), strides=(1, 1), activation='tanh',
                                padding='valid'),
         tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'),
         tf.keras.layers.Flatten(),
         tf.keras.layers.Dense(500, activation='tanh'),
         tf.keras.layers.Dense(10, activation='softmax'),
         ]
    )

    model.compile(optimizer=tf.keras.optimizers.RMSprop(),
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    est = Estimator.from_keras(keras_model=model)
    est.fit(data=mnist_train,
            batch_size=320,
            epochs=max_epoch,
            validation_data=mnist_test)

    result = est.evaluate(mnist_test)
    print(result)

    est.save_keras_model("/tmp/mnist_keras.h5")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cluster_mode', type=str, default="local",
                        help='The mode for the Spark cluster. local or yarn.')

    args = parser.parse_args()
    if args.cluster_mode == "local":
        init_orca_context(cluster_mode="local", cores=4)
        dataset_dir = "~/tensorflow_datasets"
    elif args.cluster_mode == "yarn":
        init_orca_context(cluster_mode="yarn-client", num_nodes=2, cores=2, driver_memory="6g")
        dataset_dir = "hdfs:///tensorflow_datasets"
    main(5, dataset_dir)
    stop_orca_context()

@@ -600,16 +600,45 @@
 gparse%0A%0A
+from tensorflow import keras%0A
 import t
@@ -1219,19 +1219,16 @@
 model = 
-tf.
 keras.Se
@@ -1246,19 +1246,16 @@
        %5B
-tf.
 keras.la
@@ -1317,35 +1317,32 @@
 ivation='tanh',%0A
-   
                 
@@ -1398,35 +1398,32 @@
 lid'),%0A         
-tf.
 keras.layers.Max
@@ -1485,35 +1485,32 @@
 lid'),%0A         
-tf.
 keras.layers.Con
@@ -1601,19 +1601,16 @@
         
-  
- 
 padding=
@@ -1620,35 +1620,32 @@
 lid'),%0A         
-tf.
 keras.layers.Max
@@ -1707,35 +1707,32 @@
 lid'),%0A         
-tf.
 keras.layers.Fla
@@ -1740,35 +1740,32 @@
 ten(),%0A         
-tf.
 keras.layers.Den
@@ -1801,19 +1801,16 @@
         
-tf.
 keras.la
@@ -1893,19 +1893,16 @@
 timizer=
-tf.
 keras.op

Add orca tf keras quickstart & update tf quickstart (#3244)