torch.argmax函数学习

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给出来会发生什么。

结合这两个才会知道dim在函数中的作用。下面举两个例子说明上面的第二点。

例子一：torch.argmax()函数中dim表示该维度会消失。
这个消失是什么意思？官方英文解释是：dim (int) – the dimension to reduce.
我们知道argmax就是得到最大值的序号索引，对于一个维度为(d0,d1) 的矩阵来说，我们想要求每一行中最大数的在该行中的列号，最后我们得到的就是一个维度为(d0,1) 的一矩阵。这时候，列就要消失了。

因此，我们想要求每一行最大的列标号，我们就要指定dim=1，表示我们不要列了，保留行的size就可以了。
假如我们想求每一列的最大行标，就可以指定dim=0，表示我们不要行了。

1. import torch

2. a=torch.tensor(

3. [

4. [1, 5, 5, 2],

5. [9, -6, 2, 8],

6. [-3, 7, -9, 1]

7. ])

8. b=torch.argmax(a,dim=0)

9. print(b)

10. print(a.shape)

结果

1. tensor([1, 2, 0, 1])

2. torch.Size([3, 4])

dim=0维度上为3,即在那3组数据中作比较，求得是每一列中的最大行标，因此为[1,2,0,4]

1. import torch

2. a=torch.tensor([

3. [

4. [1, 5, 5, 2],

5. [9, -6, 2, 8],

6. [-3, 7, -9, 1]

7. ],

9. [

10. [-1, 7, -5, 2],

11. [9, 6, 2, 8],

12. [3, 7, 9, 1]

13. ]])

14. b=torch.argmax(a,dim=0)

15. print(b)

16. print(a.shape)

18. """

19. tensor([[0, 1, 0, 1],

20. [1, 1, 1, 1],

21. [1, 1, 1, 1]])

22. torch.Size([2, 3, 4])"""

24. #dim=0,即将第一个维度消除，也就是将两个[3*4]矩阵只保留一个，因此要在两组中作比较，即将上下两个[3*4]的矩阵分别在对应的位置上比较

26. b=torch.argmax(a,dim=1)

27. “”“

28. tensor([[1, 2, 0, 1],

29. [1, 2, 2, 1]])

30. torch.Size([2, 3, 4])”“”

31. #dim=1，即将第二个维度消除,这么理解：矩阵维度变为[2*4];

32. """[1, 5, 5, 2],

33. [9, -6, 2, 8],

34. [-3, 7, -9, 1];

35. 纵向压缩成一维，因此变为[1,2,0,1];同理得到[1,2,2,1];"""

36. b=torch.argmax(a,dim=2)

37. """

38. tensor([[2, 0, 1],

39. [1, 0, 2]])

40. """

41. #dim=2,即将第三个维度消除，这么理解：矩阵维度变为[2*3]

42. """

43. [1, 5, 5, 2],

44. [9, -6, 2, 8],

45. [-3, 7, -9, 1];

46. 横向压缩成一维

47. [2,0,1],同理得到下面的“”“