五分鐘學會在Colab上使用免費的TPU訓練模型

哈囉大家好，雖然忙碌，還是趁空擋想跟大家分享關於TensorFlow2.1.x系列的兩三事，一般來說做機器學習模型最需要的就是運算資源，而除了GPU之外，大家一定很想使用Google所推出的Google Cloud TPU來做機器學習模型，重點它很貴，能不能免費的使用他呢？使用Colab就是首選了。

而這次TF2.1新版推出之後，TPU使用有了一個很大幅度的改版，在Colab上使用TPU已經變得非常的容易與簡單，這邊就讓Jerry老師來跟大家娓娓道來。

首先不要忘記在Edit -> Notebook settings 要把硬體調整成TPU，如下圖：

然後就會獲得TPU資源，如下圖：

首先幾個在Colab上要跑TPU的API要了解一下：

1.tf.distribute.cluster_resolver.TPUClusterResolver()

這是負責偵測TPU，可以透過該指令了解Google TPU集群的狀況。

2.tf.config.experimental_connect_to_cluster(tpu)

這是負責連結Google雲端的TPU環境

3.tf.tpu.experimental.initialize_tpu_system(tpu)

這是負責初始化TPU的裝置

所以前置作業如下：

必須設定TensorFlow2.x 以及設定tpu的變數。

接下來要使用TPU，目前模型有固定的寫法，必須寫成函式，舉例如下：

資料的部分也是要寫成函式，這邊範例是採用mnist，所以舉例如下：

最後是訓練模型的程式碼，有個重點是

4.tf.distribute.experimental.TPUStrategy

是呼叫前面TPU分配策略的指令，所以要搭配在程式碼當中

5.strategy.scope()

搭配這段程式碼，可以把模型與要計算的優化器、衡量標準放到運算引擎上

如下：

目前實測的結果，如果以mnist資料為例，

TPU：36秒

GPU(P100)：90秒

CPU：291秒

使用心得是，程式碼不用改動太多，在TensorFlow 2.x的架構下，可以很快的進行訓練，程式碼也可以很快速地調整，以後大家就可以用前所未有的速度訓練模型囉！

當然如果您很想更近一步了解Google Cloud Platform上面的TPU費用，可以參考下面連結

https://cloud.google.com/tpu/docs/types-zones

https://cloud.google.com/tpu/pricing

相關程式碼請參考

https://github.com/jerrywu2013/TensorFlow_TPU