在 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))