torch.max的用法
(max, max_indices) = torch.max(input, dim, keepdim=False)

  • 输入:
    1、input 是输入的tensor。
    2、dim 是索引的维度,dim=0寻找每一列的最大值,dim=1寻找每一行的最大值。
    3、keepdim 表示是否需要保持输出的维度与输入一样,keepdim=True表示输出和输入的维度一样,keepdim=False表示输出的维度被压缩了,也就是输出会比输入低一个维度。
  • 输出:
    1、max 表示取最大值后的结果。
    2、max_indices 表示最大值的索引。

举例说明

>import torch
>import numpy as np

>x = torch.randint(0,9,(2,4))
>print(x)
tensor([[7, 8, 7, 2],
   	    [6, 0, 3, 0]])

#取每一行的最大值,torch.max的输出结果
>y = torch.max(x, 1)  
>print(y)
torch.return_types.max(values=tensor([8, 6]),indices=tensor([1, 0])) #索引值

>y = torch.max(x, 1, keepdim=True)[0]
>print(y)
>print(np.shape(y))
# keepdim=True,输出仍然是二维的
tensor([[8],
     	[6]])
torch.Size([2, 1])

>y = torch.max(x, 1, keepdim=False)[0]
>print(y)
>print(np.shape(y))
# keepdim=False,输出变成了一维
tensor([8, 6])
torch.Size([2])
Logo

华为开发者空间,是为全球开发者打造的专属开发空间,汇聚了华为优质开发资源及工具,致力于让每一位开发者拥有一台云主机,基于华为根生态开发、创新。

更多推荐