PyTorch torch.index_reduce 関数


Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

torch.index_reducePyTorch において、ソーステンソルの値を指定された方法で指定されたインデックス位置に集約するために使用される関数です。指定された次元に沿ってdim、指定されたindexインデックス位置で指定された方法で集約されたsource値。

関数定義

torch.index_reduce(input, dim, index, source, reduce='mean', *, include_self=True)

パラメータ:

  • input(Tensor): 入力テンソル。
  • dim(int): インデックスの次元。
  • index(Tensor): 1次元の整数テンソル。集約先の位置を指定します。
  • source(Tensor): ソーステンソル。集約する値です。
  • reduce(str): 集約方法。指定可能な値は 'mean'、'prod'、'amax'、'amin' です。デフォルトは 'mean' です。
  • include_self(bool, オプション): 集約にインデックス位置自体の元の値を含めるかどうか。デフォルトは True です。

戻り値:

  • torch.Tensor: 集約後のテンソルを返します。

使用例

例

import torch

# 入力テンソルを作成
input = torch.randn(3, 3)

# インデックスとソースを作成
index = torch.tensor([0, 0, 0])
source = torch.tensor([1.0, 2.0, 3.0])

# mean で集約
output = torch.index_reduce(input, dim=0, index=index, source=source, reduce='mean')

print("入力:")
print(input)
print("nインデックス:", index)
print("ソース:", source)
print("nmean 集約結果:")
print(output)

出力結果は次のとおりです:

输入:
tensor([[ 0.3456, -0.1234,  0.5678],
        [-0.5678,  0.1234, -0.6789],
        [ 0.7890, -0.3456,  0.1234]])

索引: tensor([0, 0, 0])
源: tensor([1., 2., 3.])

mean 聚合结果:
tensor([[ 2.5237,  0.2931,  0.3374],
        [-0.5678,  0.1234, -0.6789],
        [ 0.7890, -0.3456,  0.1234]])

例

import torch

# 異なる集約方法をテスト
input = torch.ones(3)
index = torch.tensor([0, 0, 0])
source = torch.tensor([2.0, 4.0, 6.0])

# prod 集約
output_prod = torch.index_reduce(input, 0, index, source, reduce='prod')
print("prod 集約:", output_prod)

# amax 集約
output_max = torch.index_reduce(input, 0, index, source, reduce='amax')
print("amax 集約:", output_max)

# amin 集約
output_min = torch.index_reduce(input, 0, index, source, reduce='amin')
print("amin 集約:", output_min)

出力結果は次のとおりです:

prod 聚合: tensor([48.,  1.,  1.])
amax 聚合: tensor([7., 1., 1.])
amin 聚合: tensor([3., 1., 1.])

例

import torch

# include_self パラメータ
input = torch.tensor([1.0, 10.0, 100.0])
index = torch.tensor([0, 0])
source = torch.tensor([5.0, 5.0])

# 自身の値を含める(デフォルト)
output1 = torch.index_reduce(input, 0, index, source, reduce='mean', include_self=True)
print("include_self=True:", output1)

# 自身の値を含めない
output2 = torch.index_reduce(input, 0, index, source, reduce='mean', include_self=False)
print("include_self=False:", output2)

出力結果は次のとおりです:

include_self=True: tensor([ 3.6667, 10.0000, 100.0000])
include_self=False: tensor([ 5., 10., 100.])

例

import torch

# 2Dテンソルへの適用
input = torch.zeros(3, 4)
index = torch.tensor([0, 2, 2])
source = torch.randn(3, 4)

# mean 集約
output = torch.index_reduce(input, dim=0, index=index, source=source, reduce='mean')

print("入力形状:", input.shape)
print("インデックス:", index)
print("ソース形状:", source.shape)
print("結果形状:", output.shape)
print("n結果:")
print(output)

注意:torch.index_reduce元の入力テンソルは変更されず、新しいテンソルが返されます。複数のインデックスが同じ位置を指すことができ、値は指定された集約方法に従って集約されます。include_self=False元の値を除外して新しく追加された値のみを集約できるため、特定のシナリオで役立ちます。


Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

その他の拡張