Skip to content

Instantly share code, notes, and snippets.

@weakish

weakish/keras_lstm_tpu.ipynb

Created November 14, 2018 05:55
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save weakish/fe409ca4777395e31b176d12eb094b64 to your computer and use it in GitHub Desktop.
Save weakish/fe409ca4777395e31b176d12eb094b64 to your computer and use it in GitHub Desktop.
Download ZIP
Keras_LSTM_TPU.ipynb （single core）
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/fe409ca4777395e31b176d12eb094b64/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": "ff5e45e1-9022-4100-ecd0-b87afc970739"
},
"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": "2f7554d7-fb8c-48e1-8300-c587cdbb2658",
"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": "d407a429-3ed6-4d05-b2e6-d7f81927bcc5",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 590
}
},
"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",
" using_single_core=True))\n",
"\n",
"tpu_model.summary()"
],
"execution_count": 6,
"outputs": [
{
"output_type": "stream",
"text": [
"INFO:tensorflow:Querying Tensorflow master (b'grpc://10.39.177.130: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, 4761422944197289032)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 5845057563547798946)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_GPU:0, XLA_GPU, 17179869184, 15360918415596695747)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 10522857121329185991)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 9686434936540751613)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 16757935518776931547)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 14817231923316463660)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 10084428055787949492)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 3779979576175479097)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 12729353096811186431)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 85260730582304979)\n",
"INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 9422536406695747486)\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": "7ebbe3e3-6b78-44ae-de28-51a24402e3cc",
"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,\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": 7,
"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 1), [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: 1.903439998626709 secs\n",
"INFO:tensorflow:Setting weights on TPU model.\n",
"24832/25000 [============================>.] - ETA: 0s - loss: 0.6935 - acc: 0.5122INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 1), [TensorSpec(shape=(40,), dtype=tf.int32, name='core_id0'), TensorSpec(shape=(40, 500), dtype=tf.int32, name='Input_10'), TensorSpec(shape=(40, 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.82780385017395 secs\n",
"24960/25000 [============================>.] - ETA: 0s - loss: 0.6932 - acc: 0.5129INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 1), [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: 3.1653573513031006 secs\n",
"INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 1), [TensorSpec(shape=(8,), dtype=tf.int32, name='core_id_10'), TensorSpec(shape=(8, 500), dtype=tf.int32, name='Input_10'), TensorSpec(shape=(8, 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.031578302383423 secs\n",
"25000/25000 [==============================] - 41s 2ms/step - loss: 0.6933 - acc: 0.5130 - val_loss: 0.6680 - val_acc: 0.5820\n",
"Epoch 2/20\n",
"25000/25000 [==============================] - 19s 766us/step - loss: 0.5423 - acc: 0.7329 - val_loss: 0.3356 - val_acc: 0.8742\n",
"Epoch 3/20\n",
"25000/25000 [==============================] - 19s 763us/step - loss: 0.3064 - acc: 0.8768 - val_loss: 0.2149 - val_acc: 0.9148\n",
"Epoch 4/20\n",
"25000/25000 [==============================] - 19s 766us/step - loss: 0.2209 - acc: 0.9162 - val_loss: 0.1828 - val_acc: 0.9358\n",
"Epoch 5/20\n",
"25000/25000 [==============================] - 19s 766us/step - loss: 0.1749 - acc: 0.9338 - val_loss: 0.1929 - val_acc: 0.9286\n",
"Epoch 6/20\n",
"25000/25000 [==============================] - 19s 764us/step - loss: 0.1328 - acc: 0.9535 - val_loss: 0.0922 - val_acc: 0.9668\n",
"Epoch 7/20\n",
"25000/25000 [==============================] - 19s 766us/step - loss: 0.0957 - acc: 0.9689 - val_loss: 0.0625 - val_acc: 0.9802\n",
"Epoch 8/20\n",
"25000/25000 [==============================] - 19s 766us/step - loss: 0.0701 - acc: 0.9759 - val_loss: 0.0361 - val_acc: 0.9898\n",
"Epoch 9/20\n",
"25000/25000 [==============================] - 19s 765us/step - loss: 0.0499 - acc: 0.9841 - val_loss: 0.0243 - val_acc: 0.9916\n",
"Epoch 10/20\n",
"25000/25000 [==============================] - 19s 766us/step - loss: 0.0365 - acc: 0.9882 - val_loss: 0.0208 - val_acc: 0.9938\n",
"Epoch 11/20\n",
"25000/25000 [==============================] - 19s 760us/step - loss: 0.0250 - acc: 0.9920 - val_loss: 0.0135 - val_acc: 0.9958\n",
"Epoch 12/20\n",
"25000/25000 [==============================] - 19s 765us/step - loss: 0.0189 - acc: 0.9937 - val_loss: 0.0075 - val_acc: 0.9980\n",
"Epoch 13/20\n",
"25000/25000 [==============================] - 19s 760us/step - loss: 0.0163 - acc: 0.9946 - val_loss: 0.0100 - val_acc: 0.9974\n",
"Epoch 14/20\n",
"25000/25000 [==============================] - 19s 766us/step - loss: 0.0141 - acc: 0.9955 - val_loss: 0.0062 - val_acc: 0.9986\n",
"Epoch 15/20\n",
"25000/25000 [==============================] - 19s 769us/step - loss: 0.0138 - acc: 0.9958 - val_loss: 0.0025 - val_acc: 0.9996\n",
"Epoch 16/20\n",
"25000/25000 [==============================] - 19s 765us/step - loss: 0.0110 - acc: 0.9970 - val_loss: 0.0041 - val_acc: 0.9986\n",
"Epoch 17/20\n",
"25000/25000 [==============================] - 19s 768us/step - loss: 0.0090 - acc: 0.9972 - val_loss: 0.0073 - val_acc: 0.9980\n",
"Epoch 18/20\n",
"25000/25000 [==============================] - 19s 764us/step - loss: 0.0107 - acc: 0.9966 - val_loss: 0.0082 - val_acc: 0.9974\n",
"Epoch 19/20\n",
"25000/25000 [==============================] - 19s 771us/step - loss: 0.0093 - acc: 0.9969 - val_loss: 0.0088 - val_acc: 0.9980\n",
"Epoch 20/20\n",
"25000/25000 [==============================] - 19s 764us/step - loss: 0.0105 - acc: 0.9969 - val_loss: 0.0034 - val_acc: 0.9988\n",
"INFO:tensorflow:Copying TPU weights to the CPU\n",
"--- 408.33573031425476 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": "cc7a8030-edda-48a5-ed10-9bebcad4b087",
"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": 8,
"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": "ee222ee7-ed29-4900-b4c5-edb3fe77ca9f",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 52
}
},
"cell_type": "code",
"source": [
"inferencing_model.evaluate(x_test, y_test)"
],
"execution_count": 9,
"outputs": [
{
"output_type": "stream",
"text": [
"25000/25000 [==============================] - 72s 3ms/step\n"
],
"name": "stdout"
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"[0.8827069875001907, 0.85304]"
]
},
"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