基于条件向量将特定输入转发到不同的图形操作

Question

我正在尝试在TensorFlow中创建一个计算图，能够根据指定要使用的操作的辅助输入向量将输入数据的不同部分转发到不同的操作。

例如，假设您有一个输入张量 A ，其尺寸为（3,2,10,10）。我想要做的是根据向量 c = [1,0,1]将输入数据 A 转发到特定操作。特别是，第一行 A [0] 应转发到与向量 c 中的值1关联的操作，第二行 A [1] 应转发到与值0相关联的操作，依此类推。

假设每个操作都在函数中定义，其参数应如下所示：

def build_op_1(input, param1, param2, param3):
    # returns tensor operation for op_1
    pass

def build_op_2(input, param1, param2, param3, param4):
    # returns tensor operation for op_2
    pass

其中输入是 A 的第i行。如您所见，我需要为输入的第i行检索正确的操作并在其上执行。必须对 A 中包含的所有行执行此操作。

我已经看到有些操作能够管理图中的条件流（即 case ， switch 等），但我认为它们不是能够将特定数据转发到特定操作。有没有办法在TensorFlow中实现这种行为？

提前谢谢你，

的Alessandro

Answer 1

我相信你想使用tensorflow map_fn迭代行，再加上tensorflow cond来有条件地运行不同的子图。