Python @tf.function decorator将函数编译成图形时会发生什么?为什么在“渴望”模式下比这更快?

Python @tf.function decorator将函数编译成图形时会发生什么?为什么在“渴望”模式下比这更快?,python,tensorflow,tensorflow2.0,Python,Tensorflow,Tensorflow2.0,根据TensorFlow,@tf.function将一个函数编译成一个图形,并“使您获得更快执行、在GPU或TPU上运行或导出到SavedModel的好处。” 本文件中的示例说明了这一好处: lstm_cell = tf.keras.layers.LSTMCell(10) @tf.function def lstm_fn(input, state): return lstm_cell(input, state) input = tf.zeros([10, 10]) state = [tf

根据TensorFlow,
@tf.function
将一个函数编译成一个图形,并“使您获得更快执行、在GPU或TPU上运行或导出到SavedModel的好处。”

本文件中的示例说明了这一好处:

lstm_cell = tf.keras.layers.LSTMCell(10)

@tf.function
def lstm_fn(input, state):
  return lstm_cell(input, state)

input = tf.zeros([10, 10])
state = [tf.zeros([10, 10])] * 2
# warm up
lstm_cell(input, state); lstm_fn(input, state)
print("eager lstm:", timeit.timeit(lambda: lstm_cell(input, state), number=10))
print("function lstm:", timeit.timeit(lambda: lstm_fn(input, state), number=10))
输出:

eager lstm: 0.032440788112580776 
function lstm: 0.004768412094563246

“编译图”和“急切模式下的函数”有什么区别?为什么前者在执行时更快?在图模式中

,TysFooFor构建一个表示您的模型的计算图,并通过会话将其转发给C++运行时。这提供了分布式培训的好处,并在过程中优化了计算图(通过不断折叠等)。它还简化了部署到独立于平台的服务器的过程

@tf.function decorator允许用户在急切模式下获得图形执行的好处,但请注意一些注意事项(例如,tf.function更喜欢tf-ops而不是Python原语)

资料来源:

在TF2.0中默认使用的“急切执行”会立即评估操作,而无需构建图形。另一方面,图形具有一定的灵活性。但是,执行时间将取决于特定的代码。越多的代码使用C++后端,差异就越小。所有的开销都来自使用python解释器的急切模式下的程序。例如,基本矩阵乘法运算不会像更复杂的DNN应用程序那样获得更多好处