相关疑难解决方法(0)

我有一个logits尺寸张量[batch_size, num_rows, num_coordinates](即批次中的每个logit都是一个矩阵).在我的情况下,批量大小为2,有4行和4个坐标.

logits = tf.constant([[[10.0, 10.0, 20.0, 20.0],
                      [11.0, 10.0, 10.0, 30.0],
                      [12.0, 10.0, 10.0, 20.0],
                      [13.0, 10.0, 10.0, 20.0]],
                     [[14.0, 11.0, 21.0, 31.0],
                      [15.0, 11.0, 11.0, 21.0],
                      [16.0, 11.0, 11.0, 21.0],
                      [17.0, 11.0, 11.0, 21.0]]])

我想选择第一批的第一行和第二行以及第二批的第二行和第四行.

indices = tf.constant([[0, 1], [1, 3]])

所以期望的输出就是

logits = tf.constant([[[10.0, 10.0, 20.0, 20.0],
                      [11.0, 10.0, 10.0, 30.0]],
                     [[15.0, 11.0, 11.0, 21.0],
                      [17.0, 11.0, 11.0, 21.0]]])

如何使用TensorFlow执行此操作？我尝试使用tf.gather(logits, indices)但它没有返回我的预期.谢谢!