而這次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費用,可以參考下面連結
相關程式碼請參考
網誌管理員已經移除這則留言。
回覆刪除