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

Question

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

例如，假设您有一个输入张量 a 其尺寸为（3、2、10、10）。我要做的是根据向量 c = [1，0，1]将输入数据 a 转发到特定操作。特别是，应将第一行 a [0] 转发到与vector 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 的第三行。如您所见，我需要检索输入第三行的正确操作并在其上执行。必须对 a 中包含的所有行进行此操作。

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

预先感谢您，

Alessandro

Answer 1

我相信您想将tensorflow map_fn用于迭代行，并与tensorflow cond结合以有条件地运行不同的子图。