【colab】白嫖Google的tpu和gpu

先打开Google云盘.

在关联应用里搜索 Colaboratory 安装插件。

安装完成后就可以直接在Google云盘里创建colab了。

然后在 ‘代码执行程序’ 里找 ‘更改运行时类型’， 可以选择实在gpu上跑还是tpu。
（这里的gpu是Tesla T4，因为 gpu 的话网上资料一大把，这里接下来已 tpu 为例）

可以通过一下代码将你的 Google云盘挂载到 notebook上
（这里会给你一个网址，打开后是Google云盘的授权码，粘贴下来，因为这里已经授权过了所以下图中没有提示，直接挂载成功。）

from google.colab import drive
drive.mount('/content/drive')
root_PATH = './drive/My Drive/'

接下来是调用tpu的代码。

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

strategy = tf.distribute.TPUStrategy(resolver)

然后在构造模型的时候加上 with strategy.scope()，就成功将模型放到tpu上去了。

with strategy.scope():
  model = tf.keras.Sequential([
        tf.keras.layers.LSTM(units=256),
        tf.keras.layers.Dense(units=3, activation='softmax')
      ])

  model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4), loss='categorical_crossentropy', metrics=['accuracy'])

这里我自己的电脑上的显卡是2080，Google分配的tpu是8颗核心
相同的模型，两者一块跑，2080快跑到70轮时tpu已经快500轮了
不得不说tpu是真的快。