PyG:从预测矩阵中删除现有边

Question

我目前正在使用PyG开发一个推荐系统。边的定义如下:

edge_index = tensor([[   0,    0,    0,  ..., 9315, 9317, 9317],
[ 100,  448,  452,  ...,  452,    1,  307]], device='cuda:0')}

edge_index[0]包含一个学生的索引，edge_index[1]包含连接模块的索引(长度相同)。因此，edge_index[0][i]为i的源节点，edge_index[1][i]为i的目的节点。

经过模型训练后，我生成一个二维张量recs，形状为# of Students x # of Modules，其值为0-1。0 =不推荐，1 =推荐。recs可能看起来像这样:

recs = tensor([0.54, 0.23, 0.98, ..., 0.12, 0.43, 0.87],
...,
[0.43, 0.53, 0.12, ..., 0.92, 0.12, 0.53])

当然，我不想推荐一个学生已经学过的模块。是否有一种方法可以将原始图的所有边设置为零，通过使用PyG的edge_index作为坐标或其他东西?

基本上我想把recs中的特定值设置为0，如下所示:

for i in range(0, len(edge_index[0])):
recs[edge_index[0][i]][edge_index[1][i]] = 0

是否有一种方法使用张量函数来实现这一点?

Answer 1

由于您希望同时在两个轴上索引recs，因此直接实现是将for循环矢量化为:

>>> recs[edge_index[0], edge_index[1]] = 0

可以改善通过分裂edge_indextuple:

>>> recs[tuple(edge_index)] = 0