通过 Tensorflow 中的索引张量对张量进行切片 [英] Slicing a tensor by an index tensor in Tensorflow

问题描述

我有两个以下张量(请注意，它们都是两者 Tensorflow 张量，这意味着在我启动以下切片操作之前，它们实际上仍然是象征性的tf.Session()):

I have two following tensors (note that they are both Tensorflow tensors which means they are still virtually symbolic at the time I construct the following slicing op before I launch a tf.Session()):

• params:具有形状 (64,784, 256)
• indices:具有形状 (64, 784)

• params: has shape (64,784, 256)
• indices: has shape (64, 784)

我想构造一个返回以下张量的操作:

and I want to construct an op that returns the following tensor:

• output:具有形状 (64,784) 其中

• output: has shape (64,784) where

output[i,j] = params_tensor[i,j, indices[i,j] ]

Tensorflow 中最有效的方法是什么?

What is the most efficient way in Tensorflow to do so?

ps:我尝试使用 tf.gather 但无法使用它来执行我上面描述的操作.

ps: I tried with tf.gather but couldn't make use of it to perform the operation I described above.

非常感谢.
-最好的

Many thanks.
-Bests

推荐答案

您可以使用 tf.gather_nd 获得您想要的东西.最后的表达是:

You can get exactly what you want using tf.gather_nd. The final expression is:

tf.gather_nd(params, tf.stack([tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[0]), 1), [1, tf.shape(indices)[1]]), tf.transpose(tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[1]), 1), [1, tf.shape(indices)[0]])), indices], 2))

<小时>

这个表达式有以下解释:

---

This expression has the following explanation:

• tf.gather_nd 执行您的预期并使用索引从参数中收集输出

• tf.stack 结合了三个独立的张量，最后一个是索引.前两个张量指定前两个维度的排序(参数/索引的轴 0 和轴 1)

• tf.gather_nd does what you expected and uses the indices to gather the output from the params

• tf.stack combines three separate tensors, the last of which is the indices. The first two tensors specify the ordering of the first two dimensions (axis 0 and axis 1 of params/indices)

• 对于提供的示例，此顺序对于轴 0 来说只是 0, 1, 2, ..., 63，对于轴 1 来说是 0, 1, 2, ... 783.这些序列是通过 获得的分别为 tf.range(tf.shape(indices)[0]) 和 tf.range(tf.shape(indices)[1]).

• 对于提供的示例，索引的形状为 (64, 784).上面最后一点的另外两个张量需要具有相同的形状才能与 tf.stack

• 首先，使用 tf.expand_dims 为两个序列中的每一个添加一个额外的维度/轴.

• tf.tile 和 tf.transpose 的使用可以通过例子来说明:假设 params 和 index 的前两个轴的形状为 (5,3).我们希望第一个张量是:

• First, an additional dimension/axis is added to each of the two sequences using tf.expand_dims.

• The use of tf.tile and tf.transpose can be shown by example: Assume the first two axes of params and index have shape (5,3). We want the first tensor to be:

[[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]]

我们希望第二个张量是:

We want the second tensor to be:

[[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]

这两个张量的功能几乎就像为相关索引指定网格中的坐标一样.

These two tensors almost function like specifying the coordinates in a grid for the associated indices.

tf.stack 的最后一部分在新的第三个轴上组合了三个张量，因此结果与参数具有相同的 3 个轴.

The final part of tf.stack combines the three tensors on a new third axis, so that the result has the same 3 axes as params.

请记住，如果您的轴比问题中的轴多或少，则需要相应地修改 tf.stack 中坐标指定张量的数量.

Keep in mind if you have more or less axes than in the question, you need to modify the number of coordinate-specifying tensors in tf.stack accordingly.

这篇关于通过 Tensorflow 中的索引张量对张量进行切片的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！