%env GCE_PROJECT_NAME my-sample-tpu-project
%env TPU_ZONE us-central1-f
%env TPU_NAME demo-tpu
%env TPU_IP_RANGE 10.240.1.0/29
%env GCS_DATA_PATH gs://cloud-tpu-data-bucket/mnist/
%env GCS_CKPT_PATH gs://cloud-tpu-checkpoint-bucket/mnist/
    
# Automatically get bucket name from GCS paths
import os
os.environ['GCS_DATA_BUCKET'] = os.environ['GCS_DATA_PATH'][5:].split('/')[0]
os.environ['GCS_CKPT_BUCKET'] = os.environ['GCS_CKPT_PATH'][5:].split('/')[0]

"""
Explanation: Getting Started with Cloud TPUs
Before We Start
This notebook requires that you run it on a GCE VM in a GCP project that has Cloud TPU quota. If you are not using this notebook on a pre-built virtual machine image, here is how you can start a new GCE VM with the right settings:

Create a new GCE VM with the following settings:
Zone: us-central1-c or us-central1-f
Machine Type: n1-standard-8
Operating System: Ubuntu 16.04 LTS with 64 GB of persistent disk
Identity and API access: Enable "Allow full access to all Cloud APIs"
Networking > Network Tag: tpu-jupyterhub-demo
Create a network rule under VPC network > Firewall Rules in the GCE control plane with the following settings:
Name: tpu-jupyterhub-demo
Target tags: tpu-jupyterhub-demo
Source IP ranges: 0.0.0.0/0
Protocols and ports: Specified protocols and ports, tcp:6006,8888
SSH into the GCE VM you have created in Step 1 and run the following:
  a. sudo apt-get update
  b. sudo apt-get -y install python3 python3-pip
  c. sudo -H pip3 install jupyter tf-nightly google-api-python-client
Start Jupyter in the GCE VM with jupyter notebook --no-browser --ip=0.0.0.0
Navigate to http://IP.OF.MY.VM:8888/?token=THE.TOKEN.DISPLAYED.ON.THE.COMMANDLINE on your favorite browser.
Upload this notebook to your Jupyter notebook by clicking on the Upload button.

Configuration
Please modify the following environment variables as it is required for the notebook.

GCE_PROJECT_NAME: The name of the GCE project this VM (and your Cloud TPU) starts in.
TPU_ZONE: The GCE zone in which you want your Cloud TPU to start in.
TPU_NAME: The name of the Cloud TPU
TPU_IP_RANGE: The IP address range for the Cloud TPU
GCS_DATA_PATH: The GCS path where we will store sample test data for Cloud TPUs. As GCS bucket namespaces are global, you may need to change this.
GCS_CKPT_PATH: The GCS path where we will store sample checkpoint data for Cloud TPUs.  As GCS bucket namespaces are global, you may need to change this.

Note: We will grant Storage Admin permissions to Cloud TPUs on the GCS buckets specified in GCS_DATA_PATH and GCS_CKPT_PATH. We encourage you to specify new buckets or test buckets containing non-production data.
End of explanation
"""


!gcloud config set compute/zone $TPU_ZONE
!gcloud alpha compute tpus create $TPU_NAME --range=$TPU_IP_RANGE --accelerator-type=tpu-v2 --version=nightly --zone=$TPU_ZONE

"""
Explanation: Create a new Cloud TPU
You can create a new Cloud TPU by running the command in the cell below.
End of explanation
"""


!gsutil mb -c regional -l us-central1 gs://$GCS_DATA_BUCKET
!gsutil mb -c regional -l us-central1 gs://$GCS_CKPT_BUCKET

!gsutil iam ch serviceAccount:`gcloud alpha compute tpus describe $TPU_NAME | grep serviceAccount | cut -d' ' -f2`:admin gs://$GCS_DATA_BUCKET gs://$GCS_CKPT_BUCKET && echo 'Successfully set permissions!'

"""
Explanation: Create GCS Buckets for Cloud TPU
Here, we will create two GCS buckets -- one for training/test data (GCS_DATA_BUCKET), and the other for TensorFlow checkpoint and TensorBoard metric data (GCS_CKPT_BUCKET).
The first two commands creates the buckets in a single region for maximum performance, and the final command grants the Cloud TPU Service Account owner access to the bucket so that it can read from and write to the bucket.
End of explanation
"""


import os
import tensorflow as tf
from tensorflow.contrib import tpu
from tensorflow.contrib.cluster_resolver import TPUClusterResolver

def axy_computation(a, x, y):
  return a * x + y

inputs = [
    3.0,
    tf.random_uniform([3, 3], 0, 1, tf.float32),
    tf.random_uniform([3, 3], 0, 1, tf.float32),
]
tpu_computation = tpu.rewrite(axy_computation, inputs)

tpu_cluster_resolver = TPUClusterResolver([os.environ['TPU_NAME']], zone=os.environ['TPU_ZONE'], project=os.environ['GCE_PROJECT_NAME'])
tpu_grpc_url = tpu_cluster_resolver.get_master()

with tf.Session(tpu_grpc_url) as sess:
  print('Initializing TPU...')
  sess.run(tpu.initialize_system())
  print('Initializing global variables...')
  sess.run(tf.global_variables_initializer())
  print('Executing TPU operation...')
  output = sess.run(tpu_computation)
  print(output)
  print('Shutting down TPU...')
  sess.run(tpu.shutdown_system())
  print('Done!')

"""
Explanation: Connect to Cloud TPU and Run a Simple AX+Y Calculation
The following code connects to your GCE VM (which is already running the proxy), resolves the IP address of the TPU using TPU Cluster Resovers, connects to the Cloud TPU with tf.Session, and run a simple calculation on TPUs. 
You should see a 3x3 array of random numbers being printed out if the command is successful. This computation may take up to 60 seconds to run to completion due to TPU initialization overheads.
End of explanation
"""


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import sys

import tensorflow as tf

from tensorflow.contrib.learn.python.learn.datasets import mnist


def _int64_feature(value):
  return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def _bytes_feature(value):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def convert_mnist_data_to_tfrecord(data_set, name):
  """Converts a dataset to tfrecords."""
  images = data_set.images
  labels = data_set.labels
  num_examples = data_set.num_examples

  if images.shape[0] != num_examples:
    raise ValueError('Images size %d does not match label size %d.' %
                     (images.shape[0], num_examples))
  rows = images.shape[1]
  cols = images.shape[2]
  depth = images.shape[3]

  filename = os.path.join(os.environ['GCS_DATA_PATH'], name + '.tfrecords')
  print('Writing', filename)
  with tf.python_io.TFRecordWriter(filename) as writer:
    for index in range(num_examples):
      image_raw = images[index].tostring()
      example = tf.train.Example(
          features=tf.train.Features(
              feature={
                  'height': _int64_feature(rows),
                  'width': _int64_feature(cols),
                  'depth': _int64_feature(depth),
                  'label': _int64_feature(int(labels[index])),
                  'image_raw': _bytes_feature(image_raw)
              }))
      writer.write(example.SerializeToString())


def convert_mnist_to_tfrecord():
  # Get the data.
  data_sets = mnist.read_data_sets('/tmp',
                                   dtype=tf.uint8,
                                   reshape=False,
                                   validation_size=5000)

  # Convert to Examples and write the result to TFRecords.
  convert_mnist_data_to_tfrecord(data_sets.train, 'train')
  convert_mnist_data_to_tfrecord(data_sets.validation, 'validation')
  convert_mnist_data_to_tfrecord(data_sets.test, 'test')

convert_mnist_to_tfrecord()
print("Finished writing TFRecords to %s" % os.environ['GCS_DATA_PATH'])

"""
Explanation: Creating and Uploading TFRecords from MNIST Test Data
The script below downloads the MNIST data from http://yann.lecun.com/exdb/mnist/ and creates TFRecord files from it. We then upload it to the GCS data bucket (GCS_DATA_BUCKET) that we created earlier. If the run is successful, you should see a message similar to Operation completed over 3 objects/59.5 MiB. as the last line of the output.
End of explanation
"""


from tensorflow.contrib.tpu.python.tpu import tpu_config
from tensorflow.contrib.tpu.python.tpu import tpu_estimator
from tensorflow.contrib.tpu.python.tpu import tpu_optimizer

def mnist_metric_fn(labels, logits):
  """Evaluation metric Fn which runs on CPU."""
  predictions = tf.argmax(logits, 1)
  return {
      "accuracy": tf.metrics.precision(
          labels=labels, predictions=predictions),
  }


def mnist_model_fn(features, labels, mode, params):
  """A simple CNN."""
  del params

  if mode == tf.estimator.ModeKeys.PREDICT:
    raise RuntimeError("mode {} is not supported yet".format(mode))

  input_layer = tf.reshape(features, [-1, 28, 28, 1])
  conv1 = tf.layers.conv2d(
      inputs=input_layer,
      filters=32,
      kernel_size=[5, 5],
      padding="same",
      activation=tf.nn.relu)
  pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
  conv2 = tf.layers.conv2d(
      inputs=pool1,
      filters=64,
      kernel_size=[5, 5],
      padding="same",
      activation=tf.nn.relu)
  pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
  pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
  dense = tf.layers.dense(inputs=pool2_flat, units=128, activation=tf.nn.relu)
  dropout = tf.layers.dropout(
      inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
  logits = tf.layers.dense(inputs=dropout, units=10)
  onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)

  loss = tf.losses.softmax_cross_entropy(
      onehot_labels=onehot_labels, logits=logits)

  if mode == tf.estimator.ModeKeys.EVAL:
    return tpu_estimator.TPUEstimatorSpec(
        mode=mode,
        loss=loss,
        eval_metrics=(mnist_metric_fn, [labels, logits]))

  # Train.
  learning_rate = tf.train.exponential_decay(0.05,
                                             tf.train.get_global_step(), 100000,
                                             0.96)

  optimizer = tpu_optimizer.CrossShardOptimizer(
      tf.train.GradientDescentOptimizer(learning_rate=learning_rate))

  train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
  return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op)


def mnist_get_input_fn(filename):
  """Returns an `input_fn` for train and eval."""

  def input_fn(params):
    """A simple input_fn using the experimental input pipeline."""
    # Retrieves the batch size for the current shard. The # of shards is
    # computed according to the input pipeline deployment. See
    # `tf.contrib.tpu.RunConfig` for details.
    batch_size = params["batch_size"]

    def parser(serialized_example):
      """Parses a single tf.Example into image and label tensors."""
      features = tf.parse_single_example(
          serialized_example,
          features={
              "image_raw": tf.FixedLenFeature([], tf.string),
              "label": tf.FixedLenFeature([], tf.int64),
          })
      image = tf.decode_raw(features["image_raw"], tf.uint8)
      image.set_shape([28 * 28])
      # Normalize the values of the image from the range [0, 255] to [-0.5, 0.5]
      image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
      label = tf.cast(features["label"], tf.int32)
      return image, label

    dataset = tf.data.TFRecordDataset(
        filename, buffer_size=None)
    dataset = dataset.map(parser).cache().repeat()
    dataset = dataset.apply(
        tf.contrib.data.batch_and_drop_remainder(batch_size))
    images, labels = dataset.make_one_shot_iterator().get_next()
    return images, labels
  return input_fn

print("Estimator-based MNIST Model Defined Successfully")

"""
Explanation: Defining a Simple Neural Network Model with TPU Estimators
We can define a simple model with TensorFlow Estimators that can train on the MNIST dataset to identify images. As part of using TensorFlow estimators, we need to create a model function that defines the model (mnist_model_fn), and a input function to process the inputs (the result of mnist_get_input_fn).
End of explanation
"""


from tensorflow.contrib.cluster_resolver import TPUClusterResolver
import os
import time

tpu_cluster_resolver = TPUClusterResolver([os.environ['TPU_NAME']], zone=os.environ['TPU_ZONE'], project=os.environ['GCE_PROJECT_NAME'])
tpu_grpc_url = tpu_cluster_resolver.get_master()

batch_size = 128
train_file = os.path.join(os.environ['GCS_DATA_PATH'], "train.tfrecords")
train_steps = 1000

eval_file = os.path.join(os.environ['GCS_DATA_PATH'], "validation.tfrecords")
eval_steps = 100

model_dir = os.path.join(os.environ['GCS_CKPT_PATH'], str(int(time.time()))) + "/"
iterations = 50
num_shards = 8

os.environ['MNIST_MODEL_DIR'] = model_dir
  
tf.logging.set_verbosity(tf.logging.INFO)

run_config = tpu_config.RunConfig(
    master=tpu_grpc_url,
    evaluation_master=tpu_grpc_url,
    model_dir=model_dir,
    session_config=tf.ConfigProto(
        allow_soft_placement=True, log_device_placement=True),
    tpu_config=tpu_config.TPUConfig(iterations, num_shards),)

estimator = tpu_estimator.TPUEstimator(
    model_fn=mnist_model_fn,
    use_tpu=True,
    train_batch_size=batch_size,
    eval_batch_size=batch_size,
    config=run_config)

estimator.train(input_fn=mnist_get_input_fn(train_file),
                max_steps=train_steps)

if eval_steps:
  estimator.evaluate(input_fn=mnist_get_input_fn(eval_file),
                     steps=eval_steps)


"""
Explanation: Running a TPU Estimator-based Model
To run a TPU Estimator-based model, we can define some configurations (tpu_config.RunConfig) and create a TPUEstimator using the model_fn, and then call train and evaluate on the Estimator. This computation may take up to 120 seconds to run to completion due to TPU initialization overheads.
End of explanation
"""


!echo Visit http://`curl -H "Metadata-Flavor: Google" http://metadata/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip 2> /dev/null`:6006 for TensorBoard
!which pip3 && pip3 install html5lib==0.99999999  # workaround for TensorBoard dependency error in Python 3
!tensorboard --logdir=$MNIST_MODEL_DIR

"""
Explanation: Visualize Graphs in TensorBoard
You can visualize the model and training details in TensorBoard. To launch Tensorboard, simply pass in the GCS path of the model directory into TensorBoard. Make sure that you have port 6006 open.
To stop TensorBoard, click on Kernel > Interrupt.
End of explanation
"""


!yes | gcloud alpha compute tpus delete $TPU_NAME

"""
Explanation: Deleting your Cloud TPU
To delete the Cloud TPU you have created, simply run the command below.
End of explanation
"""


