[英]Google Cloud TPU — no TPU being used
I am trying to run a simple program on 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())
When I run this, it looks like the TPU is being found.当我运行它时,似乎正在找到 TPU。 However, none of the logged devices have 'TPU' in the name -- it is all on the CPU.但是,所有记录的设备都没有名称中包含“TPU”——它们都在 CPU 上。
What am I doing wrong?我究竟做错了什么?
strategy.scope()
is for model training. strategy.scope()
用于 model 训练。
If you want to run tf.matmul
on a TPU you could use either this:如果您想在 TPU 上运行tf.matmul
,您可以使用以下任一方法:
with tf.device('/TPU:0'):
c = tf.matmul(a, b)
Or或者
@tf.function
def matmul_fn(x, y):
z = tf.matmul(x, y)
return z
z = strategy.run(matmul_fn, args=(a, b))
print(z)
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.