我有一个形状为[7, 7, 2, 4]的tensor A和一个形状为[7, 7]的tensor B.
Tensor B是tensor A的argmax，它的值是0,1.
我想从A和B得到形状为[7, 7, 4]或[7, 7, 1, 4]的tensor C分.
规则是，tensor B的(i, j)个元素是tensor A的第二个维度的索引.
我怎么才能快点做呢？我试着在A[B]之前得到C，但不起作用.有谁能帮我吗？谢谢.
我有一个形状为[7, 7, 2, 4]的tensor A和一个形状为[7, 7]的tensor B.
Tensor B是tensor A的argmax，它的值是0,1.
我想从A和B得到形状为[7, 7, 4]或[7, 7, 1, 4]的tensor C分.
规则是，tensor B的(i, j)个元素是tensor A的第二个维度的索引.
我怎么才能快点做呢？我试着在A[B]之前得到C，但不起作用.有谁能帮我吗？谢谢.
好吧，我用了tf.gather_nd来解决这个问题:
tensor_C = tf.gather_nd(tensor_A, tf.expand_dims(tf.argmax(tensor_B, 2), 2), batch_dims=3)