Function create_model

tensorflow/examples/tf2_showcase/mnist.py:69–112

Model to recognize digits in the MNIST dataset. Network structure is equivalent to: https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/examples/tutorials/mnist/mnist_deep.py and https://github.com/tensorflow/models/blob/master/tutorials/image/mnist/convolutional.py But uses the

()

def create_model():
  """Model to recognize digits in the MNIST dataset.

  Network structure is equivalent to:
  https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/examples/tutorials/mnist/mnist_deep.py
  and
  https://github.com/tensorflow/models/blob/master/tutorials/image/mnist/convolutional.py
  But uses the tf.keras API.
  Returns:
    A tf.keras.Model.
  """
  # Assumes data_format == 'channel_last'.
  # See https://www.tensorflow.org/performance/performance_guide#data_formats

  input_shape = [28, 28, 1]

  l = tf.keras.layers
  max_pool = l.MaxPooling2D((2, 2), (2, 2), padding='same')
  # The model consists of a sequential chain of layers, so tf.keras.Sequential
  # (a subclass of tf.keras.Model) makes for a compact description.
  model = tf.keras.Sequential(
      [
          l.Reshape(
              target_shape=input_shape,
              input_shape=(28 * 28,)),
          l.Conv2D(2, 5, padding='same', activation=tf.nn.relu),
          max_pool,
          l.Conv2D(4, 5, padding='same', activation=tf.nn.relu),
          max_pool,
          l.Flatten(),
          l.Dense(32, activation=tf.nn.relu),
          l.Dropout(0.4),
          l.Dense(10)
      ])
  # TODO(brianklee): Remove when @kaftan makes this happen by default.
  # TODO(brianklee): remove `autograph=True` when kwarg default is flipped.
  model.call = tfe.function(model.call, autograph=True)
  # Needs to have input_signature specified in order to be exported
  # since model.predict() is never called before saved_model.export()
  # TODO(brianklee): Update with input signature, depending on how the impl of
  # saved_model.restore() pans out.
  model.predict = tfe.function(model.predict, autograph=True)
  # ,input_signature=(tensor_spec.TensorSpec(shape=[28, 28, None], dtype=tf.float32),) # pylint: disable=line-too-long
  return model


def mnist_datasets():