Tensorflow 谷歌云TPU——没有使用TPU
我正在尝试在TPU上运行一个简单的程序:Tensorflow 谷歌云TPU——没有使用TPU,tensorflow,google-cloud-platform,tpu,Tensorflow,Google Cloud Platform,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 =
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上
我做错了什么?策略。scope()
用于模型培训
如果要在TPU上运行tf.matmul
,可以使用以下任一选项:
with tf.device('/TPU:0'):
c = tf.matmul(a, b)
或
详细信息如下。能否提供有关如何运行此代码的更多上下文?我还建议您看看where is解释了如何使用TPUsit运行模型,结果证明我是在尝试做不可能的事情!根据文档,tf.Estimator API是在TPU上运行TF1.1X代码的唯一方法。
@tf.function
def matmul_fn(x, y):
z = tf.matmul(x, y)
return z
z = strategy.run(matmul_fn, args=(a, b))
print(z)