标签：index gather 函数 dim torch input out

建议先阅读官方文档，拿笔跟着给出的公式推导一次。

torch.gather官方文档

gather函数的定义为：

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor

不常用的暂时不关注，于是函数常使用的样子如下：

torch.gather(input, dim, index)

函数的大致功能，给出input，根据dim和index确认从input中取出的数据内容，和最终输出的shape

1.确定输出的shape，输出的shape跟index的shape一致

例如index是一个2x3的tensor,那么输出就是一个2x3的tensor

2.确定输出的tensor的内容：根据input,dim，index三者确定

dim确定的是在input中取数据的时候，那个维度是从index中取数据来确定索引，而不是直接顺序索引：

举例

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0

out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1

out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

out[i][j] = input[index[i][j][k]][j]  # if dim == 0

out[i][j] = input[i][index[i][j][k]]  # if dim == 1

其实这个函数一般是结合其他的函数使用的，比如torch.max，torch.min。比如调用max函数如下：

torch.max(a,dim=dim_num)

那么此时返回的是一个tuple，包含了value和index_list

那么有value = torch.gather(a, dim=dim_num, index=index_list)

标签：index,gather,函数,dim,torch,input,out
来源： https://blog.csdn.net/mxlmhgzw/article/details/120913625

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。