<jupyter_start><jupyter_text>Copyright 2019 The TensorFlow Hub Authors.Licensed under the Apache License, Version 2.0 (the "License");<jupyter_code># Copyright 2019 The TensorFlow Hub 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.
# ==============================================================================<jupyter_output><empty_output><jupyter_text>Train embeddings on TPU using Autoencoder  OverviewThis colab explores how to train autoencoders on a TPU device.For this colab, consider the following scenario: you have an image classification model that you want to improve by adding some additional features. The features that you can add to the model could be image embeddings that can be separately trained on a TPU.This example uses a fully-connected one layer model as the model that you want to make better with additional features trained on a TPU.  Learning objectivesIn this Colab, you will learn how to* Build a fully-connected one layer model to classify images* Build an autoencoder and train on those images, in an unsupervised fashion, to produce image embeddings* Retrain a fully-connected one layer model with additonal features, the embeddings  Check that you have enabled TPUs in this notebook<jupyter_code>import os
assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

tf.logging.set_verbosity(tf.logging.ERROR)

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']<jupyter_output><empty_output><jupyter_text>Check that the Tensorflow version is at least 1.14<jupyter_code>assert float('.'.join(tf.__version__.split('.')[:2])) >= 1.14, 'Make sure that Tensorflow version is at least 1.14'<jupyter_output><empty_output><jupyter_text>Get data<jupyter_code>(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()<jupyter_output><empty_output><jupyter_text>Function to visualize our images and pick the first image from the test set<jupyter_code>def show_img(img):
  plt.figure()
  plt.imshow(img)
  plt.grid(False)
  plt.show()

img = 0<jupyter_output><empty_output><jupyter_text>The first image from the test set is the number 7<jupyter_code>show_img(x_test[img].reshape(28, 28))<jupyter_output><empty_output><jupyter_text>MNIST setupThere are 10 classes (one for each digit) and each image is 28 by 28 pixels<jupyter_code>NUM_CLASSES = 10

# input image dimensions
IMG_ROWS, IMG_COLS = 28, 28

x_train = x_train.reshape(x_train.shape[0], IMG_ROWS, IMG_COLS, 1)
x_test = x_test.reshape(x_test.shape[0], IMG_ROWS, IMG_COLS, 1)

x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train = x_train / 255.0
x_test = x_test / 255.0<jupyter_output><empty_output><jupyter_text>Original modelHere is a contrived example where the training happens only on the corners of the MNIST image.Suppose that your original model, the fully-connected one layer network, was too computationally heavy, in terms of resources, and thus you could only afford to train on parts of the images. Instead of training on 28 by 28 pixels (784 pixels), you train on 14 by 14 pixels (196 pixels). This colab will later show that just by adding 49 more pixels to each training example, the size of each embedding, accuracy can be significantly increased.This way you introduce minimal changes to an original model while gaining benefits from a heavy computational task that you can be offload to a TPU.<jupyter_code>x_train_corners = x_train[:, :14, :14, :]
x_test_corners = x_test[:, :14, :14, :]<jupyter_output><empty_output><jupyter_text>The first image corner from the test set<jupyter_code>show_img(x_test_corners[img].reshape(14, 14))<jupyter_output><empty_output><jupyter_text>Create a model with one fully-connected layer<jupyter_code>def get_model(input_shape):
  ip = tf.keras.layers.Input(shape=input_shape)
  x = tf.keras.layers.Flatten()(ip)
  x = tf.keras.layers.Dense(NUM_CLASSES, activation='sigmoid')(x)
  
  model = tf.keras.models.Model(ip, x)
  return model
model0 = get_model(x_train_corners[0].shape)
model0.compile(
    optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.05), 
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'])<jupyter_output><empty_output><jupyter_text>Train and evaluate the fully-connected one layer model on CPUAs expected, the model performs poorly, but it does train fairly quickly. Expected accuracy is 65%.<jupyter_code>model0.fit(x_train_corners, y_train, epochs=3)
model0.evaluate(x_test_corners, y_test)<jupyter_output><empty_output><jupyter_text>Create an autoencoder and make sure to get back an encoder as well<jupyter_code>def get_autoencoder_and_encoder(input_shape):
  ip = tf.keras.layers.Input(shape=input_shape)

  x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(ip)
  x = tf.keras.layers.MaxPooling2D((2, 2), padding='same')(x)
  x = tf.keras.layers.Conv2D(1, (3, 3), activation='relu', padding='same')(x)
  encoded = tf.keras.layers.MaxPooling2D((2, 2), padding='same')(x)

  x = tf.keras.layers.Conv2DTranspose(1, (3, 3), activation='relu', strides=2, padding='same')(encoded)
  x = tf.keras.layers.Conv2DTranspose(32, (3, 3), activation='relu', strides=2, padding='same')(x)
  
  decoded = tf.keras.layers.Conv2DTranspose(1, (3, 3), activation='sigmoid', padding='same')(x)

  autoencoder = tf.keras.models.Model(ip, outputs=decoded)
  encoder = tf.keras.models.Model(ip, encoded)
  
  return autoencoder, encoder<jupyter_output><empty_output><jupyter_text>Train the autoencoder on TPUThis is a computationally resource expensive operation that can be offloaded to the TPU.<jupyter_code>resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER)
tf.contrib.distribute.initialize_tpu_system(resolver)
strategy = tf.contrib.distribute.TPUStrategy(resolver)

with strategy.scope():
  autoencoder, encoder = get_autoencoder_and_encoder(x_train[0].shape)
  autoencoder.compile(
    optimizer='adam',
    loss='binary_crossentropy')
  
  autoencoder.fit(
    x_train,
    x_train, 
    batch_size=128,
    epochs=3,
    steps_per_epoch=468,
    validation_data=(x_test, x_test))<jupyter_output><empty_output><jupyter_text>Produce image embeddingsNow that the autoencoder is trained, you can use the encoder part to produce image embeddings.<jupyter_code>x_train_embeddings = encoder.predict(x_train)
x_test_embeddings = encoder.predict(x_test)<jupyter_output><empty_output><jupyter_text>Produce image reconstructionsLet's visually see the quality of our autoencoder to see how the number 7 from above is reconstructed.<jupyter_code>x_test_hat = autoencoder.predict(x_test[:8])<jupyter_output><empty_output><jupyter_text>Reconstructed number 7This looks like the number 7 so now you can be confident in the quality of our embeddings. The autoencoder learned to compress and uncompress information accurately.<jupyter_code>show_img(x_test_hat[img].reshape(28, 28))<jupyter_output><empty_output><jupyter_text>Check the original imageRemember, the image in the previous section is the reconstructed image. Compare it to the original image, as shown here.<jupyter_code>show_img(x_test[0].reshape(28, 28))<jupyter_output><empty_output><jupyter_text>Examine the embedding for the number 7<jupyter_code>show_img(x_test_embeddings[0].reshape(7, 7))<jupyter_output><empty_output><jupyter_text>Augment the corners datasetThe following code augments the corners dataset with embeddings trained on TPU.<jupyter_code>x_train_augmented = np.concatenate([x_train_corners.reshape(60000, 14*14, 1), x_train_embeddings.reshape(60000, 7*7, 1)], axis=1)
x_test_augmented = np.concatenate([x_test_corners.reshape(10000, 14*14, 1), x_test_embeddings.reshape(10000, 7*7, 1)], axis=1)<jupyter_output><empty_output><jupyter_text>Retrain the original modelAt this point, you can train the original model using the augmented dataset. You should verify that the TPU embeddings augmented model works better than without embeddings. Expected accuracy is 87%.<jupyter_code>model1 = get_model(x_train_augmented[0].shape)
model1.compile(
    optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.05),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'])

model1.fit(x_train_augmented, y_train, epochs=3)
model1.evaluate(x_test_augmented, y_test)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10292463442069741, 'fraction_numerical': 0.022609673790776153, 'mean_word_length': 6.541136556403732, 'pattern_counts': {'":': 0, '<': 81, '<?xml version=': 0, '>': 84, 'https://': 0, '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': '20', 'dir': 'jupyter-structured-clean-dedup', 'id': '306522', 'license': 'Apache-2.0', 'n_tokens_mistral': 3036, 'n_tokens_neox': 2838, 'n_words': 825, 'path': 'tools/colab/autoencoder_embeddings.ipynb', 'repo_name': 'dreadlord1984/tpu', 'size': '8890'}