本文通过实例解析PyTorch中的torch.gather函数，该函数用于根据指定维度上的索引从输入张量中获取值并创建新张量。以一个二维张量为例，展示了如何使用该函数按照行索引选取元素，常用于深度学习中的类别概率最大值获取。
PyTorch中的torch.gather函数，可以帮助我们在指定维度上按照索引从输入张量中获取对应的值，并输出一个新的张量。这个函数的操作可能比较抽象，下面我就举一个简单的例子来说明其用法。
我们假设有一个二维张量input，其形状为(4, 3)，其中每个元素都是一个整数。我们需要从input中按照行的索引获取对应的值，并将它们组成一个新的张量output。具体来说，我们需要按照以下方式获取：
output[i][j] = input[index[i][j]
订阅专栏 解锁全文
扫一扫
评论
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
查看更多评论
添加红包
