<jupyter_start><jupyter_text>Copyright 2018 The TensorFlow Authors.<jupyter_code>#@title 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
#
# https://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.<jupyter_output><empty_output><jupyter_text>Playing around with the TPU tutorial notebook before training transformer on TPU<jupyter_code>import tensorflow as tf
import os
import tensorflow_datasets as tfds
import time
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
strategy = tf.distribute.TPUStrategy(resolver)
def create_model():
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),
       tf.keras.layers.Conv2D(256, 3, activation='relu'),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(256, activation='relu'),
       tf.keras.layers.Dense(128, activation='relu'),
       tf.keras.layers.Dense(10)])
def get_dataset(batch_size, is_training=True):
  split = 'train' if is_training else 'test'
  dataset, info = tfds.load(name='mnist', split=split, with_info=True,
                            as_supervised=True, try_gcs=True)

  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255.0

    return image, label

  dataset = dataset.map(scale)

  # Only shuffle and repeat the dataset in training. The advantage to have a
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so users don't need to think about scaling the gradients
  # based on the actual batch size.
  if is_training:
    dataset = dataset.shuffle(10000)
    dataset = dataset.repeat()

  #dataset = dataset.batch(batch_size)
  dataset = dataset.batch(batch_size).prefetch(8)

  return dataset
with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

batch_size = 200
steps_per_epoch = 60000 // batch_size
validation_steps = 10000 // batch_size

train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset, 
          validation_steps=validation_steps,verbose = 2)
#TPU 11s/epoch
#no cache
#no pretech

"""Epoch 1/5
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.
  1/300 [..............................] - ETA: 22:40 - loss: 2.3021 - sparse_categorical_accuracy: 0.1300WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0026s vs `on_train_batch_end` time: 0.0308s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0026s vs `on_train_batch_end` time: 0.0308s). Check your callbacks.
299/300 [============================>.] - ETA: 0s - loss: 0.1350 - sparse_categorical_accuracy: 0.9583WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0014s vs `on_test_batch_end` time: 0.0115s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0014s vs `on_test_batch_end` time: 0.0115s). Check your callbacks.
300/300 [==============================] - 18s 60ms/step - loss: 0.1347 - sparse_categorical_accuracy: 0.9585 - val_loss: 0.0443 - val_sparse_categorical_accuracy: 0.9863
Epoch 2/5
300/300 [==============================] - 11s 36ms/step - loss: 0.0337 - sparse_categorical_accuracy: 0.9896 - val_loss: 0.0404 - val_sparse_categorical_accuracy: 0.9881
Epoch 3/5
300/300 [==============================] - 11s 36ms/step - loss: 0.0189 - sparse_categorical_accuracy: 0.9938 - val_loss: 0.0465 - val_sparse_categorical_accuracy: 0.9871
Epoch 4/5
300/300 [==============================] - 11s 37ms/step - loss: 0.0122 - sparse_categorical_accuracy: 0.9959 - val_loss: 0.0542 - val_sparse_categorical_accuracy: 0.9859
Epoch 5/5
300/300 [==============================] - 11s 36ms/step - loss: 0.0109 - sparse_categorical_accuracy: 0.9964 - val_loss: 0.0484 - val_sparse_categorical_accuracy: 0.9869
<tensorflow.python.keras.callbacks.History at 0x7fcf79ed3400>"""

#TPU 11s/epoch
#with cache
#with pretech

"""
Epoch 1/5
  1/300 [..............................] - ETA: 15:57 - loss: 2.3030 - sparse_categorical_accuracy: 0.1100WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0025s vs `on_train_batch_end` time: 0.0316s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0025s vs `on_train_batch_end` time: 0.0316s). Check your callbacks.
299/300 [============================>.] - ETA: 0s - loss: 0.1310 - sparse_categorical_accuracy: 0.9584WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_test_batch_end` time: 0.0129s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_test_batch_end` time: 0.0129s). Check your callbacks.
300/300 [==============================] - 16s 54ms/step - loss: 0.1306 - sparse_categorical_accuracy: 0.9585 - val_loss: 0.0464 - val_sparse_categorical_accuracy: 0.9853
Epoch 2/5
300/300 [==============================] - 11s 36ms/step - loss: 0.0342 - sparse_categorical_accuracy: 0.9891 - val_loss: 0.0398 - val_sparse_categorical_accuracy: 0.9874
Epoch 3/5
300/300 [==============================] - 11s 36ms/step - loss: 0.0188 - sparse_categorical_accuracy: 0.9940 - val_loss: 0.0415 - val_sparse_categorical_accuracy: 0.9869
Epoch 4/5
300/300 [==============================] - 11s 36ms/step - loss: 0.0131 - sparse_categorical_accuracy: 0.9958 - val_loss: 0.0455 - val_sparse_categorical_accuracy: 0.9864
Epoch 5/5
300/300 [==============================] - 11s 37ms/step - loss: 0.0087 - sparse_categorical_accuracy: 0.9970 - val_loss: 0.0475 - val_sparse_categorical_accuracy: 0.9876
<tensorflow.python.keras.callbacks.History at 0x7fcf79fbe320>"""

#Something is bottlenecking each batch. 
#TPU 9s/epoch
#with cache
#with pretech
#verbose = 2

"""Epoch 1/5
WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0019s vs `on_train_batch_end` time: 0.0290s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0019s vs `on_train_batch_end` time: 0.0290s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0014s vs `on_test_batch_end` time: 0.0116s). Check your callbacks.
WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0014s vs `on_test_batch_end` time: 0.0116s). Check your callbacks.
300/300 - 15s - loss: 0.1381 - sparse_categorical_accuracy: 0.9587 - val_loss: 0.0423 - val_sparse_categorical_accuracy: 0.9851
Epoch 2/5
300/300 - 9s - loss: 0.0339 - sparse_categorical_accuracy: 0.9891 - val_loss: 0.0395 - val_sparse_categorical_accuracy: 0.9877
Epoch 3/5
300/300 - 9s - loss: 0.0190 - sparse_categorical_accuracy: 0.9938 - val_loss: 0.0381 - val_sparse_categorical_accuracy: 0.9881
Epoch 4/5
300/300 - 9s - loss: 0.0113 - sparse_categorical_accuracy: 0.9959 - val_loss: 0.0455 - val_sparse_categorical_accuracy: 0.9877
Epoch 5/5
300/300 - 10s - loss: 0.0101 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0435 - val_sparse_categorical_accuracy: 0.9890
<tensorflow.python.keras.callbacks.History at 0x7fcf6f98bc88>"""

# Create the model, optimizer and metrics inside strategy scope, so that the
# variables can be mirrored on each device.
with strategy.scope():
  model = create_model()
  optimizer = tf.keras.optimizers.Adam()
  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      'training_accuracy', dtype=tf.float32)

# Calculate per replica batch size, and distribute the datasets on each TPU
# worker.
per_replica_batch_size = batch_size // strategy.num_replicas_in_sync

train_dataset = strategy.experimental_distribute_datasets_from_function(
    lambda _: get_dataset(per_replica_batch_size, is_training=True))

@tf.function
def train_step(iterator):
  """The step function for one training step"""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  strategy.run(step_fn, args=(next(iterator),))
steps_per_eval = 10000 // batch_size

train_iterator = iter(train_dataset)
for epoch in range(5):
  tstart = time.time()
  print('Epoch: {}/5'.format(epoch))

  for step in range(steps_per_epoch):
    train_step(train_iterator)
  print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
  tend = time.time()
  print('Time per epoch: {}/5'.format(tend-tstart))
  training_loss.reset_states()
  training_accuracy.reset_states()
@tf.function
def train_multiple_steps(iterator, steps):
  
  """The step function for one training step"""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  for _ in tf.range(steps):
    strategy.run(step_fn, args=(next(iterator),))

# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get 
# retraced if the value changes.
train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))

print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))<jupyter_output>Current step: 3300, training loss: 0.0059, accuracy: 99.81%
{'fraction_non_alphanumeric': 0.13627206060117655, 'fraction_numerical': 0.06350229672012249, 'mean_word_length': 4.790947270181988, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 13, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '20515', 'license': 'Apache-2.0', 'n_tokens_mistral': 4706, 'n_tokens_neox': 4119, 'n_words': 1017, 'path': 'TF_TPU_test.ipynb', 'repo_name': 'kyle-gao/ML_ipynb', 'size': '12409'}