如何在TensorFlow的tf.function中从函数集合随机选取函数调用

问题原因 #

tf.function会将Python代码转换为TensorFlow静态计算图，图构建阶段Python原生列表的索引必须是编译期可确定的静态值，不能传入动态生成的tf.Tensor类型索引，因此会抛出类型错误。

解决方案 #

方案1：使用tf.switch_case实现动态分支（适用于函数逻辑差异大的场景） #

TensorFlow提供的tf.switch_case原生支持按Tensor类型的索引动态选择执行分支，是该需求的标准实现方式。
示例代码如下：

import tensorflow as tf

# 你的24个旋转函数保持原有定义即可
@tf.function
def func1(tensor):
    # 替换为实际旋转逻辑
    return tensor

@tf.function
def func2(tensor):
    # 替换为实际旋转逻辑
    return tensor

# ... 省略func3到func24的定义

@tf.function
def apply(tensor):
    list_of_funcs = [func1, func2, ..., func24]
    # 生成0~23的随机整数，注意maxval设为24（左闭右开区间）
    rand_idx = tf.random.uniform(shape=[], minval=0, maxval=24, dtype=tf.int32)
    # 构造分支函数列表，必须加f=cur_func的默认参数避免闭包晚绑定问题
    branch_fns = [lambda f=cur_func: f(tensor) for cur_func in list_of_funcs]
    # 按随机索引执行对应分支函数
    return tf.switch_case(branch_index=rand_idx, branch_fns=branch_fns)

注意事项

lambda表达式中必须添加f=cur_func的默认参数，否则会触发Python闭包的晚绑定特性，导致所有分支最终都调用列表最后一个函数。

方案2：参数化逻辑替换分支（适用于函数仅参数不同的场景，更高效） #

如果你的24个旋转函数只是参数不同、逻辑完全一致，可以直接将所有旋转参数预存为Tensor常量，随机采样参数后直接计算，完全避免分支开销，运行效率更高。
示例代码如下：

# 预存24种旋转对应的参数，比如3x3旋转矩阵，形状为[24, 3, 3]
ROTATION_MATS = tf.constant([
    # 填入你24种旋转对应的旋转矩阵参数
])

@tf.function
def apply_rotation(tensor):
    # 随机采样参数索引
    rand_idx = tf.random.uniform(shape=[], minval=0, maxval=24, dtype=tf.int32)
    # 取出对应旋转参数
    current_rot = ROTATION_MATS[rand_idx]
    # 执行旋转操作，替换为你实际的计算逻辑即可
    return tf.tensordot(tensor, current_rot, axes=1)

内容的提问来源于stack exchange，提问作者Matt Lyon