torch.max

torch.operator , , 评论

torch.max

官方文档: pytorch

pytorch的torch.max函数有两种用法: 一种是不带任何参数的，他返回的是tensor中所有元素的最大值；第二种是带参数的，他返回最大值和对应的索引，这里分别介绍

第一种

torch.max(input)

Returns the maximum value of all elements in the input tensor.

返回所有元素的最大值

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
a = torch.rand((2, 3, 4))
b = torch.max(a)

'''
a
tensor([[[0.0615, 0.7280, 0.3775, 0.9397],
         [0.4438, 0.1904, 0.1330, 0.5213],
         [0.5054, 0.9993, 0.3916, 0.7125]],

        [[0.2660, 0.9682, 0.8774, 0.0991],
         [0.6073, 0.2710, 0.8681, 0.8080],
         [0.4416, 0.5851, 0.0409, 0.3030]]])

b
tensor(0.9993)
'''

第二种

torch.max(input, dim, keepdim=False, *, out=None)

返回对应维度元素的最大值以及对应的索引

  • input: 需要求最大值的tensor
  • dim: 求最大值的维度，参考normalize，第几维度，就是求那一维度的最大值，例如tensor维度为(3,4,5)，dim=1，那么就是求第一维度，每一组有四个元素的最大值
  • keepdim: 是否保持以前的维度数量，而不是大小

下面举两个例子

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
a = torch.rand((2, 3, 4))
b = torch.max(a, dim=1)

'''
a
tensor([[[0.9463, 0.7360, 0.0537, 0.5537],
         [0.0405, 0.5534, 0.4716, 0.0500],
         [0.2060, 0.6482, 0.8448, 0.0171]],

        [[0.7581, 0.5042, 0.1960, 0.0650],
         [0.2035, 0.8434, 0.2853, 0.1112],
         [0.9866, 0.3353, 0.2400, 0.5643]]])

b
torch.return_types.max(
values=tensor([[0.9463, 0.7360, 0.8448, 0.5537],
        [0.9866, 0.8434, 0.2853, 0.5643]]),
indices=tensor([[0, 0, 2, 0],
        [2, 1, 1, 2]]))
'''

这个例子可以看出，dim=1计算的是3那一维的最大值，可以看出0.9463=max(0.9463, 0.0405, 0.2060)
然后我们通过下面这个例子看一下keepdim的作用

1
2
3
4
5
6
7
8
9
10
11
12
13
c = torch.max(a, dim=1, keepdim=True)

'''
c
torch.return_types.max(
values=tensor([[[0.9463, 0.7360, 0.8448, 0.5537]],
        [[0.9866, 0.8434, 0.2853, 0.5643]]]),
indices=tensor([[[0, 0, 2, 0]],
        [[2, 1, 1, 2]]]))

c[0].size()
torch.Size([2, 1, 4])
'''

可以发现最后c的维度还是三维，不过求max的那一维度从3变成了1

harryteaustc

computer vision