PyTorch torch.take_along_dim 関数
Pytorch torch リファレンスマニュアルtorch.take_along_dimは、PyTorch で指定した次元に沿ってインデックス位置の要素を取得するための関数です。これは、indicesで指定されたインデックスに従い、dim次元に沿って、input中から値を取得します。
関数定義
torch.take_along_dim(input, indices, dim)
パラメータ:
input(Tensor): 入力テンソル。indices(Tensor): インデックステンソル。取得する要素の位置を指定します。形状は input と dim 次元で互換性がなければなりません。dim(int): 沿う次元。
戻り値:
torch.Tensor: インデックスに従って取り出された要素で構成される新しいテンソルを返します。
使用例
例
import torch
# 2Dテンソルを作成
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
print("元のテンソル:")
print(x)
# dim=1 に沿って要素を取得
indices = torch.tensor([[0, 1, 2],
[2, 1, 0],
[0, 0, 0]])
y = torch.take_along_dim(x, indices, dim=1)
print("nインデックス:")
print(indices)
print("ndim=1 に沿って取り出した要素:")
print(y)
# 2Dテンソルを作成
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
print("元のテンソル:")
print(x)
# dim=1 に沿って要素を取得
indices = torch.tensor([[0, 1, 2],
[2, 1, 0],
[0, 0, 0]])
y = torch.take_along_dim(x, indices, dim=1)
print("nインデックス:")
print(indices)
print("ndim=1 に沿って取り出した要素:")
print(y)
出力結果は次のとおりです：
原始张量:
tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
索引:
tensor([[0, 1, 2],
[2, 1, 0],
[0, 0, 0]])
沿 dim=1 取出的元素:
tensor([[1, 2, 3],
[6, 5, 4],
[7, 7, 7]])
例
import torch
# dim=0 に沿って要素を取得
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# 各列で異なる行を取得
indices = torch.tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
y = torch.take_along_dim(x, indices, dim=0)
print("元のテンソル:")
print(x)
print("nインデックス:")
print(indices)
print("ndim=0 に沿って取り出した要素:")
print(y)
# dim=0 に沿って要素を取得
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# 各列で異なる行を取得
indices = torch.tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
y = torch.take_along_dim(x, indices, dim=0)
print("元のテンソル:")
print(x)
print("nインデックス:")
print(indices)
print("ndim=0 に沿って取り出した要素:")
print(y)
出力結果は次のとおりです：
原始张量:
tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
索引:
tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
沿 dim=0 取出的元素:
tensor([[1, 5, 9],
[7, 2, 6],
[4, 8, 3]])
</p>
<div class="example">
<h2 class="example">实例</h2>
<div class="example_code">
<span style="color: Green;font-weight:bold;">import</span> torch<br />
<br />
<span style="color: #a50"># 在3D张量上使用</span><br />
x <span style="color: Gray;">=</span> torch.<span style="color: #05a;">arange</span><span style="color: Olive;">(</span><span style="color: Maroon;">24</span><span style="color: Olive;">)</span>.<span style="color: #05a;">reshape</span><span style="color: Olive;">(</span><span style="color: Maroon;">2</span><span style="color: Gray;">,</span> <span style="color: Maroon;">3</span><span style="color: Gray;">,</span> <span style="color: Maroon;">4</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"原始形状:"</span><span style="color: Gray;">,</span> x.<span style="color: #05a;">shape</span><span style="color: Olive;">)</span><br />
<br />
<span style="color: #a50"># 沿 dim=1 取元素</span><br />
indices <span style="color: Gray;">=</span> torch.<span style="color: #05a;">tensor</span><span style="color: Olive;">(</span><span style="color: Olive;">[</span><span style="color: Olive;">[</span><span style="color: Maroon;">0</span><span style="color: Gray;">,</span> <span style="color: Maroon;">1</span><span style="color: Gray;">,</span> <span style="color: Maroon;">2</span><span style="color: Olive;">]</span><span style="color: Gray;">,</span><br />
<span style="color: Olive;">[</span><span style="color: Maroon;">2</span><span style="color: Gray;">,</span> <span style="color: Maroon;">0</span><span style="color: Gray;">,</span> <span style="color: Maroon;">1</span><span style="color: Olive;">]</span><span style="color: Olive;">]</span><span style="color: Olive;">)</span><br />
y <span style="color: Gray;">=</span> torch.<span style="color: #05a;">take_along_dim</span><span style="color: Olive;">(</span>x<span style="color: Gray;">,</span> indices<span style="color: Gray;">,</span> dim<span style="color: Gray;">=</span><span style="color: Maroon;">1</span><span style="color: Olive;">)</span><br />
<br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"索引形状:"</span><span style="color: Gray;">,</span> indices.<span style="color: #05a;">shape</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"结果形状:"</span><span style="color: Gray;">,</span> y.<span style="color: #05a;">shape</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"n结果:"</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span>y<span style="color: Olive;">)</span><br />
</div>
</div>
<p>输出结果为：</p>
<pre>
原始形状: torch.Size([2, 3, 4])
索引形状: torch.Size([2, 3])
结果形状: torch.Size([2, 3, 4])
结果:
tensor([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[16, 17, 18, 19],
[12, 13, 14, 15],
[20, 21, 22, 23]]])
例
import torch
# 応用：バッチ次元に沿って特定の要素を選択
# 例：Attentionメカニズムで特定のkey-valueを選択
batch_size = 2
num_heads = 3
seq_len = 4
head_dim = 5
# queryのAttention重みをシミュレート
attn_weights = torch.randn(batch_size, num_heads, seq_len)
# 各ヘッドのtop-kのインデックスを取得
k = 2
indices = torch.argsort(attn_weights, dim=-1, descending=True)[..., :k]
print("インデックス形状:", indices.shape)
# valueテンソルをシミュレート
value = torch.randn(batch_size, num_heads, seq_len, head_dim)
# seq_len次元に沿って対応するvalueを取得
selected_value = torch.take_along_dim(value, indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), dim=2)
print("Value形状:", value.shape)
print("選択されたValue形状:", selected_value.shape)
# 応用：バッチ次元に沿って特定の要素を選択
# 例：Attentionメカニズムで特定のkey-valueを選択
batch_size = 2
num_heads = 3
seq_len = 4
head_dim = 5
# queryのAttention重みをシミュレート
attn_weights = torch.randn(batch_size, num_heads, seq_len)
# 各ヘッドのtop-kのインデックスを取得
k = 2
indices = torch.argsort(attn_weights, dim=-1, descending=True)[..., :k]
print("インデックス形状:", indices.shape)
# valueテンソルをシミュレート
value = torch.randn(batch_size, num_heads, seq_len, head_dim)
# seq_len次元に沿って対応するvalueを取得
selected_value = torch.take_along_dim(value, indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), dim=2)
print("Value形状:", value.shape)
print("選択されたValue形状:", selected_value.shape)
出力結果は次のとおりです：
索引形状: torch.Size([2, 3, 4]) Value形状: torch.Size([2, 3, 4, 5]) 选择的Value形状: torch.Size([2, 3, 2, 5])
注意：
torch.take_along_dim次元に沿ったインデックス指定が可能で、これはtorch.takeそれよりも柔軟です。後者は常にテンソルを1次元とみなすためです。