PyTorch torch.argmax 函数
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))
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])