python - 使用 Cloud TPU 节点对 Tensorboard 进行预测调用分析

Question

我一直在尝试使用 Cloud TPU v2-8 节点分析自定义神经网络模型的预测调用。

重要的是，我的预测调用大约需要 2 分钟才能完成，并且我使用按 TFRecord 批处理划分的数据来完成此操作。

我按照官方文档“Profile your model with Cloud TPU Tools”并 try catch 配置文件:

1. 使用 Tensorboard UI和
2. “programatic way ”带有 tf.profiler.experimental.start() 和 tf.profilier.experimental.stop() 包装预测调用，但在这两种情况下我都没有成功。

# TPU Node connection is done before...

# TPU at this point is already running
logdir_path = "logs/predict"
tf.profiler.experimental.start(logdir_path)
# Tensorflow predict call here
tf.profiler.experimental.stop()

我可以在两种情况下生成一些数据(Tensorboard UI 和探查器调用)，但是当我尝试在指向 logdir 路径的 Tensorboard 中打开它时，我收到一条“当前数据集没有事件的仪表板”消息。

是否有任何方法可以使用在 Cloud TPU 节点中运行的模型来分析 Tensorflow/Keras 预测调用？



奇怪的事实 - Tensorflow 文档和 Cloud TPU 文档似乎存在不一致:在 Tensorflow Optimization Docs 中我们可以看到，TPU 硬件不支持 tf.profiler.experimental.start/stop 调用，但在 Google Cloud docs 中这是在 TPU 中捕获配置文件的推荐方法。

配置:

• tensorflow 2.6.1
• 张量板2.9.1
• Python 3.8
• 云 TPU 节点 v2-8

Answer 1

1. 请检查日志目录中的跟踪文件。如果它们太小，您在跟踪过程中可能会遇到一些问题。
2. 请确保您输入了正确的命令。 $ tensorboard --logdir logs/predict
3. 使用 tf.profiler.experimental.client.start(...) 尝试另一种分析方法，如 TF profiler Docs 所示。 下面是代码片段。

import tensorflow as tf
from threading import Thread

def call_trace(tpu_resolver):  # This should be called asynchronously
  # a profiler service has been started in the TPU worker at port 8466
  service_addr = ":".join(tpu_resolver.get_master().split(":")[:-1] +
                          ["8466"])  # need to change for TPU pod
  tf.profiler.experimental.client.trace(service_addr=service_addr,
                                        logdir="gs://your_logdir",
                                        duration_ms=5000)

tpu_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(...)
# Other initialization codes

thr = Thread(target=call_trace, args=(tpu_resolver,))
thr.start()
# Codes you want to execute on the cloud TPU node
thr.join()

然后打开tensorboard进行可视化。

$ tensorboard --logdir gs://your_logdir