torch_scatter.scatter详解

最新推荐文章于 2025-06-27 14:16:18 发布

原创最新推荐文章于 2025-06-27 14:16:18 发布 · 1.5w 阅读

31 ·

CC 4.0 BY-SA版权

文章标签：

#python

pytorch 专栏收录该内容

18 篇文章

订阅专栏

本文深入探讨了torch_scatter.scatter_mean函数的使用，该函数通过src和index张量创建新的张量，按照index相同值对src元素进行维度上的平均计算。在给出的例子中，展示了一个二维张量在指定维度上如何利用scatter_mean进行均值运算，这对于处理分组数据或聚合计算非常有用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

scatter方法通过src和index两个张量来获得一个新的张量。

torch_scatter.scatter(src: torch.Tensor, index: torch.Tensor, dim: int = - 1, out: Optional[torch.Tensor] = None, dim_size: Optional[int] = None, reduce: str = 'sum') → torch.Tensor

原理如图，根据index，将index相同值对应的src元素进行对应定义的计算，dim为在第几维进行相应的运算。e.g.scatter_sum即进行sum运算，scatter_mean即进行mean运算。
在这里插入图片描述
e.g.x=scatter_mean(data.x, data.batch, dim=0)
我们给定一个二维张量x[952,21]为src，一维张量batch[952]为index，scatter_mean则将batch相同元素对应的x元素在0维上进行mean计算。具体过程如下：