【Python-torch】两个tensor相同元素对应位置求和

【Python-torch】两个tensor相同元素对应位置加权求和：一个很简单的需求：tensor A中元素为 tensor b中元素对应的 label，想把 label 相同的b中元素，除上这个label下元素的和

1. 需求

• A 为lable，B为一个tensor，

一个很简单的需求：tensor A 中元素为 tensor b中元素对应的 label，想把 label 相同的b中元素，除上这个label下元素的和。

• 这个需求很多地方其实都有用到，最多用的是图注意力网络attention那块，求注意力因子
• 网上找了很多地方，都没找到。
• 终于找到了。

2. 需要安装包

• 需要安装torcch-scatter，可以参考：安装torcch-scatter系列

3. 解决代码

import torch
from torch_scatter import scatter_sum

A = torch.tensor([1, 2, 0, 4, 2])

B = torch.tensor([2, 4, 6, 8, 10], dtype=torch.float32)

# 将 B 中的元素按照标签 A 进行求和
sum_result = scatter_sum(B, A, dim=0)

print(sum_result)

# 创建一个新的张量，用于存储除以元素和后的结果
C = B / sum_result[A]

print(C)