code: def main():
      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()
  
imports: import tensorflow
  import numpy
  from zoo import init_nncontext
  from zoo.tfpark import TFDataset, TFEstimator
  from zoo.tfpark import ZooOptimizer
label_desc: Use image node_id 2 for calling a global function with example usage: main() and returns: dataset, tf, tf