时间戳

首页 > 解决方案 > 谷歌云 TPU——没有使用 TPU

tensorflow - 谷歌云 TPU——没有使用 TPU

问题描述

我正在尝试在 TPU 上运行一个简单的程序：

import tensorflow as tf

tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
print("Device:", tpu.master())
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)

a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

with strategy.scope():
    c = tf.matmul(a, b)
    print("c device: ", c.device)
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
    print(c.eval())

当我运行它时，似乎正在找到 TPU。但是，所有记录的设备的名称中都没有“TPU”——它们都在 CPU 上。

我究竟做错了什么？

标签： tensorflowgoogle-cloud-platformtpu

解决方案

strategy.scope()用于模型训练。

如果您想tf.matmul在 TPU 上运行，您可以使用以下任一方法：

with tf.device('/TPU:0'):
  c = tf.matmul(a, b)

或者

@tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)

详情在这里。

推荐阅读