code: def main(max_epoch, data_num):
      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()
  
      # get data, pre-process and create TFDataset
      (train_images_data, train_labels_data) = mnist.read_data_sets(
          "/tmp/mnist", "train"
      )
      (test_images_data, test_labels_data) = mnist.read_data_sets(
          "/tmp/mnist", "train"
      )
  
      train_images_data = (
          train_images_data[:data_num] - mnist.TRAIN_MEAN
      ) / mnist.TRAIN_STD
      train_labels_data = train_labels_data[:data_num].astype(np.int)
      test_images_data = (
          test_images_data[:data_num] - mnist.TRAIN_MEAN
      ) / mnist.TRAIN_STD
      test_labels_data = (test_labels_data[:data_num]).astype(np.int)
      dataset = TFDataset.from_ndarrays(
          (train_images_data, train_labels_data),
          batch_size=360,
          val_tensors=(test_images_data, test_labels_data),
      )
  
      # construct the model from TFDataset
      images, labels = dataset.tensors
  
      with slim.arg_scope(lenet.lenet_arg_scope()):
          logits, end_points = lenet.lenet(
              images, num_classes=10, is_training=True
          )
  
      loss = tf.reduce_mean(
          tf.losses.sparse_softmax_cross_entropy(
              logits=logits, labels=labels
          )
      )
  
      acc = accuracy(logits, labels)
  
      # create a optimizer
      optimizer = TFOptimizer.from_loss(
          loss,
          Adam(1e-3),
          metrics={"acc": acc},
          model_dir="/tmp/lenet/",
      )
      # kick off training
      optimizer.optimize(end_trigger=MaxEpoch(max_epoch))
  
      saver = tf.train.Saver()
      saver.save(optimizer.sess, "/tmp/lenet/model")
  
file_path: BigDL/python/orca/example/tfpark/tf_optimizer/train.py
imports: import tensorflow
  from bigdl.dllib.nncontext import init_nncontext
  from bigdl.orca.tfpark import TFOptimizer, TFDataset
  from bigdl.dllib.optim.optimizer import *
  from bigdl.dllib.utils.common import *
  import numpy
  from bigdl.dllib.feature.dataset import mnist
  from bigdl.dllib.feature.dataset.transformer import *
  from bigdl.dllib.utils.log4Error import invalidInputError
  import os
  import sys
  import argparse
  from nets import lenet
label_desc: Use image node_id 2 for calling a global function with example usage: main(max_epoch, data_num) without return types