我如何在 tensorflow 1.x 中使用 colab 的 TPU 关于 tensorflow 2.x 的代码？

Question

tensorflow 2.x 代码如下：

将张量流导入为 tf

导入操作系统

将 tensorflow_datasets 导入为 tfds

解析器=tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://'+ os.environ['COLAB_TPU_ADDR'])

tf.config.experimental_connect_to_cluster(解析器)

tf.tpu.experimental.initialize_tpu_system(解析器)

策略 = tf.distribute.experimental.TPUStrategy(resolver)

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 tf.device('/TPU:0'):

c = tf.matmul(a, b)

print("c 设备：", c.device)

打印（c）

@tf.function

def matmul_fn(x, y):

z = tf.matmul(x, y)

返回 z

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

打印（z）

我的 1.x 代码如下：

%tensorflow_version 1.x

将张量流导入为 tf

导入操作系统

将 tensorflow_datasets 导入为 tfds

tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']

解析器 = tf.distribute.cluster_resolver.TPUClusterResolver(tpu_address)

tf.config.experimental_connect_to_cluster(解析器)

tf.tpu.experimental.initialize_tpu_system(解析器)

策略 = tf.distribute.experimental.TPUStrategy(resolver)

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 tf.device('/TPU:0'):

c = tf.matmul(a, b)

打印（c）

def matmul_fn(x, y):

z = tf.matmul(x, y)

返回 z

使用 tf.Session() 作为 sess：

使用 strategy.scope()：

z = strategy.experimental_run_v2(matmul_fn, args=(a, b))

打印(sess.run(z))

最后，我对在 colab 上的 tensorlfow 1.x 中使用 TPU 感到很困惑。

Answer 1

要在 TPU 上创建变量，您可以在 strategy.scope() 上下文管理器中创建它们。修正后的TensorFlow 2.x代码如下：

import tensorflow as tf
import os

resolver =tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://'+ os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
  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]])

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

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

这会在所有 TPU 副本上运行 tf.function 并给出以下结果：

PerReplica:{
  0: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  1: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  2: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  3: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  4: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  5: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  6: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  7: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)
}

要在 TF 1.x 中的 TPU 上评估相同的函数，您需要将 strategy.run 更改为 strategy.experimental_run_v2，在 TPU 上创建 tf.Session，并在值列表上调用 sess.run()由experimental_run_v2返回。

初始设置是相同的。将with 块更改为以下内容：

with strategy.scope():
  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]])

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

  with tf.Session(tpu_address) as sess:
    z = strategy.experimental_run_v2(matmul_fn, args=(a, b))
    print(sess.run(z.values))

这给出了以下结果：

(array([[22., 28.],
       [49., 64.]], dtype=float32), array([[22., 28.],
       [49., 64.]], dtype=float32), array([[22., 28.],
       [49., 64.]], dtype=float32), array([[22., 28.],
       [49., 64.]], dtype=float32), array([[22., 28.],
       [49., 64.]], dtype=float32), array([[22., 28.],
       [49., 64.]], dtype=float32), array([[22., 28.],
       [49., 64.]], dtype=float32), array([[22., 28.],
       [49., 64.]], dtype=float32))

我希望这能回答您的问题。有关在 TPU 上运行 TensorFlow 的更多信息，请参阅TensorFlow TPU guide。有关在不使用 Keras 的情况下使用分发策略的更多信息，请参阅guide on distribution strategies with custom training loops