{"cells":[{"metadata":{},"cell_type":"markdown","source":"I am writing this kernel mostly to explain myself how TPUs work. TPUs or Tensor Processing Units are accelerators that are specifically designed to carry out tasks related to deep learning with ease. However, just setting up TPU as the accelerator in one's Kaggle kernel is not sufficient. I, thus, try to perform a simple image classification problem on the Flowers dataset, as suggested by [Phil Culliton](https://www.kaggle.com/philculliton).","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true,"_kg_hide-output":false},"cell_type":"code","source":"import tensorflow as tf\nfrom kaggle_datasets import KaggleDatasets\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nprint(\"Tensorflow version \" + tf.__version__)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The following code requests the system for a TPU in the `try` block. If the TPU is available, well and good. If not, the system allocated either GPU or CPU to work with the code. Thus, it is essential to check or detect which accelerator is being used.","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Running on TPU ', tpu.master())\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"`REPLICAS: 8` lets us know that this kernel will be using something similar to 8 physical TPUs accross the board. Phil explains something extremely important for handling TPUs. It is that they require the data to be stored in a separate location, sort of co-mapped with the hardware of the TPU. In other words, they need to be stored in a container or \"bucket\", quite close to the TPU. So the dataset is shipped to the bucket next to the TPU. And hence, `KaggleDatasets()` transfers the dataset to the GCS (Google Cloud Services) bucket.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"GCS_DS_PATH = KaggleDatasets().get_gcs_path()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now that all pre-requisites are ready, we can carry out the deep learning tasks efficiently. The image size is kept to be `[192,192]`, at which a GPU will easily run out of memory. I set the number of generations to be 50 and the batch size to be 16 per hardware. This means that a specific number of pieces of data get sent to the TPU as a particular instance of time.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"IMAGE_SIZE = [192, 192]\nEPOCHS = 50\nBATCH_SIZE = 16 * strategy.num_replicas_in_sync\n\nNUM_TRAINING_IMAGES = 12753\nNUM_TEST_IMAGES = 7382\nSTEPS_PER_EPOCH = NUM_TRAINING_IMAGES // BATCH_SIZE\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"It is important to know that a single TPU board contains 4 TPU chips and each TPU chip contains 2 TPU cores. These cores are where all the matrix multiplications happen with regard to neural networks. These hardware structures are designed to carry out matrix multiplications on extremely large datasets, but at the same time, very fast. This is why they are ideal for using in Deep Learning since we can thus iterate our models for more time if we train them quite fast.","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"class\": tf.io.FixedLenFeature([], tf.int64),\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    return image, label\n\ndef read_unlabeled_tfrecord(example):\n    UNLABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"id\": tf.io.FixedLenFeature([], tf.string),\n    }\n    example = tf.io.parse_single_example(example, UNLABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    idnum = example['id']\n    return image, idnum\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False\n\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO)\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls=AUTO)\n    return dataset\n\ndef get_training_dataset():\n    dataset = load_dataset(tf.io.gfile.glob(GCS_DS_PATH + '/tfrecords-jpeg-192x192/train/*.tfrec'), labeled=True)\n    dataset = dataset.repeat()\n    dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_validation_dataset():\n    dataset = load_dataset(tf.io.gfile.glob(GCS_DS_PATH + '/tfrecords-jpeg-192x192/val/*.tfrec'), labeled=True, ordered=False)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_test_dataset(ordered=False):\n    dataset = load_dataset(tf.io.gfile.glob(GCS_DS_PATH + '/tfrecords-jpeg-192x192/test/*.tfrec'), labeled=False, ordered=ordered)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ntraining_dataset = get_training_dataset()\nvalidation_dataset = get_validation_dataset()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"However, that does not mean that we can take all data and dump it onto the TPU. There ought to be a meticulous handling of this data and a proper way in which data is delivered. This is made possible thanks to `tfrecords` which takes, in this case, the key pixels from the image and shoves it into a file for the TPU and sends it from the bucket like so. This helps in the maximum utilisation of the TPU, mostly because, if the way in which data is delivered is not optimised, it becomes a bottleneck for the efficient usage of the TPU.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"Now, we train the model using a pre-trained model which further helps in saving unnecessary time.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():    \n    pretrained_model = tf.keras.applications.VGG16(weights='imagenet', include_top=False ,input_shape=[*IMAGE_SIZE, 3])\n    pretrained_model.trainable = False\n    \n    model = tf.keras.Sequential([\n        pretrained_model,\n        tf.keras.layers.GlobalAveragePooling2D(),\n        tf.keras.layers.Dense(104, activation='softmax')\n    ])\n        \nmodel.compile(\n    optimizer='adam',\n    loss = 'sparse_categorical_crossentropy',\n    metrics=['sparse_categorical_accuracy']\n)\n\nhistorical = model.fit(training_dataset, \n          steps_per_epoch=STEPS_PER_EPOCH, \n          epochs=EPOCHS, \n          validation_data=validation_dataset)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"plt.plot(historical.history['loss'])\nplt.plot(historical.history['val_loss'])\nplt.xlabel('Epochs')\nplt.ylabel('Loss')\nplt.legend(['Loss', 'Validation Loss'])","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"plt.plot(historical.history['sparse_categorical_accuracy'])\nplt.plot(historical.history['val_sparse_categorical_accuracy'])\nplt.xlabel('Epochs')\nplt.ylabel('Sparse Categorical Accuracy')\nplt.legend(['Accuracy', 'Validation Accuracy'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"And so, there we go, in around 200 seconds, the model loss and accuracy are converged.","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"test_ds = get_test_dataset(ordered=True)\n\nprint('Computing predictions')\ntest_images_ds = test_ds.map(lambda image, idnum: image)\nprobabilities = model.predict(test_images_ds)\npredictions = np.argmax(probabilities, axis=-1)\nprint(predictions)\n","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}