我們在網上隻發現比較少的資訊與資源，最開始介紹 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.BatchNormalization(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
model.add(tf.keras.layers.Dropout(0.25))
model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
model.add(tf.keras.layers.Dropout(0.25))
model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
model.add(tf.keras.layers.Dropout(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.Dropout(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.AdamOptimizer(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-7c00961fe d69