PyTorch torch.argmax 関数

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

torch.argmaxこれは、PyTorch で次元に沿った最大値のインデックスを返すために使用される関数です。

関数定義

torch.argmax(input, dim, keepdim)

使用例

実例

import torch

x = torch.tensor([[1, 3, 2], [4, 1, 3]])

# グローバル最大値のインデックス
print("グローバル最大インデックス:", torch.argmax(x))

# dim=0 に沿った最大値インデックス
print("dim=0 最大インデックス:", torch.argmax(x, dim=0))

# dim=1 に沿った最大値インデックス
print("dim=1 最大インデックス:", torch.argmax(x, dim=1))

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

全局最大索引: tensor(4)
dim=0 最大索引: tensor([1, 0, 1])
dim=1 最大索引: tensor([1, 0])

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

その他の拡張