tensorflow2.0之tf.function and auto-graph

tf.function可以将python式的函数转换成tensorflow 能运行的函数

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)
    return scale * tf.where(is_positive, z, alpha * tf.nn.elu(z))


print(scaled_elu(tf.constant(-3.)))
print(scaled_elu(tf.constant([-3., -2.5])))

scaled_elu_tf = tf.function(scaled_elu)
print(scaled_elu_tf(tf.constant(-3.)))
print(scaled_elu_tf(tf.constant([-3., -2.5])))

print(scaled_elu_tf.python_function is scaled_elu)

可以将scaled_elu转换成tf运行的代码

这样做的方法是速度运行速度快，特别是在GPU加速的tensorflow运行的比一般python函数更快

%timeit scaled_elu(tf.random.normal((1000, 1000)))
%timeit scaled_elu_tf(tf.random.normal((1000, 1000)))

#运行时间
#25.4 ms ± 1.75 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
#23.5 ms ± 454 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

第二种使用tf.function

# 1 + 1/2 + 1/2^2 + ... + 1/2^n

@tf.function
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))

这样也可以将python函数声明成tensorflow可以运行的函数

同样可以将python代码直接转换成tensorflow的graph函数，使用tf.autigraph.to_code()

def display_tf_code(func):
    code = tf.autograph.to_code(func)
    from IPython.display import display, Markdown
    display(Markdown('```python\n{}\n```'.format(code)))

我们可以使用在tf.function 中使用input_signature设置输入函数参数的类型

@tf.function(input_signature=[tf.TensorSpec([None], tf.int32, name='x')])
def cube(z):
    return tf.pow(z, 3)

try:
    print(cube(tf.constant([1., 2., 3.])))
except ValueError as ex:
    print(ex)
    
print(cube(tf.constant([1, 2, 3])))

input_signature的第二个作用是：使用tf.function转换而来的python函数只有包含input_signature才能被保存为savemodel