torch.gather函数

最新推荐文章于 2024-05-29 01:07:50 发布

zexalzz

最新推荐文章于 2024-05-29 01:07:50 发布

阅读量504

点赞数

文章标签： python pytorch Powered by 金山文档

本文链接：https://blog.csdn.net/qq_49863991/article/details/128710146

版权

文章介绍了torch.gather函数的工作原理，它是如何根据指定索引从矩阵中选取元素的。在推荐系统库Recbole中的应用示例展示了如何利用此函数选择用户交互过的热门物品，并计算recall。pos_matrix记录用户-物品交互，topk_idx选择热门物品的索引，pos_idx和pos_len_list用于计算召回率。

摘要由CSDN通过智能技术生成

1. 基本描述

基于对官方文档的理解，gather是通过输入一组序列矩阵来查找输入矩阵中对应的元素，最后返回对应元素组成的矩阵，因此序列矩阵和输入矩阵大小应一致。查找的方式可以通过参数 dim(int) 决定。

import torch
t = torch.tensor([[1, 2], [3, 4]])
torch.gather(t, dim=1, index=torch.tensor([[0, 0], [1, 0]])) # dim=1, 按行索引，[0, 0]返回第一行第一个元素, [1, 0]返回第二行第二个元素和第二行第一个元素

tensor([[1, 1],

[4, 3]])

2. 实际应用

在recbole中collector.py中有

pos_matrix = torch.zeros_like(scores_tensor, dtype=torch.int)
pos_matrix[positive_u, positive_i] = 1
pos_len_list = pos_matrix.sum(dim=1, keepdim=True) # 对每一行数据求和
pos_idx = torch.gather(pos_matrix, dim=1, index=topk_idx)
result = torch.cat((pos_idx, pos_len_list), dim=1)

pos_matrix是用户-物品矩阵，标为1的位置代表用户和物品有过交互

topk_idx代表选出的top个物品的序号

现在从pos-matrix中挑选出热门物品

pos_idx表示挑选出的热门物品，pos_len_list记录了每个用户交互过的物品数；pos_idx中的为1的数量除以pos_len_list记录的物品数即为recall