在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模型中使用这些操作，因此想在图形模式下评估它们的性能。

首选方法是什么？

我的谷歌搜索一无所获，我不知道如何使用像 tf 1.x 中的会话。

Answer 1

您要查找的是tf.function。查看本教程和此文档。

正如教程所说，在 TensorFlow 2 中，默认情况下打开了预先执行。用户界面直观而灵活(运行一次性操作要容易得多，速度也快得多(，但这可能会以牺牲性能和可部署性为代价。要获得高性能和可移植的模型，请使用 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))