torch.argmax()用法示例

argmax函数torch.argmax(input, dim=None, keepdim=False) 返回指定维度最大值的序号,dim给定的定义是:the demention to reduce,就是把dim这个维度,变成这个维度的最大值的index。
参数含义:
1)dim表示不同维度。特别的在dim=0表示二维矩阵中的列,dim=1在二维矩阵中的行。广泛的来说,我们不管一个矩阵是几维的,比如一个矩阵维度如下:(d0,d1,…,dn−1) ,那么dim=0就表示对应到d0 也就是第一个维度,dim=1表示对应到也就是第二个维度,以此类推。

2)知道dim的值是什么意思还不行,还要知道函数中这个dim给出来会发生什么。

Example one:二维数组
看例子推理一下

x = torch.randn(2, 4, 5)
print(x)
'''
tensor([[[-1.2204, -0.6428, -0.2278,  0.5589,  1.1589],
         [ 0.4235,  1.9663,  0.5055, -1.3472,  1.3523],
         [ 1.4220,  0.7886, -1.0821,  0.6268, -0.9465],
         [-0.3950,  1.3275,  0.3369,  1.0224, -0.9944]],

        [[ 0.6024, -0.2604, -0.8631,  0.8113, -0.3140],
         [ 0.3487, -0.1941, -0.3955, -0.1719, -1.3734],
         [ 0.2467, -0.4268, -1.3428,  0.7346,  1.0932],
         [-0.5799,  0.0976, -1.9403, -0.2643,  0.7657]]])
'''

Example two:三维数组
1)dim=0时,返回每个元素在各个通道最大值的索引号
2)dim=1时,返回n排每一列最大值的索引(n是通道数)
3)dim=2时,返回n排每一排最大值的索引(n是通道数)

x = torch.randn(2, 4, 5)
print(x)
'''
tensor([[[-1.2204, -0.6428, -0.2278,  0.5589,  1.1589],
         [ 0.4235,  1.9663,  0.5055, -1.3472,  1.3523],
         [ 1.4220,  0.7886, -1.0821,  0.6268, -0.9465],
         [-0.3950,  1.3275,  0.3369,  1.0224, -0.9944]],

        [[ 0.6024, -0.2604, -0.8631,  0.8113, -0.3140],
         [ 0.3487, -0.1941, -0.3955, -0.1719, -1.3734],
         [ 0.2467, -0.4268, -1.3428,  0.7346,  1.0932],
         [-0.5799,  0.0976, -1.9403, -0.2643,  0.7657]]])
'''

# dim=0,将第一个维度消除,也就是将两个[4*5]矩阵只保留一个,因此要在上下两个[3*4]的矩阵分别在对应位置上比较
y0 = torch.argmax(x, dim=0)
print(y0)
'''
tensor([[1, 1, 0, 1, 0],
        [0, 0, 0, 1, 0],
        [0, 0, 0, 1, 1],
        [0, 0, 0, 0, 1]])
'''

# dim=1,将第二个维度消除,也就是将四个[2*5]矩阵只保留一个
y1 = torch.argmax(x, dim=1)
print(y1)
'''
tensor([[2, 1, 1, 3, 1],
        [0, 3, 1, 0, 2]])
'''
# dim= 2
y2 = torch.argmax(x, dim=2)
print(y2)
'''
tensor([[4, 1, 0, 1],
        [3, 0, 4, 4]])
'''
  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值