tf.function:可以将普通的python函数转成tensorflow里的图
auto-graph:是tf.function所依赖的机制
将普通的python函数转成tensorflow里图结构的优势:
更快,tensorflow图结构里有很多的优化,所以一个普通的python函数转成tensorflow里的图实现的函数后,速度会提升。
除了使用tf.function来显式的把一个python函数转成tensorflow里的图,还可以用别的方法(添加@tf.function标注)把python函数转成tensorflow里的图,下面一一介绍。
1、 tf.function函数转换
核心代码展示:
# tf.function and auto-graph.
def scaled_elu(z, scale=1.0, alpha=1.0):
# z >= 0 ? scale * z : scale * alpha * tf.nn.elu(z)
is_positive = tf.greater_equal(z, 0.0) #判断z是否大于0
return scale * tf.where(is_positive, z, alpha * tf.nn.elu(z)) #用where实现三项表达式
print(scaled_elu(tf.constant(-3.)))
print(scaled_elu(tf.constant([-3., -2.5])))
#一行代码将python语法实现的函数变成tensorflow中图实现的函数
scaled_elu_tf = tf.function(scaled_elu)
print(scaled_elu_tf(tf.constant(-3.)))
print(scaled_elu_tf(tf.constant([-3., -2.5])))
#对于转化后的函数可以通过 找回它的python函数
print(scaled_elu_tf.python_function is scaled_elu)
%timeit scaled_elu(tf.random.normal((1000, 1000)))
%timeit scaled_elu_tf(tf.random.normal((1000, 1000)))
运行结果为:
tf.Tensor(-0.95021296, shape=(), dtype=float32) tf.Tensor([-0.95021296 -0.917915 ], shape=(2,), dtype=float32) tf.Tensor(-0.95021296, shape=(), dtype=float32) tf.Tensor([-0.95021296 -0.917915 ], shape=(2,), dtype=float32) True 946 µs ± 12.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) 838 µs ± 23.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
从上面的运行结果可以看出:1、普通python函数的运行结果与转换成tensorflow图实现的函数的运行结果相同
2、一个普通的python函数转成tensorflow里的图实现的函数后,速度会提升
2、@tf.function函数转换
核心代码展示:
# 1 + 1/2 + 1/2^2 + ... + 1/2^n
#@tf.function #标注:将下面的函数转成tensorflow里的图
def converge_to_2(n_iters): #迭代次数
total = tf.constant(0.)
increment = tf.constant(1.)
for _ in range(n_iters):
total += increment
increment /= 2.0
return total
print(converge_to_2(20))
未使用@tf.function输出结果为:
tf.Tensor(1.9999981, shape=(), dtype=float32)
# 1 + 1/2 + 1/2^2 + ... + 1/2^n
@tf.function #标注:将下面的函数转成tensorflow里的图
def converge_to_2(n_iters): #迭代次数
total = tf.constant(0.)
increment = tf.constant(1.)
for _ in range(n_iters):
total += increment
increment /= 2.0
return total
print(converge_to_2(20))
使用@tf.function后输出结果为:
tf.Tensor(1.9999981, shape=(), dtype=float32)
综上:用@tf.function也可以将一个普通的python函数转成tensorflow里的图实现的函数。
此外,使用to_code方法可以将python函数转成可以生成tensorflow图结构的中间代码展示出来,如下:
def display_tf_code(func): #传进来的是一个python函数
"""to_code函数可以将普通的Python函数转成tensorflow的代码,tensorflow的代码可以生成tensorflow里的图结构。"""
code = tf.autograph.to_code(func)
#展示tensorflow的代码
from IPython.display import display, Markdown
display(Markdown('```python\n{}\n```'.format(code)))
display_tf_code(scaled_elu)
输出结果为(输出的是转化成的tensorflow代码,scaled_elu函数在前面有定义):
def tf__scaled_elu(z, scale=None, alpha=None): do_return = False retval_ = ag__.UndefinedReturnValue() with ag__.FunctionScope('scaled_elu', 'scaled_elu_scope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as scaled_elu_scope: is_positive = ag__.converted_call(tf.greater_equal, scaled_elu_scope.callopts, (z, 0.0), None, scaled_elu_scope) do_return = True retval_ = scaled_elu_scope.mark_return_value(scale * ag__.converted_call(tf.where, scaled_elu_scope.callopts, (is_positive, z, alpha * ag__.converted_call(tf.nn.elu, scaled_elu_scope.callopts, (z,), None, scaled_elu_scope)), None, scaled_elu_scope)) do_return, return ag__.retval(retval_)
#使用变量,初始化必须在函数外面
var = tf.Variable(0.)
@tf.function
def add_21():
return var.assign_add(21) # # 给变量+21+=
print(add_21())
输出结果:
tf.Tensor(21.0, shape=(), dtype=float32)