【PyTorch】torch.gather() 用法

gather常被用于image做mask的操作中，对哪些地方进行赋值0/1

API：

torch.gather --- PyTorch 2.2 documentation

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

gather()的意义：

顾名思义，聚集、集合：gather的作用是根据索引index查找，然后将在input中的查找结果返回。

1. input：要被查找的tensor
2. dim：要进行查找的维度
3. index：要进行查找的索引

返回的维度和index的维度一致

例子：

例1：

此时假设input就是src，dim = 0，index就是index，最后的输出是result

dim = 0 的时候，按照列进行查找：

index的1在src的第一列中找到1，result的第1行的第1列就是1

index的2在src的第二列中找到5，result的第1行的第2列就是5

index的3在src的第三列中找到9，result的第1行的第3列就是9

....

例2：

此时假设input就是src，dim = 1，index就是index，最后的输出是result

dim = 1 的时候，按照行进行查找

pytorch中的gather()函数_.gather()-CSDN博客