tensorflow 中的 tf.function 是否优化了 运行 时间?

Does tf.function in tensorflow optimize run time?

我读到使用 tf. function 可以通过创建图表和 lazy evaluation 优化 运行 时间。以下是进行矩阵乘法的示例代码:

import tensorflow as tf

def random_sample(x1,x2):
    return tf.matmul(x1,x2)

@tf.function
def random_sample_optimized(x1,x2):
    return tf.matmul(x1,x2)

x1 = tf.constant(tf.random.normal(shape=(3999,29999)))
x2 = tf.constant(tf.random.normal(shape=(29999,3999)))

正在计算 运行 时间:

import time
start = time.time()
op1 = random_sample(x1,x2)
end = time.time()
print (end-start) ## op ~avg = 7 secs

start = time.time()
op2 = random_sample_optimized(x1,x2)
end = time.time()
print (end-start) ##op ~avg = 9.5 secs

不仅使用 tf.function 时平均值高,而且每个 运行 上的个体 run_time 在使用 tf.function 时也很高。

关于我是否正确使用 tf.function 或仅在构建复杂神经网络时提供优化的任何建议?

对于复杂的计算,tf.function 通常会 运行 更快。引用书中的以下几行 - 'Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow' 以及示例:

A TF Function will usually run much faster than the original Python function, especially if it performs complex computations. However, in this trivial example, the computation graph is so small that there is nothing at all to optimize, so tf_cube() actually runs much slower than cube().

def cube(x):
    return x ** 3

>>> cube(2)
8
>>> cube(tf.constant(2.0))
<tf.Tensor: id=18634148, shape=(), dtype=float32, numpy=8.0>

现在,让我们使用 tf.function() 将此 Python 函数转换为 TensorFlow 功能:

>>> tf_cube = tf.function(cube)
>>> tf_cube
<tensorflow.python.eager.def_function.Function at 0x1546fc080>

>>> tf_cube(2)
<tf.Tensor: id=18634201, shape=(), dtype=int32, numpy=8>
>>> tf_cube(tf.constant(2.0))
<tf.Tensor: id=18634211, shape=(), dtype=float32, numpy=8.0>