gather常被用于image做mask的操作中,对哪些地方进行赋值0/1
API:
torch.gather — PyTorch 2.2 documentation
torch.gather(input, dim, index, out=None) → Tensor
gather()的意义:
顾名思义,聚集、集合:gather的作用是根据索引index查找,然后将在input中的查找结果返回。
- input:要被查找的tensor
- dim:要进行查找的维度
- 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博客