Skip to content

Instantly share code, notes, and snippets.

@weakish

weakish/keras_lstm_tpu.ipynb

Created November 14, 2018 09:04
Show Gist options
  • Save weakish/94551bc0609a506e8808e788cfe067ac to your computer and use it in GitHub Desktop.
Save weakish/94551bc0609a506e8808e788cfe067ac to your computer and use it in GitHub Desktop.
Download ZIP
Keras_LSTM_TPU.ipynb (rerun)
Display the source blob
Display the rendered blob
Raw
keras_lstm_tpu.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Keras_LSTM_TPU.ipynb",
"version": "0.3.2",
"provenance": [],
"collapsed_sections": [],
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "TPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/weakish/94551bc0609a506e8808e788cfe067ac/keras_lstm_tpu.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"metadata": {
"id": "CB43mV-TD1vb",
"colab_type": "text"
},
"cell_type": "markdown",
"source": [
"# Tutorial - [How to train Keras model x20 times faster with TPU for free](https://www.dlology.com/blog/how-to-train-keras-model-x20-times-faster-with-tpu-for-free/)"
]
},
{
"metadata": {
"id": "ya06BE0ZU526",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
"import tensorflow as tf\n",
"from tensorflow.keras.datasets import imdb\n",
"from tensorflow.keras.preprocessing import sequence\n",
"from tensorflow.python.keras.layers import Input, LSTM, Bidirectional, Dense, Embedding"
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "_uSZchXTVOHr",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 52
},
"outputId": "7cb8ec24-8e74-4ef0-d04b-92af765ea75c"
},
"cell_type": "code",
"source": [
"\n",
"# Number of words to consider as features\n",
"max_features = 10000\n",
"# Cut texts after this number of words (among top max_features most common words)\n",
"maxlen = 500\n",
"\n",
"# Load data\n",
"(x_train, y_train), (x_test, y_test) = imdb.load_data(num_words=max_features)\n",
"\n",
"# Reverse sequences\n",
"x_train = [x[::-1] for x in x_train]\n",
"x_test = [x[::-1] for x in x_test]\n",
"\n",
"# Pad sequences\n",
"x_train = sequence.pad_sequences(x_train, maxlen=maxlen)\n",
"x_test = sequence.pad_sequences(x_test, maxlen=maxlen)"
],
"execution_count": 2,
"outputs": [
{
"output_type": "stream",
"text": [
"Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/imdb.npz\n",
"17465344/17464789 [==============================] - 0s 0us/step\n"
],
"name": "stdout"
}
]
},
{
"metadata": {
"id": "p35nSfjbVVBE",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
"def make_model(batch_size=None):\n",
" source = Input(shape=(maxlen,), batch_size=batch_size, dtype=tf.int32, name='Input')\n",
" embedding = Embedding(input_dim=max_features, output_dim=128, name='Embedding')(source)\n",
" # lstm = Bidirectional(LSTM(32, name = 'LSTM'), name='Bidirectional')(embedding)\n",
" lstm = LSTM(32, name = 'LSTM')(embedding)\n",
" predicted_var = Dense(1, activation='sigmoid', name='Output')(lstm)\n",
" model = tf.keras.Model(inputs=[source], outputs=[predicted_var])\n",
" model.compile(\n",
" optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01),\n",
" loss='binary_crossentropy',\n",
" metrics=['acc'])\n",
" return model\n",
" "
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "bivVZS0jZhxg",
"colab_type": "code",
"outputId": "2adda45d-53aa-44cc-c12d-abc98da284cb",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 278
}
},
"cell_type": "code",
"source": [
"tf.keras.backend.clear_session()\n",
"training_model = make_model(batch_size = 128)\n",
"training_model.summary()"
],
"execution_count": 4,
"outputs": [
{
"output_type": "stream",
"text": [
"_________________________________________________________________\n",
"Layer (type) Output Shape Param # \n",
"=================================================================\n",
"Input (InputLayer) (128, 500) 0 \n",
"_________________________________________________________________\n",
"Embedding (Embedding) (128, 500, 128) 1280000 \n",
"_________________________________________________________________\n",
"LSTM (LSTM) (128, 32) 20608 \n",
"_________________________________________________________________\n",
"Output (Dense) (128, 1) 33 \n",
"=================================================================\n",
"Total params: 1,300,641\n",
"Trainable params: 1,300,641\n",
"Non-trainable params: 0\n",
"_________________________________________________________________\n"
],
"name": "stdout"
}
]
},
{
"metadata": {
"id": "xeGmIQX2aJvw",
"colab_type": "code",
"outputId": "c406bac6-5d12-4161-b471-3620d921fe4f",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 610
}
},
"cell_type": "code",
"source": [
"import os\n",
"# This address identifies the TPU we'll use when configuring TensorFlow.\n",
"TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
"tf.logging.set_verbosity(tf.logging.INFO)\n",
"\n",
"tpu_model = tf.contrib.tpu.keras_to_tpu_model(\n",
" training_model,\n",
" strategy=tf.contrib.tpu.TPUDistributionStrategy(\n",
" tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))\n",
"\n",
"tpu_model.summary()"
],
"execution_count": 5,
"outputs": [
{
"output_type": "stream",
"text": [
"INFO:tensorflow:Querying Tensorflow master (b'grpc://10.67.53.98:8470') for TPU system metadata.\n",
"INFO:tensorflow:Found TPU system:\n",
"INFO:tensorflow:*** Num TPU Cores: 8\n",
"INFO:tensorflow:*** Num TPU Workers: 1\n",
"INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 11157655388029336326)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 17822332713669640552)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_GPU:0, XLA_GPU, 17179869184, 11393582042076572559)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 17555043945696025512)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 12242177310840364991)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 10430744784801167142)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 8513005328315714657)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 10337892639234683075)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 5952952897402739363)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 2691509858094533235)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 10401006571233641223)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 18251553439053879553)\n",
"WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.\n",
"_________________________________________________________________\n",
"Layer (type) Output Shape Param # \n",
"=================================================================\n",
"Input (InputLayer) (128, 500) 0 \n",
"_________________________________________________________________\n",
"Embedding (Embedding) (128, 500, 128) 1280000 \n",
"_________________________________________________________________\n",
"LSTM (LSTM) (128, 32) 20608 \n",
"_________________________________________________________________\n",
"Output (Dense) (128, 1) 33 \n",
"=================================================================\n",
"Total params: 1,300,641\n",
"Trainable params: 1,300,641\n",
"Non-trainable params: 0\n",
"_________________________________________________________________\n"
],
"name": "stdout"
}
]
},
{
"metadata": {
"id": "XlSm1vd5bteH",
"colab_type": "code",
"outputId": "fc074e70-3a3c-484b-9e0f-8c1ce8d18603",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1148
}
},
"cell_type": "code",
"source": [
"import time\n",
"start_time = time.time()\n",
"\n",
"history = tpu_model.fit(x_train, y_train,\n",
" epochs=20,\n",
" batch_size=128 * 8,\n",
" validation_split=0.2)\n",
"tpu_model.save_weights('./tpu_model.h5', overwrite=True)\n",
"\n",
"print(\"--- %s seconds ---\" % (time.time() - start_time))\n"
],
"execution_count": 6,
"outputs": [
{
"output_type": "stream",
"text": [
"Train on 25000 samples, validate on 5000 samples\n",
"Epoch 1/20\n",
"INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name='core_id0'), TensorSpec(shape=(128, 500), dtype=tf.int32, name='Input_10'), TensorSpec(shape=(128, 1), dtype=tf.float32, name='Output_target_30')]\n",
"INFO:tensorflow:Overriding default placeholder.\n",
"INFO:tensorflow:Remapping placeholder for Input\n",
"INFO:tensorflow:Started compiling\n",
"INFO:tensorflow:Finished compiling. Time elapsed: 2.7104790210723877 secs\n",
"INFO:tensorflow:Setting weights on TPU model.\n",
"23552/25000 [===========================>..] - ETA: 0s - loss: 0.6933 - acc: 0.4987INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(53,), dtype=tf.int32, name='core_id0'), TensorSpec(shape=(53, 500), dtype=tf.int32, name='Input_10'), TensorSpec(shape=(53, 1), dtype=tf.float32, name='Output_target_30')]\n",
"INFO:tensorflow:Overriding default placeholder.\n",
"INFO:tensorflow:Remapping placeholder for Input\n",
"INFO:tensorflow:Started compiling\n",
"INFO:tensorflow:Finished compiling. Time elapsed: 3.582460403442383 secs\n",
"24576/25000 [============================>.] - ETA: 0s - loss: 0.6933 - acc: 0.4974INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name='core_id_10'), TensorSpec(shape=(128, 500), dtype=tf.int32, name='Input_10'), TensorSpec(shape=(128, 1), dtype=tf.float32, name='Output_target_30')]\n",
"INFO:tensorflow:Overriding default placeholder.\n",
"INFO:tensorflow:Remapping placeholder for Input\n",
"INFO:tensorflow:Started compiling\n",
"INFO:tensorflow:Finished compiling. Time elapsed: 4.2674195766448975 secs\n",
"INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(113,), dtype=tf.int32, name='core_id_10'), TensorSpec(shape=(113, 500), dtype=tf.int32, name='Input_10'), TensorSpec(shape=(113, 1), dtype=tf.float32, name='Output_target_30')]\n",
"INFO:tensorflow:Overriding default placeholder.\n",
"INFO:tensorflow:Remapping placeholder for Input\n",
"INFO:tensorflow:Started compiling\n",
"INFO:tensorflow:Finished compiling. Time elapsed: 5.533883571624756 secs\n",
"25000/25000 [==============================] - 34s 1ms/step - loss: 0.6933 - acc: 0.4982 - val_loss: 0.6932 - val_acc: 0.5010\n",
"Epoch 2/20\n",
"25000/25000 [==============================] - 3s 127us/step - loss: 0.6932 - acc: 0.5011 - val_loss: 0.6934 - val_acc: 0.4938\n",
"Epoch 3/20\n",
"25000/25000 [==============================] - 3s 126us/step - loss: 0.6932 - acc: 0.5025 - val_loss: 0.6931 - val_acc: 0.5040\n",
"Epoch 4/20\n",
"25000/25000 [==============================] - 3s 125us/step - loss: 0.6930 - acc: 0.5098 - val_loss: 0.6948 - val_acc: 0.4938\n",
"Epoch 5/20\n",
"25000/25000 [==============================] - 3s 129us/step - loss: 0.6911 - acc: 0.5275 - val_loss: 0.7052 - val_acc: 0.4942\n",
"Epoch 6/20\n",
"25000/25000 [==============================] - 3s 130us/step - loss: 0.6930 - acc: 0.5746 - val_loss: 0.6602 - val_acc: 0.5996\n",
"Epoch 7/20\n",
"25000/25000 [==============================] - 3s 130us/step - loss: 0.6667 - acc: 0.6392 - val_loss: 0.6198 - val_acc: 0.6632\n",
"Epoch 8/20\n",
"25000/25000 [==============================] - 3s 129us/step - loss: 0.6193 - acc: 0.6846 - val_loss: 0.5424 - val_acc: 0.7636\n",
"Epoch 9/20\n",
"25000/25000 [==============================] - 3s 130us/step - loss: 0.5670 - acc: 0.7307 - val_loss: 0.4381 - val_acc: 0.8160\n",
"Epoch 10/20\n",
"25000/25000 [==============================] - 3s 131us/step - loss: 0.4917 - acc: 0.7830 - val_loss: 0.4848 - val_acc: 0.7656\n",
"Epoch 11/20\n",
"25000/25000 [==============================] - 3s 132us/step - loss: 0.4536 - acc: 0.8017 - val_loss: 0.3684 - val_acc: 0.8488\n",
"Epoch 12/20\n",
"25000/25000 [==============================] - 3s 131us/step - loss: 0.3902 - acc: 0.8357 - val_loss: 0.3333 - val_acc: 0.8598\n",
"Epoch 13/20\n",
"25000/25000 [==============================] - 3s 129us/step - loss: 0.3253 - acc: 0.8705 - val_loss: 0.5172 - val_acc: 0.7606\n",
"Epoch 14/20\n",
"25000/25000 [==============================] - 3s 129us/step - loss: 0.2829 - acc: 0.8894 - val_loss: 0.2152 - val_acc: 0.9224\n",
"Epoch 15/20\n",
"25000/25000 [==============================] - 3s 128us/step - loss: 0.2361 - acc: 0.9064 - val_loss: 0.2454 - val_acc: 0.8990\n",
"Epoch 16/20\n",
"25000/25000 [==============================] - 3s 129us/step - loss: 0.1912 - acc: 0.9269 - val_loss: 0.2027 - val_acc: 0.9144\n",
"Epoch 17/20\n",
"25000/25000 [==============================] - 3s 129us/step - loss: 0.1643 - acc: 0.9390 - val_loss: 0.1004 - val_acc: 0.9678\n",
"Epoch 18/20\n",
"25000/25000 [==============================] - 3s 131us/step - loss: 0.1237 - acc: 0.9550 - val_loss: 0.0662 - val_acc: 0.9780\n",
"Epoch 19/20\n",
"25000/25000 [==============================] - 3s 130us/step - loss: 0.0999 - acc: 0.9627 - val_loss: 0.0565 - val_acc: 0.9834\n",
"Epoch 20/20\n",
"25000/25000 [==============================] - 3s 130us/step - loss: 0.0687 - acc: 0.9768 - val_loss: 0.0294 - val_acc: 0.9938\n",
"INFO:tensorflow:Copying TPU weights to the CPU\n",
"--- 99.58249568939209 seconds ---\n"
],
"name": "stdout"
}
]
},
{
"metadata": {
"id": "tLjeqllMZzv5",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
"# history = tpu_model.fit(x_train, y_train,\n",
"# epochs=20,\n",
"# batch_size=128 * 8,\n",
"# validation_split=0.2)\n",
"# tpu_model.save_weights('./tpu_model.h5', overwrite=True)"
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "gpcLs6PYatf5",
"colab_type": "code",
"outputId": "df394c7b-a040-48e4-dca3-52c326042802",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 278
}
},
"cell_type": "code",
"source": [
"inferencing_model = make_model(batch_size=None)\n",
"inferencing_model.load_weights('./tpu_model.h5')\n",
"inferencing_model.summary()"
],
"execution_count": 7,
"outputs": [
{
"output_type": "stream",
"text": [
"_________________________________________________________________\n",
"Layer (type) Output Shape Param # \n",
"=================================================================\n",
"Input (InputLayer) (None, 500) 0 \n",
"_________________________________________________________________\n",
"Embedding (Embedding) (None, 500, 128) 1280000 \n",
"_________________________________________________________________\n",
"LSTM (LSTM) (None, 32) 20608 \n",
"_________________________________________________________________\n",
"Output (Dense) (None, 1) 33 \n",
"=================================================================\n",
"Total params: 1,300,641\n",
"Trainable params: 1,300,641\n",
"Non-trainable params: 0\n",
"_________________________________________________________________\n"
],
"name": "stdout"
}
]
},
{
"metadata": {
"id": "KZyqMpMecAGp",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
""
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "2cqqkN84dJQk",
"colab_type": "code",
"outputId": "20a5908b-0f12-485a-9828-2ab249df1e2e",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 52
}
},
"cell_type": "code",
"source": [
"inferencing_model.evaluate(x_test, y_test)"
],
"execution_count": 8,
"outputs": [
{
"output_type": "stream",
"text": [
"25000/25000 [==============================] - 73s 3ms/step\n"
],
"name": "stdout"
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"[0.7603418525123596, 0.80328]"
]
},
"metadata": {
"tags": []
},
"execution_count": 8
}
]
},
{
"metadata": {
"id": "XMR_TncEeFKx",
"colab_type": "code",
"outputId": "f2daa67c-562b-48b2-8e66-d0f22a5ba1d4",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 159
}
},
"cell_type": "code",
"source": [
"tpu_model.evaluate(x_test, y_test, batch_size=128 * 8)"
],
"execution_count": 9,
"outputs": [
{
"output_type": "stream",
"text": [
"24576/25000 [============================>.] - ETA: 0sINFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(53,), dtype=tf.int32, name='core_id_10'), TensorSpec(shape=(53, 500), dtype=tf.int32, name='Input_10'), TensorSpec(shape=(53, 1), dtype=tf.float32, name='Output_target_30')]\n",
"INFO:tensorflow:Overriding default placeholder.\n",
"INFO:tensorflow:Remapping placeholder for Input\n",
"INFO:tensorflow:Started compiling\n",
"INFO:tensorflow:Finished compiling. Time elapsed: 9.659732580184937 secs\n",
"25000/25000 [==============================] - 16s 628us/step\n"
],
"name": "stdout"
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"[0.7600244810867309, 0.8032800001716613]"
]
},
"metadata": {
"tags": []
},
"execution_count": 9
}
]
},
{
"metadata": {
"id": "EiQqm2R_DvIV",
"colab_type": "text"
},
"cell_type": "markdown",
"source": [
"## Download the trained model weights to your local file system"
]
},
{
"metadata": {
"id": "UzMykUsdDjiD",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
"from google.colab import files\n",
"\n",
"files.download('./tpu_model.h5')"
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "apwRGvwWDnau",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
""
],
"execution_count": 0,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment