python - Tensorflow 可以进行条件计算吗？

Question

我正在使用 MxNet 开发条件计算框架。假设我们的小批量中有 N 个样本。我需要使用伪代码在我的计算图中执行此类操作:

x = graph.Variable("x")
y = graph.DoSomeTranformations(x)
# The following operation generates a Nxk sized matrix, k responses for each sample.
z = graph.DoDecision(y)
for i in range(k):
   argmax_sample_indices_for_i = graph.ArgMaxIndices(z, i)
   y_selected_samples = graph.TakeSelectedSample(y, argmax_sample_indices_for_i )
   result = graph.DoSomeTransformations(y_selected_samples)

我想要实现的目标如下:获得 y 后，我应用决策函数(这可以是 D 到 k 全连接层，其中 D 是数据维度)并为我的每个样本获得 k 个激活N 大小的小批量。然后，我想根据每个样本的最大激活的列索引，动态地将我的小批量分成 k 个不同的部分(k 可以是 2、3、一个小数字)。我假设的“graph.ArgMaxIndices”函数执行此操作，给定 z、Nxk 大小的矩阵和 i，该函数查找沿 i 列给出最大激活的样本索引并返回它们的索引。 (请注意，我寻找与“graph.ArgMaxIndices”给出等效结果的任何系列或函数组合，而不是单个函数)。最后，对于每个 i，我选择具有最大激活值的样本并对它们应用特定的转换。目前，据我所知，MxNet 在其符号网络中不支持此类条件计算。因此，我在每次决策后构建单独的符号图，并且必须为每个小批量分割编写单独的簿记 - 条件图结构，这会产生 1) 维护和开发非常复杂且繁琐的代码 2) 训练和评估期间运行性能下降。

我的问题是，我可以使用 Tensorflow 的符号运算符来完成上述操作吗？它是否允许人们根据一项标准选择小批量的子集？是否有一个函数或一系列函数相当于上面伪代码中的“graph.ArgMaxIndices”？ (给定 Nxk 矩阵和列索引 i，返回在第 k 列具有最大激活值的行索引)。

Answer 1

您可以在 Tensorflow 中做到这一点。

我认为最好的方法是使用面具和 tf.boolean_mask k 次，第 i 个掩码由 tf.equal(i, tf.argmax(z, axis=-1)) 给出

x = graph.Variable("x")
y = graph.DoSomeTranformations(x)
# The following operation generates a Nxk sized matrix, k responses for each sample.
z = graph.DoDecision(y)
max_indices = tf.argmax(z, axis=-1)
for i in range(k):
   argmax_sample_indices_for_i = tf.equal(i, max_indices)
   y_selected_samples = tf.boolean_mask(y, mask=argmax_sample_indices_for_i )
   result = graph.DoSomeTransformations(y_selected_samples)