code: # Copyright 2020 The ElasticDL Authors. All rights reserved.
  # 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 numpy as np
  import PIL.Image
  import tensorflow as tf
  
  from elasticdl.python.common.constants import Mode
  
  
  def custom_model():
      inputs = tf.keras.Input(shape=(28, 28), name="image")
      x = tf.keras.layers.Reshape((28, 28, 1))(inputs)
      x = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu")(x)
      x = tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu")(x)
      x = tf.keras.layers.BatchNormalization()(x)
      x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
      x = tf.keras.layers.Dropout(0.25)(x)
      x = tf.keras.layers.Flatten()(x)
      outputs = tf.keras.layers.Dense(10)(x)
  
      return tf.keras.Model(inputs=inputs, outputs=outputs, name="mnist_model")
  
  
  def prepare_data_for_a_single_file(file_object, filename):
      """
      :param filename: training data file name
      :param file_object: a file object associated with filename
      """
      label = int(filename.split("/")[-2])
      image = PIL.Image.open(file_object)
      numpy_image = np.array(image)
      example_dict = {
          "image": tf.train.Feature(
              float_list=tf.train.FloatList(value=numpy_image.flatten())
          ),
          "label": tf.train.Feature(
              int64_list=tf.train.Int64List(value=[label])
          ),
      }
      example = tf.train.Example(
          features=tf.train.Features(feature=example_dict)
      )
      return example.SerializeToString()
  
  
  def loss(labels, predictions):
      labels = tf.reshape(labels, [-1])
      return tf.reduce_mean(
          input_tensor=tf.nn.sparse_softmax_cross_entropy_with_logits(
              logits=predictions, labels=labels
          )
      )
  
  
  def optimizer(lr=0.01):
      return tf.optimizers.SGD(lr)
  
  
  def feed(dataset, mode, _):
      def _parse_data(record):
          if mode == Mode.PREDICTION:
              feature_description = {
                  "image": tf.io.FixedLenFeature([28, 28], tf.float32)
              }
          else:
              feature_description = {
                  "image": tf.io.FixedLenFeature([28, 28], tf.float32),
                  "label": tf.io.FixedLenFeature([1], tf.int64),
              }
          r = tf.io.parse_single_example(record, feature_description)
          features = {
              "image": tf.math.divide(tf.cast(r["image"], tf.float32), 255.0)
          }
          if mode == Mode.PREDICTION:
              return features
          else:
              return features, tf.cast(r["label"], tf.int32)
  
      dataset = dataset.map(_parse_data)
  
      if mode == Mode.TRAINING:
          dataset = dataset.shuffle(buffer_size=1024)
      return dataset
  
  
  def eval_metrics_fn():
      return {
          "accuracy": lambda labels, predictions: tf.equal(
              tf.argmax(predictions, 1, output_type=tf.int32),
              tf.cast(tf.reshape(labels, [-1]), tf.int32),
          )
      }
  
apis: tensorflow.keras.layers.Flatten tensorflow.reshape tensorflow.train.Int64List tensorflow.nn.sparse_softmax_cross_entropy_with_logits tensorflow.keras.layers.Conv2D tensorflow.io.parse_single_example tensorflow.train.Features tensorflow.keras.layers.BatchNormalization tensorflow.io.FixedLenFeature tensorflow.keras.layers.Dense tensorflow.keras.Input tensorflow.keras.layers.Dropout tensorflow.keras.layers.MaxPooling2D tensorflow.cast tensorflow.keras.Model tensorflow.argmax numpy.array tensorflow.optimizers.SGD tensorflow.keras.layers.Reshape