Google Colab TPU版本

如何解决Google Colab TPU版本

如何在Google Colab中打印我正在使用哪个TPU版本以及TPU有多少内存？

随着我得到以下输出

result = []
@gear
def consumer():

    data = once(val=[255],tout=Uint[10])
    res = data | Fixp[11,16]
    collect(res,result=result)


consumer()
sim()

输出

tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)

tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)

我想要类似INFO:tensorflow:Initializing the TPU system: grpc://10.123.109.90:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.123.109.90:8470 INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Finished initializing TPU system. INFO:tensorflow:Finished initializing TPU system. WARNING:absl:`tf.distribute.experimental.TPUStrategy` is deprecated,please use the non experimental symbol `tf.distribute.TPUStrategy` instead. INFO:tensorflow:Found TPU system: INFO:tensorflow:Found TPU system: INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:cpu:0,cpu,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:cpu:0,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_cpu:0,XLA_cpu,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:cpu:0,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0,TPU,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYstem:0,TPU_SYstem,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_cpu:0,0) 之类的东西。

我希望有人能帮助我。期待得到答案！

解决方法

This answer应该满足您的需求。

---

import os
from tensorflow.python.profiler import profiler_client

tpu_profile_service_address = os.environ['COLAB_TPU_ADDR'].replace('8470','8466')
print(profiler_client.monitor(tpu_profile_service_address,100,2))

输出如下：

  Timestamp: 22:23:03
  TPU type: TPU v2
  Utilization of TPU Matrix Units (higher is better): 0.000%

TPUv2每核8GB，TPUv3每核16GB HBM（https://cloud.google.com/tpu）。