图模式下的tensorflow函数(tf.function)在返回对象时速度较慢

Question

我有一个用@tf.function包装的函数，当它不返回任何对象时，这个函数很快。但当它返回一个对象时，速度会明显变慢。

有什么方法可以提高性能，即使它必须返回一个对象？

下面是一个玩具示例

from tqdm import tqdm
import tensorflow as tf
import numpy as np
model = tf.keras.Sequential([
tf.keras.layers.Input((4,)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10000)
])
model.compile(
optimizer="rmsprop",
loss="sparse_categorical_crossentropy",
metrics=["sparse_categorical_accuracy"],
)
np.random.seed(0)
xy = tf.data.Dataset.from_tensor_slices((np.random.random((200000,4)), np.random.randint(1000, size=(200000,1)))).batch(200)
@tf.function
def _no_return(batch):
x, y = batch
y_pred = model(x, training=False)
r = tf.math.top_k(y_pred, k=20).indices
# True indicate that item is the correct prediction
r = tf.cast(tf.equal(r, tf.expand_dims(tf.cast(y, tf.int32), 1)), tf.float32)
# rank of the correct prediction, rank = 9999999+1 if no correction prediction within topk
r = tf.add((tf.reduce_sum(r, 1)-1) * -9999999, tf.cast(tf.argmax(r, 1) + 1, tf.float32))
@tf.function
def _return(batch):
x, y = batch
y_pred = model(x, training=False)
r = tf.math.top_k(y_pred, k=20).indices
# True indicate that item is the correct prediction
r = tf.cast(tf.equal(r, tf.expand_dims(tf.cast(y, tf.int32), 1)), tf.float32)
# rank of the correct prediction, rank = 9999999+1 if no correction prediction within topk
r = tf.add((tf.reduce_sum(r, 1)-1) * -9999999, tf.cast(tf.argmax(r, 1) + 1, tf.float32))
return r
#
# this run 1428.27 iterations/s 
#
for batch in tqdm(xy):
_no_return(batch)
#
# this run 135.61 iterations/s 
#
for batch in tqdm(xy):
_return(batch)

Answer 1

这是因为autograph。在_return函数的情况下，因为返回了r值，所以它正在生成完整的图(以r作为叶节点(。但是，在_no_return函数的情况下，没有返回值，签名不将r设置为叶节点，并且因为没有叶节点，所以整个图中没有任何内容

因此，你观察到这种加速是因为函数没有计算任何东西。

添加任何叶节点(如下面的演示所示(

def _no_return(batch):
# your old code
tf.print(r) # this will create another leaf node, 
# or you can try
tf.summary.scalar(r)

在CCD_ 8和CCD_。此外，如果您删除@tf.function装饰器，那么您也将看到相同的性能。

让我知道你是否还有任何疑问。