【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)