Tensorflow 使用tf.function 一定提高计算效率吗？

最新推荐文章于 2021-06-17 17:13:35 发布

风华明远

最新推荐文章于 2021-06-17 17:13:35 发布

阅读量473

点赞数

文章标签：机器学习 tensorflow

本文链接：https://blog.csdn.net/weixin_42272768/article/details/105987852

版权

在很多的讲解中，都提到了在TF2.0中，使用tf.function会提高计算效率。针对神经网络计算，这个说法是正确的。但是针对纯矩阵运算，使用tf.function反而会降低计算效率。猜测是因为使用tf.function产生了额外的系统调用造成的。因为纯矩阵运算使用的是numpy的库。而tf.function使用Graph，一定会有额外的开销。这些开销不足以让矩阵运算提高效率。以上为猜测。


import tensorflow as tf
import timeit	

conv_layer = tf.keras.layers.Conv2D(100, 3)	
@tf.function	
def conv_fn(image):	
  return conv_layer(image)	
image = tf.zeros([1, 200, 200, 100])	
# warm up	
conv_layer(image); conv_fn(image)	
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=100))	
print("Function conv:", timeit.timeit(lambda: conv_fn(image)