code: def main():
      args = parser.parse_args()
      cluster_mode = args.cluster_mode
      if cluster_mode.startswith("yarn"):
          hadoop_conf = os.environ.get("HADOOP_CONF_DIR")
          invalidInputError(
              hadoop_conf is not None,
              "Directory path to hadoop conf not found for yarn-client mode. Please "
              "set the environment variable HADOOP_CONF_DIR",
          )
          spark_conf = (
              create_spark_conf()
              .set("spark.executor.memory", "5g")
              .set("spark.executor.cores", 2)
              .set("spark.executor.instances", 2)
              .set("spark.driver.memory", "2g")
          )
          if cluster_mode == "yarn-client":
              sc = init_nncontext(
                  spark_conf,
                  cluster_mode="yarn-client",
                  hadoop_conf=hadoop_conf,
              )
          else:
              sc = init_nncontext(
                  spark_conf,
                  cluster_mode="yarn-cluster",
                  hadoop_conf=hadoop_conf,
              )
      else:
          sc = init_nncontext()
  
      def model_fn(features, labels, mode):
          from nets import lenet
  
          slim = tf.contrib.slim
          with slim.arg_scope(lenet.lenet_arg_scope()):
              logits, end_points = lenet.lenet(
                  features, num_classes=10, is_training=True
              )
  
          if (
              mode == tf.estimator.ModeKeys.EVAL
              or mode == tf.estimator.ModeKeys.TRAIN
          ):
              loss = tf.reduce_mean(
                  tf.losses.sparse_softmax_cross_entropy(
                      logits=logits, labels=labels
                  )
              )
  
              optimizer = ZooOptimizer(tf.train.AdamOptimizer())
              train_op = optimizer.minimize(loss)
              return tf.estimator.EstimatorSpec(
                  mode, predictions=logits, loss=loss, train_op=train_op
              )
          else:
              return tf.estimator.EstimatorSpec(
                  mode, predictions=logits
              )
  
      def input_fn(mode):
          if mode == tf.estimator.ModeKeys.TRAIN:
              training_data = get_data("train")
              dataset = TFDataset.from_ndarrays(
                  training_data, batch_size=320
              )
          elif mode == tf.estimator.ModeKeys.EVAL:
              testing_data = get_data("test")
              dataset = TFDataset.from_ndarrays(
                  testing_data, batch_per_thread=80
              )
          else:
              images, _ = get_data("test")
              dataset = TFDataset.from_ndarrays(
                  images, batch_per_thread=80
              )
  
          return dataset
  
      estimator = TFEstimator.from_model_fn(
          model_fn, model_dir="/tmp/estimator"
      )
  
      estimator.train(input_fn, steps=10)
  
      metrics = estimator.evaluate(input_fn, ["acc"])
      print(metrics)
  
      predictions = estimator.predict(input_fn)
  
      print(predictions.first())
      print("finished...")
      sc.stop()
  
file_path: BigDL/python/orca/example/tfpark/estimator/estimator_dataset.py
imports: import tensorflow
  import numpy
  from bigdl.dllib.nncontext import init_nncontext
  from bigdl.orca.tfpark import TFDataset, TFEstimator
  from bigdl.orca.tfpark import ZooOptimizer
  from bigdl.dllib.utils.common import *
  from bigdl.dllib.utils.log4Error import invalidInputError
  import os
  import argparse
label_desc: Use image node_id 2 for calling a global function with example usage: main() and returns: dataset, tf, tf