如何使用 tf.gather 从第一个轴切片值？答案

【问题标题】：How can I use tf.gather to slice values from the first axis?如何使用 tf.gather 从第一个轴切片值？
【发布时间】：2021-06-02 15:24:02
【问题描述】：

我希望下面的块返回一个形状为(128,64) 或(1,128,64) 的张量，因为我告诉它从第一个轴收集值并且我

tf.gather(h_states # shape: (3,128,64)
          ,indices # shape: (128), values are integers between 0 and 2
          ,axis = 0
          ,batch_dims=0)

相反，它返回一个 (128,128,64) 张量。我究竟做错了什么？我怎样才能让它从第一个轴中选择子张量？

谢谢

【问题讨论】：

您希望返回一个形状为 (128,64) 的张量。您想从 (3,128,64) 张量中获取一行吗？
是（如果“行”指的是第一个维度）。我想得到由索引张量确定的行。
您可以简单地对张量进行切片，例如：yourTensor[ indexYouWant , : , : ]

标签： python-3.x tensorflow

【解决方案1】：

你不必使用gather方法，对张量进行切片非常简单。

yourTensor[ indexYouWant , : , : ]                 # shape(128,64)
yourTensor[ indexYouWant:indexYouWant+1 , : , : ]  # shape(1,128,64)

如果你想要多行（例如：索引0和2），你可以使用gather方法：

tf.gather(yourTensor , indices = [0,2] , axis = 0)

【讨论】：