在 TensorFlow 中获取由另一个张量部分索引的切片的好方法是什么? [英] What is a nice way to obtain a slice partially indexed by another tensor in TensorFlow?

问题描述

假设我们有一个第一维未知的张量 x(例如 [?, 32, 32, 3])和另一个张量 i 实际上是一个标量.有没有一种很好的方法来获得 x 的 i 切片被第一维分割，例如，获得维度 [32, 32,3]?我是 TensorFlow 的新手，只能想出这个极其笨拙的解决方案.

Suppose we have a tensor x with unknown first dimension (for example [?, 32, 32, 3]), and another tensor i that is actually a scalar. Is there a nice way to obtain the i-th slice of x split by first dimension, for example, to get a tensor of dimension [32, 32, 3]? I'm new to TensorFlow and was only able to come up with this extremely clumsy solution.

index = tf.concat(0, [i, tf.constant([0, 0, 0], tf.int64)])
size = [1, x.get_shape()[1].value, x.get_shape()[2].value, x.get_shape()[3].value]
result = tf.unpack(tf.slice(x, index, size))[0]

推荐答案

您可以利用 -1 是 tf.slice() size 参数，意思是该维度中的所有剩余元素".然后，假设 i 是一个标量(而不是一个长度为 1 的向量，因为它似乎在你的代码片段中)，你可以这样做:

You can take advantage of the fact that -1 is a special argument to the tf.slice() size argument, meaning "all remaining elements in that dimension". Then, assuming i is a scalar (and not a length-1 vector as it seems to be in your code snippet), you can do:

result = tf.squeeze(tf.slice(x, tf.pack([index, 0, 0, 0]), [1, -1, -1, -1]), [0])

或者，您可以使用 tf.gather() 从第零维的张量中选择一个或多个切片.在这种情况下，i 必须是一个向量:

Alternatively, you can use tf.gather() to select one or more slices from a tensor on the zeroth dimension. In this case, i must be a vector:

i = tf.expand_dims(i, 0)  # Converts `i` to a vector if it is a scalar.
result = tf.squeeze(tf.gather(x, i), [0])

在这两种情况下，tf.Squeeze() op 去掉第 0 维得到一个三维结果.

In both cases, the tf.squeeze() op removes the 0th dimension to give a three-dimensional result.

这篇关于在 TensorFlow 中获取由另一个张量部分索引的切片的好方法是什么?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！