在 tf 2.x 中以图形模式运行 TensorFlow op

Question

我想对一些 TensorFlow 操作进行基准测试（例如在它们之间或针对 PyTorch）。 但是大多数时候我会写一些类似的东西：

import numpy as np
import tensorflow as tf

tf_device = '/GPU:0'
a = np.random.normal(scale=100, size=shape).astype(np.int64)
b = np.array(7).astype(np.int64)

with tf.device(tf_device):
  a_tf = tf.constant(a)
  b_tf = tf.constant(b)

%timeit tf.math.floormod(a_tf, b_tf)

这种方法的问题在于它在急切模式下进行计算（我认为特别是它必须执行 GPU 到 CPU 放置）。 最终，我想在tf.keras model 中使用这些操作，因此想评估它们在图形模式下的性能。

首选的方法是什么？

我的谷歌搜索没有任何结果，我不知道如何使用 tf 1.x 中的会话。

Answer 1

您正在寻找的是tf.function 。 检查本教程和本文档。

正如教程所说，在 TensorFlow 2 中，Eager Execution 默认开启。 用户界面直观且灵活（一次性操作更容易、更快捷），但这可能会以牺牲性能和可部署性为代价。 要获得高性能和可移植的模型，请使用 tf.function 从您的程序中制作图表。

检查此代码：

import numpy as np
import tensorflow as tf
import timeit

tf_device = '/GPU:0'

shape = [100000]
a = np.random.normal(scale=100, size=shape).astype(np.int64)
b = np.array(7).astype(np.int64)

@tf.function
def experiment(a_tf, b_tf):
  tf.math.floormod(a_tf, b_tf)

with tf.device(tf_device):
  a_tf = tf.constant(a)
  b_tf = tf.constant(b)

# warm up
experiment(a_tf, b_tf)
print("In graph mode:", timeit.timeit(lambda: experiment(a_tf, b_tf), number=10))
print("In eager mode:", timeit.timeit(lambda: tf.math.floormod(a_tf, b_tf), number=10))