pytorch torch.scatter_reduce函数介绍

PyTorch torch.scatter_reduce 函数

torch.scatter_reduce 是 PyTorch 中的一种高级操作,用于在特定维度上将源张量的值按索引归约到目标张量中。它结合了 scatterreduce 操作,非常适合处理需要对特定索引进行归约(如求和、最大值等)的场景。


函数签名

torch.scatter_reduce(input, dim, index, src, reduce, *, include_self=True, out=None)

参数说明

  1. input:

    • 目标张量,表示归约操作的初始值。
  2. dim:

    • 指定在目标张量 input 中进行归约操作的维度。
  3. index:

    • 张量,表示目标张量中归约操作的索引位置。
    • index 的形状必须与 src 兼容,或者可以广播成 src 的形状。
  4. src:

    • 源张量,提供要归约到 input 中的值。
  5. reduce:

    • 指定归约操作的类型,支持以下选项:
      • "sum":按索引进行求和。
      • "prod":按索引进行乘积。
      • "mean":按索引计算平均值。
      • "amax":按索引取最大值。
      • "amin":按索引取最小值。
  6. include_self (可选, 默认 True):

    • 是否在归约时包括 input 中的原始值。
    • 如果为 False,只使用 src 中的值进行归约。
  7. out (可选):

    • 用于存储结果的张量。如果提供,将直接修改此张量。

返回值

返回一个张量,包含归约操作的结果,形状与 input 相同。


示例

1. 按索引求和 (reduce="sum")
import torch

input = torch.zeros(3, 5)
index = torch.tensor([[0, 1, 2, 0, 1],
                      [1, 2, 0, 1, 2]])
src = torch.tensor([[10., 20., 30., 40., 50.],
                    [1., 2., 3., 4., 5.]])

result = torch.scatter_reduce(input, dim=1, index=index, src=src, reduce="sum")
print(result)

输出

tensor([[50., 70., 30.,  0.,  0.],
        [ 3.,  5.,  7.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.]])
2. 按索引取最大值 (reduce="amax")
result = torch.scatter_reduce(input, dim=1, index=index, src=src, reduce="amax")
print(result)

输出

tensor([[40., 50., 30.,  0.,  0.],
        [ 3.,  4.,  5.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.]])
3. 使用 include_self=False
result = torch.scatter_reduce(input, dim=1, index=index, src=src, reduce="sum", include_self=False)
print(result)

输出:

tensor([[50., 70., 30.,  0.,  0.],
        [ 3.,  5.,  7.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.]])

注意事项

  1. index 范围

    • index 的值必须在 [0, input.shape[dim]) 范围内,否则会引发错误。
  2. 广播规则

    • indexsrc 必须具有相同的形状,或者可以通过广播匹配。
  3. 性能优化

    • torch.scatter_reduce 对于稀疏更新和归约非常高效,避免了循环操作。

应用场景

  • 聚合数据(如按索引分组求和或求最大值)。
  • 构造稀疏张量。
  • 实现自定义的归约操作(如图神经网络中的消息传递)。
相关推荐
qq_15321452648 分钟前
Openai Dashboard可视化微调大语言模型
人工智能·语言模型·自然语言处理·chatgpt·nlp·gpt-3·transformer
明月清风徐徐10 分钟前
Miniconda + VSCode 的Python环境搭建
ide·vscode·python
笨鸟笃行23 分钟前
爬虫第七篇数据爬取及解析
开发语言·爬虫·python
青松@FasterAI25 分钟前
【Arxiv 大模型最新进展】PEAR: 零额外推理开销,提升RAG性能!(★AI最前线★)
人工智能
java1234_小锋29 分钟前
一周学会Flask3 Python Web开发-response响应格式
开发语言·python·flask·flask3
大数据追光猿31 分钟前
Python中的Flask深入认知&搭建前端页面?
前端·css·python·前端框架·flask·html5
java1234_小锋31 分钟前
一周学会Flask3 Python Web开发-flask3模块化blueprint配置
开发语言·python·flask·flask3
huoyingcg32 分钟前
武汉火影数字|VR沉浸式空间制作 VR大空间打造
人工智能·科技·vr·虚拟现实·增强现实
莫忘初心丶33 分钟前
python flask 使用教程 快速搭建一个 Web 应用
前端·python·flask
冷冷清清中的风风火火1 小时前
本地部署DeepSeek的硬件配置建议
人工智能·ai