首页 > 技术文章 > argmax()、max()介绍

Allen-mat 2021-03-28 19:29 原文

WechatIMG515

在开始之前先明确“行轴”和“列轴”的方向如上图。图中“行轴”索引范围为0-2,轴长为3,“列轴”索引范围0-3,轴长为4。

对张量的操作

import torch
t = torch.tensor([[3,4,1,0],[12,1,2,3],[52,34,88,11]])
t
# 输出为:
tensor([[ 3,  4,  1,  0],
        [12,  1,  2,  3],
        [52, 34, 88, 11]])
--------------------------------------------------------------------------------------------------------------
t.max() # 相当于将张量进行flatten后取出最大值
# 输出为:
tensor(88)
--------------------------------------------------------------------------------------------------------------
t.max(dim=0) # 沿张量行轴进行操作,返回每列的最大值及其对应索引
# 输出为:
(values=tensor([52, 34, 88, 11]),
indices=tensor([2, 2, 2, 2]))
--------------------------------------------------------------------------------------------------------------
t.max(dim=1) # 沿张量列轴进行操作,返回每行的最大值及其对应索引
# 输出为:
(values=tensor([ 4, 12, 88]),
indices=tensor([1, 0, 2]))
--------------------------------------------------------------------------------------------------------------
t.max(dim=0)[0] 
# 输出每个行轴的最大值:
tensor([52, 34, 88, 11])
t.max(dim=0)[1]
# 输出每个行轴最大值对应的索引:
tensor([2, 2, 2])
--------------------------------------------------------------------------------------------------------------
t.argmax() # 类似max(),但仅取最大值索引
# 输出为:
tensor(10) # flatten后最大值88对应的索引为10

对numpy数组的操作

# 用max()时,与对tensor的操作唯一不同是不再返回索引值
import numpy as np
t = np.array([[3,4,1,0],[12,1,2,3],[52,34,88,11]])
t.max(axis=0) 
# 输出为:
array([52, 34, 88, 11])
--------------------------------------------------------------------------------------------------------------
# 用argmax()时,没有区别

推荐阅读