定义和运行基准测试

本指南介绍了如何定义和运行 TensorFlow 基准测试。这些基准测试以 TestResults 格式存储输出。这些基准测试添加到 TensorFlow GitHub 库后,我们会每天对我们的连续版本运行这些测试,并在我们的信息中心显示一张图:https://benchmarks-dot-tensorflow-testing.appspot.com/。

定义基准测试

定义 TensorFlow 基准测试时,需要从 tf.test.Benchmark 类进行扩展,并调用 self.report_benchmark 方法。例如,请看下面的示例基准测试代码:

import time

import tensorflow as tf

# Define a class that extends from tf.test.Benchmark.
class SampleBenchmark(tf.test.Benchmark):

  # Note: benchmark method name must start with `benchmark`.
  def benchmarkSum(self):
    with tf.Session() as sess:
      x = tf.constant(10)
      y = tf.constant(5)
      result = tf.add(x, y)

      iters = 100
      start_time = time.time()
      for _ in range(iters):
        sess.run(result)
      total_wall_time = time.time() - start_time

      # Call report_benchmark to report a metric value.
      self.report_benchmark(
          name="sum_wall_time",
          # This value should always be per iteration.
          wall_time=total_wall_time/iters,
          iters=iters)

if __name__ == "__main__":
  tf.test.main()

查看 SampleBenchmark 的完整示例。

以上示例中要注意的要点:

  • 基准测试类从 tf.test.Benchmark 扩展而来。
  • 每个基准测试方法都应以 benchmark 前缀开头。
  • 基准测试方法会调用 report_benchmark 来报告指标值。

使用 Python 运行基准测试

通过 --benchmarks 标记,可以使用 Python 运行基准测试。将打印一个 BenchmarkEntries 原型。

python sample_benchmark.py --benchmarks=SampleBenchmark

也可以将该标记设为 --benchmarks=.--benchmarks=all

(请确保已安装 Tensorflow,以便成功导入 import tensorflow as tf 一行中的程序包。如需安装说明,请参阅安装 TensorFlow。使用 bazel 运行基准测试时,可以不执行此步骤。)

添加 bazel 目标

对于在 TensorFlow Github 库中定义的基准测试,我们会添加一个名为 tf_py_logged_benchmark 的特殊目标。tf_py_logged_benchmark 应封装一个常规的 py_test 目标。运行 tf_py_logged_benchmark 将会打印一个 TestResults 原型。定义 tf_py_logged_benchmark 也可以让我们对 TensorFlow 连续版本运行该测试。

首先,定义一个常规的 py_test 目标。请参见以下示例:

py_test(
  name = "sample_benchmark",
  srcs = ["sample_benchmark.py"],
  srcs_version = "PY2AND3",
  deps = [
    "//tensorflow:tensorflow_py",
  ],
)

您可以通过传递 --benchmarks 标记,在 py_test 目标中运行基准测试。该基准测试应该只会打印一个 BenchmarkEntries 原型。

bazel test :sample_benchmark --test_arg=--benchmarks=all

现在请添加 tf_py_logged_benchmark 目标(如果有)。该目标会将 --benchmarks=all 传递给封装的 py_test 目标,并为我们的 TensorFlow 连续版本提供一种存储输出的方式。TensorFlow 代码库中应该有 tf_py_logged_benchmark 目标。

load("//tensorflow/tools/test:performance.bzl", "tf_py_logged_benchmark")

tf_py_logged_benchmark(
    name = "sample_logged_benchmark",
    target = "//tensorflow/examples/benchmark:sample_benchmark",
)

使用以下命令运行该基准测试目标:

bazel test :sample_logged_benchmark