PyTorch torch.take 関数

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

torch.takePyTorch で指定されたインデックス位置の要素を取得する関数です。入力テンソルを1次元配列として扱い、指定されたインデックス位置の要素を返します。

関数定義

torch.take(input, index)

引数:

input(Tensor): 入力テンソル。
index(Tensor): 整数インデックスのテンソル。取得する要素の位置を指定します。

戻り値:

torch.Tensor: 指定されたインデックス位置の要素から構成される新しいテンソルを返します。

使用例

例

import torch

# 2Dテンソルを作成
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])

print("元のテンソル:")
print(x)

# 2Dテンソルを1次元配列として扱い、インデックス0-8
# インデックス0, 4, 8の位置の要素を取得
index = torch.tensor([0, 4, 8])
y = torch.take(x, index)

print("nインデックス:", index)
print("取り出した要素:", y)

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

原始张量:
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])

索引: tensor([0, 4, 8])
取出的元素: tensor([1, 5, 9])

例

import torch

# 3Dテンソル
x = torch.arange(24).reshape(2, 3, 4)

print("元の3Dテンソル:")
print(x)
print("形状:", x.shape)

# 複数のインデックスの要素を取得
index = torch.tensor([0, 1, 2, 10, 20, 23])
y = torch.take(x, index)

print("nインデックス:", index)
print("取り出した要素:", y)

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

原始3D张量:
tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8, 9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])
形状: torch.Size([2, 3, 4])

索引: tensor([ 0,  1,  2, 10, 20, 23])
取出的元素: tensor([ 0,  1,  2, 10, 20, 23])

例

import torch

# 負のインデックスを使用
x = torch.tensor([[1, 2, 3],
[4, 5, 6]])

# 負のインデックスは末尾から数える
index = torch.tensor([0, -1]) # 最初と最後の要素
y = torch.take(x, index)

print("元の:", x)
print("インデックス [0, -1]:", y)

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

原始: tensor([[1, 2, 3],
        [4, 5, 6]])
索引 [0, -1]: tensor([1, 6])

例

import torch

# 要素をランダムに選択
x = torch.randn(10, 10)

# 10個のインデックスをランダムに生成
index = torch.randint(0, 100, (10,))
print("ランダムインデックス:", index)

# 対応する位置の要素を取得
selected = torch.take(x, index)

print("元のテンソルの形状:", x.shape)
print("選択された要素の形状:", selected.shape)
print("選択された要素:", selected)

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

随机索引: tensor([12, 45, 67, 82, 35, 59, 92,  7, 28, 73])
原始张量形状: torch.Size([10, 10])
选中的元素形状: torch.Size([10])
选中的元素: tensor([ 0.2345, -0.1234,  0.5678,  1.2345, -0.6789,  0.8901, -0.3456,  0.1234, -0.5678,  0.7890])

注意：torch.take入力テンソルをフラット化された1次元テンソルとしてインデックス付けします。インデックスは有効範囲内（0 から numel-1）にある必要があります。負のインデックス（-1 は最後の要素を表します）も使用できます。

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

その他の拡張