PyTorch torch.index_copy 関数


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

torch.index_copyは、PyTorch において、ソーステンソルを指定されたインデックス位置にコピーするための関数です。指定された次元に沿ってdim、index指定されたインデックス位置にコピーしますsourceの値。

とtorch.index_add異なる点は、index_copy上書きであり、累積ではないことです。

関数定義

torch.index_copy(input, dim, index, source)

パラメータ:

  • input(Tensor): 入力テンソル。
  • dim(int): インデックスの次元。
  • index(Tensor): 1次元の整数テンソルで、コピー先の位置を指定します。
  • source(Tensor): ソーステンソル、コピーする値。

戻り値:

  • torch.Tensor: 変更されたテンソルを返します。

使用例

例

import torch

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

# インデックスとソースを作成
index = torch.tensor([0, 2, 3])
source = torch.randn(3, 5)

# dim=0 に沿ってコピー
output = torch.index_copy(input, dim=0, index=index, source=source)

print("入力:")
print(input)
print("nソース:")
print(source)
print("n位置 [0, 2, 3] にコピーした後の結果:")
print(output)

出力結果は:

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

源:
tensor([[-1.2345,  0.5678, -1.2345,  0.5678, -1.2345],
        [ 1.5678, -0.6789,  1.5678, -0.6789,  1.5678],
        [-0.8901,  1.2345, -0.8901,  1.2345, -0.8901]])

复制到位置 [0, 2, 3] 后的结果:
tensor([[-1.2345,  0.5678, -1.2345,  0.5678, -1.2345],
        [-0.5678,  0.1234, -0.6789,  0.2345, -0.1234],
        [ 1.5678, -0.6789,  1.5678, -0.6789,  1.5678],
        [-0.8901,  1.2345, -0.8901,  1.2345, -0.8901]])

例

import torch

# dim=1 に沿ってコピー
input = torch.zeros(3, 5)
index = torch.tensor([1, 3])
source = torch.tensor([[10, 20, 30, 40, 50],
                        [60, 70, 80, 90, 100]])

output = torch.index_copy(input, dim=1, index=index, source=source)

print("入力:")
print(input)
print("nソース:")
print(source)
print("n位置 [1, 3] にコピーした後の結果:")
print(output)

出力結果は:

输入:
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

源:
tensor([[ 10.,  20.,  30.,  40.,  50.],
        [ 60.,  70.,  80.,  90., 100.]])

复制到位置 [1, 3] 后的结果:
tensor([[  0.,  10.,   0.,  20.,   0.],
        [  0.,  60.,   0.,  70.,   0.],
        [  0.,   0.,   0.,   0.,   0.]])

例

import torch

# 大きなテンソルを構築
# 複数の小さなバッチの結果を1つの大きなバッチにマージする必要があると仮定

# ターゲットテンソル
batch_size = 8
feature_dim = 4
output = torch.zeros(batch_size, feature_dim)

# 複数の小さなバッチの結果をシミュレート
mini_batches = [
    torch.randn(2, feature_dim),
    torch.randn(3, feature_dim),
    torch.randn(1, feature_dim)
]

# 各バッチを配置する位置
indices = [0, 2, 5]

# 各バッチを順番にコピー
for idx, batch in zip(indices, mini_batches):
    # 対応するサイズのインデックスを作成
    index = torch.arange(idx, idx + len(batch))
    output = torch.index_copy(output, dim=0, index=index, source=batch)

print("最終出力形状:", output.shape)
print(output)

出力結果は:

最终输出形状: torch.Size([8, 4])
tensor([[ 0.1234, -0.5678,  0.8901, -0.2345],
        [ 0.6789, -0.1234, -0.5678,  0.3456],
        [ 1.2345, -0.8901,  0.1234, -0.6789],
        [-0.3456,  0.5678, -0.1234,  0.8901],
        [ 1.5678, -0.2345,  0.6789, -0.1234],
        [-0.8901,  0.3456,  0.5678, -0.8901],
        [ 0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000]])

注意:torch.index_copy元の入力テンソルを変更せず、新しいテンソルを返します。index_copy上書き操作であり、torch.index_addの累積操作とは異なります。


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

その他の拡張