<jupyter_start><jupyter_text>Copyright 2018 The TensorFlow Authors.Licensed under the Apache License, Version 2.0 (the "License"). MNIST, with TensorFlow 2.0  This notebook is still under construction! Please come back later. This notebook trains a simple MNIST model, demonstrating a basic workflow using TensorFlow 2.0 APIs.The basic workflow consists of:- defining a model- preprocessing your data into a tf.data.Dataset- training over a dataset  - using tf.GradientTape to compute gradients  - using stateful tf.metrics.* APIs to record metrics of interest  - logging those metrics with tf.summary.* APIs so that they can be viewed in TensorBoard  - using tf.train.Checkpoint to save and restore weights- export a SavedModel using tf.saved_model (This SavedModel is a portable representation of the model, and can be imported into C++, JS, Python without knowledge of the original TensorFlow code.)- reimport that SavedModel and demonstrate its usage in Python.<jupyter_code>!pip install tf-nightly-2.0-preview
from __future__ import absolute_import, division, print_function

import os
import time

import numpy as np
import tensorflow as tf
#TODO(brianklee): remove these when new modules are exported.
from tensorflow.python.ops import summary_ops_v2

print(tf.__version__)<jupyter_output><empty_output><jupyter_text>Define a convolution-based model, using Keras APIs.<jupyter_code>def create_model():
  # Assumes data_format == 'channel_last'.
  # See https://www.tensorflow.org/performance/performance_guide#data_formats

  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.
  return tf.keras.Sequential([
      l.Reshape(
          target_shape=[28, 28, 1],
          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)])

# Define a loss function and accuracy function
def compute_loss(logits, labels):
  return tf.reduce_mean(
      tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels))


def compute_accuracy(logits, labels):
  predictions = tf.argmax(logits, axis=1, output_type=tf.int64)
  labels = tf.cast(labels, tf.int64)
  return tf.reduce_mean(
      tf.cast(tf.equal(predictions, labels), dtype=tf.float32))

# Create the model and optimizer
model = create_model()

# TODO(brianklee): Switch out when keras optimizer v2 is published.
optimizer = tf.compat.v1.train.MomentumOptimizer(0.01, 0.5)
# optimizer = tf.keras.optimizers.SGD(lr=0.01, momentum=0.5)

# Create a step variable to track how much we've trained.
global_step = tf.Variable(0, dtype=tf.int64)<jupyter_output><empty_output><jupyter_text>Download and create tf.data.DatasetsLet's load the MNIST dataset into TF datasets. This lets us use all sorts of useful transformations like batching and shuffling.<jupyter_code># Set up datasets
def mnist_datasets():
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
  # Numpy defaults to dtype=float64; TF defaults to float32. Stick with float32.
  x_train, x_test = x_train / np.float32(255), x_test / np.float32(255)
  y_train, y_test = y_train.astype(np.int64), y_test.astype(np.int64)
  train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
  test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
  return train_dataset, test_dataset

train_ds, test_ds = mnist_datasets()
train_ds = train_ds.shuffle(60000).batch(100)
test_ds = test_ds.batch(100)

print('Dataset will yield tensors of the following shape: {}'.format(train_ds.output_shapes))<jupyter_output><empty_output><jupyter_text>Configure trainingOur train() function does a few things. It iterates over our training dataset, computing the gradients for each batch and then applying them to the model variables. Along the way, we periodically output summaries.<jupyter_code># TODO(brianklee): Enable @tf.function on the training loop when zip, enumerate
# are supported by autograph.
def train(model, optimizer, dataset, global_step, num_steps=None, log_freq=10):
  """Trains model on `dataset` using `optimizer`."""
  start = time.time()
  # Metrics are stateful. They accumulate values and return a cumulative
  # result when you call .result(). Clear accumulated values with .reset_states()
  avg_loss = tf.metrics.Mean('loss', dtype=tf.float32)
  avg_accuracy = tf.metrics.Mean('accuracy', dtype=tf.float32)
  # Datasets can be iterated over like any other Python iterable.
  # tf.function() will convert Python iteration over tf.Datasets into
  # TensorFlow dataset iteration ops.
  for (batch, (images, labels)) in enumerate(dataset):
    if num_steps is not None and batch > num_steps:
      break
    # Record the operations used to compute the loss, so that the gradient
    # of the loss with respect to the variables can be computed.
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = compute_loss(logits, labels)
      accuracy = compute_accuracy(logits, labels)
    grads = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(
        zip(grads, model.variables), global_step=global_step)
    avg_loss(loss)
    avg_accuracy(accuracy)
    if batch % log_freq == 0:
      summary_ops_v2.scalar('loss', avg_loss.result(), step=global_step)
      summary_ops_v2.scalar('accuracy', avg_accuracy.result(), step=global_step)
      avg_loss.reset_states()
      avg_accuracy.reset_states()
      rate = log_freq / (time.time() - start)
      print('Step #%d\tLoss: %.6f (%d steps/sec)' % (batch, loss, rate))
      start = time.time()

def test(model, dataset, global_step):
  """Perform an evaluation of `model` on the examples from `dataset`."""
  avg_loss = tf.metrics.Mean('loss', dtype=tf.float32)
  avg_accuracy = tf.metrics.Mean('accuracy', dtype=tf.float32)

  for (images, labels) in dataset:
    logits = model(images, training=False)
    avg_loss(compute_loss(logits, labels))
    avg_accuracy(compute_accuracy(logits, labels))
  print('Model test set loss: {:0.4f} accuracy: {:0.2f}%'.format(
      avg_loss.result(), avg_accuracy.result() * 100))
  summary_ops_v2.scalar('loss', avg_loss.result(), step=global_step)
  summary_ops_v2.scalar('accuracy', avg_accuracy.result(), step=global_step)<jupyter_output><empty_output><jupyter_text>Configure model directoryWe'll use one directory to save all of our relevant artifacts (summary logs, checkpoints, SavedModel exports, etc.)<jupyter_code># Where to save checkpoints, tensorboard summaries, etc.
MODEL_DIR = '/tmp/tensorflow/mnist'

def apply_clean():
  if tf.io.gfile.exists(MODEL_DIR):
    print('Removing existing model dir: {}'.format(MODEL_DIR))
    tf.io.gfile.rmtree(MODEL_DIR)

# Optional: wipe the existing directory

apply_clean()<jupyter_output><empty_output><jupyter_text>Let's configure where we'll output our summaries from training. Remember how we called tf.summary.scalar(...) in our train() function? By using our summary_writer in a `with` block, we can catch those generated summaries and direct them to a file.You can see the summaries with `tensorboard --logdir=`<jupyter_code>train_dir = os.path.join(MODEL_DIR, 'summaries', 'train')
test_dir = os.path.join(MODEL_DIR, 'summaries', 'eval')
train_summary_writer = summary_ops_v2.create_file_writer(
  train_dir, flush_millis=10000)
test_summary_writer = summary_ops_v2.create_file_writer(
  test_dir, flush_millis=10000, name='test')<jupyter_output><empty_output><jupyter_text>Configure checkpointsThe Checkpoint object helps manage which tf.Variables will be saved to and restored from checkpoint files.A checkpoint differs from a SavedModel because it must additionally keep track of training-related state, like momentum variables for a momentum-based optimizer or things like the global step. Additionally, it only stores weights, so you'll need the original code to define the computation using those weights.<jupyter_code>checkpoint_dir = os.path.join(MODEL_DIR, 'checkpoints')
checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')
checkpoint = tf.train.Checkpoint(
  model=model, optimizer=optimizer, global_step=global_step)
# Restore variables on creation if a checkpoint exists.
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))<jupyter_output><empty_output><jupyter_text>TrainNow that we've set up train() and test(), let's create a model and train it for some number of epochs.<jupyter_code>NUM_TRAIN_EPOCHS = 1

for i in range(NUM_TRAIN_EPOCHS):
  start = time.time()
  with train_summary_writer.as_default():
    train(model, optimizer, train_ds, global_step)
  end = time.time()
  print('\nTrain time for epoch #{} ({} total steps): {}'.format(
      i + 1, global_step.numpy(), end - start))
  with test_summary_writer.as_default():
    test(model, test_ds, global_step)
  checkpoint.save(checkpoint_prefix)<jupyter_output><empty_output><jupyter_text>Export a SavedModel<jupyter_code>export_path = os.path.join(MODEL_DIR, 'export')
tf.saved_model.save(model, export_path)<jupyter_output><empty_output><jupyter_text>Restore and run the SavedModelYou can restore any SavedModel and call it without reference to the original source code. APIs for importing and transforming SavedModels exist a variety of languages. See the [guide](https://www.tensorflow.org/guide/saved_model) for more.<jupyter_code>def import_and_eval():
  restored_model = tf.saved_model.restore(export_path)
  _, (x_test, y_test) = tf.keras.datasets.mnist.load_data()
  x_test = x_test / np.float32(255)
  y_predict = restored_model(x_test)
  accuracy = compute_accuracy(y_predict, y_test)
  print('Model accuracy: {:0.2f}%'.format(accuracy.numpy() * 100))

# TODO(brianklee): Activate after v2 import is implemented.
# import_and_eval()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10361301881158555, 'fraction_numerical': 0.013934507813277595, 'mean_word_length': 5.141809290953545, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 42, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '232990', 'license': 'Apache-2.0', 'n_tokens_mistral': 3331, 'n_tokens_neox': 3185, 'n_words': 918, 'path': 'site/en/2/guide/_tf2.ipynb', 'repo_name': 'colinmorris/docs', 'size': '10047'}