我有两个以下形状的张量:
tensor1 => shape(10, 99, 106)
tensor2 => shape(10, 99)
tensor2
包含0 - 105
范围内的值,我希望用它来切片tensor1
的最后一个维度并获取形状的tensor3
tensor3 => shape(10, 99, 99)
我尝试过使用:
tensor4 = tf.gather(tensor1, tensor2)
# this causes tensor4 to be of shape (10, 99, 99, 106)
另外,使用
tensor4 = tf.gather_nd(tensor1, tensor2)
# gives the error: last dimension of tensor2 (which is 99) must be
# less than the rank of the tensor1 (which is 3).
我正在寻找类似于numpy的cross_indexing的东西。
答案 0 :(得分:1)
您可以使用tf.map_fn
:
tensor3 = tf.map_fn(lambda u: tf.gather(u[0],u[1],axis=1),[tensor1,tensor2],dtype=tensor1.dtype)
您可以将此行视为在tensor1
和tensor2
的第一维上运行的循环,并且对于它们应用的第一个维度中的每个索引i
{{1在tf.gather
和tensor1[i,:,:]
。