tf.function中if-else分支均执行的原因及解决方法

为什么tf.function中if-else两个分支都会执行？ #

• tf.function的核心是将Python代码转换为静态计算图，第一次调用时会触发"图追踪"过程。
• Python原生的if/else语句依赖Python标量值做判断，但tf.rank(t)返回的是TensorFlow张量（不是Python整数），Python无法在图追踪阶段解析这个张量的实际值，因此会将if和else两个分支的代码全部执行一遍，用于构建完整的计算图。
• 当实际运行时，比如输入rank=0的标量，else分支中tf.math.reduce_sum(..., axis=1)的操作会因为张量没有axis=1维度而抛出维度不匹配的错误。

如何让函数g在tf.function中正常工作？ #

需要使用TensorFlow原生的动态控制流操作tf.cond，它能基于张量条件动态选择执行分支，是计算图支持的分支逻辑。同时注意将Python的print替换为tf.print（前者在图追踪阶段执行，后者是图操作，仅在对应分支运行时执行）。

修改后的代码如下：

import tensorflow as tf

def g(t):
    rank = tf.rank(t)
    
    # 定义两个分支的逻辑
    def handle_rank0():
        tf.print('Rank=0')
        return tf.math.reduce_sum(tf.math.exp(t))
    
    def handle_higher_rank():
        tf.print('Rank=higher')
        return tf.math.reduce_sum(tf.math.exp(t), axis=1)
    
    # 使用tf.cond根据张量条件选择分支
    return tf.cond(tf.equal(rank, 0), handle_rank0, handle_higher_rank)

@tf.function
def Test(t):
    return g(t)

测试验证 #

• 调用Test(0.5)：输出Rank=0，返回tf.exp(0.5)的正确结果
• 调用Test(tf.constant([[1.0, 2.0], [3.0, 4.0]]))：输出Rank=higher，返回[tf.exp(1)+tf.exp(2), tf.exp(3)+tf.exp(4)]的正确结果

内容的提问来源于stack exchange，提问作者Jasper Rou