pytorch里的argmax和argmin的理解

最新推荐文章于 2025-05-29 10:43:49 发布

顶会accept

最新推荐文章于 2025-05-29 10:43:49 发布

阅读量647

点赞数

CC 4.0 BY-SA版权

文章标签：深度学习 pytorch

本文链接：https://blog.csdn.net/loki2018/article/details/83009473

本文详细介绍了PyTorch中argmin和argmax函数的使用方法，包括如何获取张量中所有元素的最大值/最小值的索引，以及在特定维度上进行操作的技巧。通过实例演示了不同参数设置下函数的行为，帮助读者深入理解这两个函数的功能。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

返回输入张量中所有元素的最大值/最小值的索引。

torch.argmin(dim=None, keepdim=False) → LongTensor
torch.argmax(input, dim, keepdim=False) → LongTensor

dim：用于指定在哪个维度进行操作
keepdim：若为True,则输出的索引的维数和输入的张量的维度保持一致，一般默认为False

实验如下

a = torch.randn(4, 4)
print(a)
print(a.flatten())
print(torch.argmax(a))  # 不指定dim，则返回张量展开为一维最大值的索引
"""
tensor([[ 1.0284, -0.1810,  2.2535,  1.0803],
        [ 0.0180,  0.8234, -1.8942, -0.0986],
        [-0.7464,  0.5619,  1.1605,  0.3683],
        [-0.7363, -0.2394, -0.6425, -0.5612]])
tensor([ 1.0284, -0.1810,  2.2535,  1.0803,  0.0180,  0.8234, -1.8942, -0.0986,
        -0.7464,  0.5619,  1.1605,  0.3683, -0.7363, -0.2394, -0.6425, -0.5612])
tensor(2)
"""
a = torch.randn(4, 4)
print(a)
print(torch.argmax(a), dim=1)
"""
tensor([[-0.3381, -0.4999, -1.1872, -0.6996],
        [-1.1637, -1.9946, -1.5603,  3.1711],
        [ 0.1105, -1.4377,  0.0382, -0.3079],
        [ 2.2044,  0.9992, -1.6193, -0.2622]])
tensor([0, 3, 0, 0])
"""