tensorflow2.0之tf.function and auto-graph

ゞ浴缸里的玫瑰 2023-03-13 12:25 13阅读 0赞

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