Colab提供了免費TPU，機器之心幫你試了試

我們在網上只發現比較少的資訊與資源，最開始介紹 Colab 免費 TPU 的內容還是谷歌開發者 Sam Wittevee 最近的演講 PPT。因此本文的測試和探索都是基於官方文件和例項所實現的，還有很多 TPU 特性沒有考慮到，感興趣的讀者可查閱文末的參考資料，瞭解更多 Colab 免費 TPU 的特性。

本文所有的測試程式碼與結果都可以訪問：https://colab.research.google.com/drive/1DpUCBm58fruGNRtQL_DiSVbT90spdZgm

試驗 Colab 免費 TPU

首先我們需要確保 Colab 筆記本中執行時型別選擇的是 TPU，同時分配了 TPU 資源。因此依次選擇選單欄中的「runtime」和「change runtime type」就能彈出以下對話方塊：

為了確保 Colab 給我們分配了 TPU 計算資源，我們可以執行以下測試程式碼。如果輸出 ERROR 項，則表示目前的執行時並沒有調整到 TPU，如果輸出 TPU 地址及 TPU 裝置列表，則表示 Colab 已經為我們分配了 TPU 計算資源。

如果檢視以下測試程式碼的正常輸出，Colab 會為「TPU 執行時」分配 CPU 和 TPU，其中分配的 TPU 工作站有八個核心，因此在後面配置的 TPU 策略會選擇 8 條並行 shards。

import os
import pprint
import tensorflow as tf

if 'COLAB_TPU_ADDR' not in os.environ:
print('ERROR: Not connected to a TPU runtime')
else:
tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print ('TPU address is', tpu_address)

with tf.Session(tpu_address) as session:
devices = session.list_devices()

print('TPU devices:')
pprint.pprint(devices)

目前，Colab 一共支援三種執行時，即 CPU、GPU（K80）和 TPU（據說是 TPU v2）。但我們不太瞭解 Colab 中的 GPU 和 TPU 在深度模型中的表現如何，當然後面會用具體的任務去測試，不過現在我們可以先用相同的運算試試它們的效果。因此我們首先嚐試用簡單的卷積運算測試它們的迭代時間。

在測試不同的硬體時，需要切換到不同的執行時。如下先定義 128 張隨機生成的 256×256 影象，然後定義 256 個 5×5 的卷積核後就能執行卷積運算，其中魔術函式 %timeit 會自動多次執行，以產生一個更為精確的平均執行時間。

import tensorflow as tf 
import numpy as np
import timeit

tf.reset_default_graph()
img = np.random.randn(128, 256, 256, 3).astype(np.float32)
w = np.random.randn(5, 5, 3, 256).astype(np.float32)
conv = tf.nn.conv2d(img, w, [1,2,2,1], padding='SAME')

with tf.Session() as sess:
# with tf.device("/gpu:0") as dev:
%timeit sess.run(conv)

然而，是我們想當然了，使用 TPU 執行運算似乎需要特定的函式與運算，它不像 CPU 和 GPU 那樣可以共用相同的程式碼。分別選擇 CPU、GPU 和 TPU 作為執行時狀態，執行上面的程式碼並迭代一次所需要的時間分別為：2.44 s、280 ms、2.47 s。從這裡看來，僅修改執行時狀態，並不會真正呼叫 TPU 資源，真正實現運算的還是 CPU。隨後我們發現 TF 存在一個神奇的類 tf.contrib.tpu，似乎真正呼叫 TPU 資源必須使用它改寫模型。

因此，根據文件與呼叫示例，我們將上面的卷積測試程式碼改為了以下形式，併成功地呼叫了 TPU。此外，因為每次都需要重新連線不同的執行時，所以這裡的程式碼都保留了庫的匯入。雖然程式碼不太一樣，但直覺上它的計算量應該和上面的程式碼相同，因此大致上能判斷 Colab 提供的 GPU、TPU 速度對比。

import tensorflow as tf 
import numpy as np
import timeit
import os

tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']

tf.reset_default_graph()
def conv_op():
img =np.random.randn(128, 256, 256, 3).astype(np.float32)
conv_w = np.random.randn(5, 5, 3, 256).astype(np.float32)
conv = tf.nn.conv2d(img, conv_w, [1,2,2,1], padding='SAME')

tpu_ops = tf.contrib.tpu.batch_parallel(conv_op, [], num_shards=8)

with tf.Session(tpu_address) as sess:
sess.run(tf.contrib.tpu.initialize_system())
sess.run(tpu_ops)
%timeit sess.run(tpu_ops)
sess.run(tf.contrib.tpu.shutdown_system())

執行後出現了非常意外的結果，這樣的卷積運算每一次迭代只需要 1.22 ms。如下圖所示，很可能存在變數快取等其它因素造成了一定程度的緩慢，但 TPU 的速度無可置疑地快。因此如果在 Colab 上測試模型，我們就更希望使用免費的 TPU，不過使用 TPU 需要改模型程式碼，這又比較麻煩。

儘管簡單的卷積運算 TPU 要比 K80 快很多，但這隻能給我們一個大致的猜想，因此我們需要測試完整的模型。注意在 tf.contrib.tpu 類中，它還提供了兩種使用 TPU 的簡單方法，即直接使用 Keras 介面和使用 TPUEstimator 構建模型。

在 tf.contrib.tpu 的文件中，我們發現 tf.contrib.tpu.keras_to_tpu_model 方法可以直接將 Keras 模型與對應的權重複制到 TPU，並返回 TPU 模型。該方法在輸入 Keras 模型和在多個 TPU 核心上的訓練策略後，能輸出一個 Keras TPU 模型的例項，且可分配到 TPU 進行運算。

除此之外，另外一種呼叫 TPU 計算資源的方法是 tf.contrib.tpu.TPUEstimator，對於修正我們原來的TensorFlow 模型以適用 TPU，它可能是一種更方便的方式。根據文件所示，TPUEstimator 類繼承自 Estimator 類，因此它不僅支援在 TPU 上運算，同時還支援 CPU 和 GPU 的運算。TPUEstimator 隱藏了非常多在 TPU 上訓練的細節，例如為多個 TPU 核心複製多個輸入和模型等。

TPU 呼叫文件地址：https://www.tensorflow.org/api_docs/python/tf/contrib/tpu

對比 TPU 與 GPU 的計算速度

為了簡單起見，這裡僅使用 Fashion-MNIST 資料集與簡單的 5 層 卷積神經網路 測試不同的晶片效能。這個模型是基於 Keras 構建的，因為除了模型轉換與編譯，Keras 模型在 TPU 和 GPU 的訓練程式碼都是一樣的，且用 Keras 模型做展示也非常簡潔。

幾天前谷歌 Colab 團隊發了一版使用 Keras 呼叫 TPU 的教程，因此我們就藉助它測試 TPU 的訓練速度。對於 GPU 的測試，我們可以修改該模型的編譯與擬合部分，並呼叫 GPU 進行訓練。所以整個訓練的資料獲取、模型結構、超引數都是一樣的，不一樣的只是硬體。

教程地址：https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/fashion_mnist.ipynb

以下是整個測試的公共部分，包含了訓練資料的獲取和模型架構。Keras 的模型程式碼非常好理解，如下第一個卷積層首先採用了批歸一化，然後用 64 個 5×5 的卷積核實現卷積運算，注意這裡採用的啟用函式都是指數線性單元（ELU）。隨後對卷積結果做 2×2 的最大池化，並加上一個隨機丟棄率為 0.25 的Dropout層，最後得出的結果就是第一個卷積層的輸出。

import tensorflow as tf
import numpy as np
import timeit

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

model = tf.keras.models.Sequential()

# 以下為第一個卷積層
model.add(tf.keras.layers.<mark data-type="technologies" data-id="e3995dcb-2c96-4671-b7ed-d5f2abae3c08">Batch</mark><mark data-type="technologies" data-id="a815978e-ab05-479e-b6b4-30176b3d5427">Normalization</mark>(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.Max<mark data-type="technologies" data-id="0a4cedf0-0ee0-4406-946e-2877950da91d">Pooling</mark>2D(pool_size=(2, 2), strides=(2,2)))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.25))

model.add(tf.keras.layers.<mark data-type="technologies" data-id="e3995dcb-2c96-4671-b7ed-d5f2abae3c08">Batch</mark><mark data-type="technologies" data-id="a815978e-ab05-479e-b6b4-30176b3d5427">Normalization</mark>(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.Max<mark data-type="technologies" data-id="0a4cedf0-0ee0-4406-946e-2877950da91d">Pooling</mark>2D(pool_size=(2, 2)))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.25))

model.add(tf.keras.layers.<mark data-type="technologies" data-id="e3995dcb-2c96-4671-b7ed-d5f2abae3c08">Batch</mark><mark data-type="technologies" data-id="a815978e-ab05-479e-b6b4-30176b3d5427">Normalization</mark>(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.Max<mark data-type="technologies" data-id="0a4cedf0-0ee0-4406-946e-2877950da91d">Pooling</mark>2D(pool_size=(2, 2), strides=(2,2)))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.25))

model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(256))
model.add(tf.keras.layers.Activation('elu'))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.5))
model.add(tf.keras.layers.Dense(10))
model.add(tf.keras.layers.Activation('softmax'))
model.summary()

在定義模型後，TPU 需要轉化模型與編譯模型。如下所示，keras_to_tpu_model 方法需要輸入正常 Keras 模型及其在 TPU 上的分散式策略，這可以視為「TPU 版」的模型。完成模型的轉換後，只需要像一般 Keras 模型那樣執行編譯並擬合數據就可以了。

注意兩個模型的超引數，如學習率、批量大小和 Epoch 數量等都設定為相同的數值，且損失函式和最優化器等也採用相同的方法。

import os
tpu_model = tf.contrib.tpu.keras_to_tpu_model(
model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
)
)
tpu_model.compile(
optimizer=tf.train.Adam<mark data-type="technologies" data-id="fa50298e-1a85-4af0-ae96-a82708f4b610">Optimizer</mark>(learning_rate=1e-3, ),
loss=tf.keras.losses.sparse_categorical_crossentropy,
metrics=['sparse_categorical_accuracy']
)


def train_gen(batch_size):
while True:
offset = np.random.randint(0, x_train.shape[0] - batch_size)
yield x_train[offset:offset+batch_size], y_train[offset:offset + batch_size]


%time tpu_model.fit_generator(train_gen(1024), epochs=5, steps_per_epoch=100, validation_data=(x_test, y_test))

最後在使用 GPU 訓練模型時，我們會刪除模型轉換步驟，並保留相同的編譯和擬合部分。訓練的結果如下所示，Colab 提供的 TPU 要比 GPU 快 3 倍左右，一般 TPU 訓練 5 個 Epoch 只需要 40 多秒，而 GPU 需要 2 分多鐘。

Colab 使用免費 TPU 訓練的資訊摘要。

Colab 使用免費 GPU 訓練的資訊摘要。

最後，Colab 確實提供了非常強勁的免費 TPU，而且使用 Keras 或 TPUEstimator 也很容易重新搭建或轉換已有的TensorFlow 模型。機器之心只是簡單地試用了 Colab 免費 TPU，還有很多特性有待讀者的測試，例如支援 TPU 的 PyTorch 1.0 或迴圈神經網路在 TPU 上的效能等。

參考資料：

• 文件：https://www.tensorflow.org/api_docs/python/tf/contrib/tpu

• 官方示例（Keras）：https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb

• 官方示例（TPUEstimator）：https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpuestimator.ipynb

• Sam Wittevee PPT：https://www.dropbox.com/s/jg7j07unw94wbom/TensorFlow%20Keras%20Colab%20TPUs.pdf?dl=0

• Ceshine Lee 部落格：https://medium.com/the-artificial-impostor/keras-for-tpus-on-google-colaboratory-free-7c00961fed69

相關資料

推薦文章

如何將模型部署到安卓移動端，這裡有一份簡單教程

用TensorFlow Estimator實現文字分類

詳解Wassertein GAN：使用Keras在MNIST上的實現