pytorch 基本函数中的 dim【详细说明】：以torch.argmax为例

pytorch中有很多自带函数，掌握好这些函数，程序写起来当然非常有效率。

在这些pytorch函数中，很多都有dim这个控制参数，但是我们很难明白这个含义是什么。
本文试着总结一下：
1）dim的不同值表示不同维度。特别的在dim=0表示二维中的行，dim=1在二维矩阵中表示列。广泛的来说，我们不管一个矩阵是几维的，比如一个矩阵维度如下：$(d_0,d_1,...,d_{n-1})$，那么dim=0就表示对应到$d_0$也就是第一个维度，dim=1表示对应到$d_1$也就是第二个维度，依次类推。
2）知道dim的值是什么意思还不行，还要知道函数中这个dim给出来会发生什么。

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

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

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

import torch
import os
import numpy as np
# os.environ['CUDA_VISIBLE_DEVICES'] = '1'

a = torch.rand((3,4))

print(a.size())
print(a)

b = torch.argmax(a, dim=1)
print(b)
print(b.size())

输出：

torch.Size([3, 4])
tensor([[0.8338, 0.6953, 0.7558, 0.5803],
        [0.2105, 0.7638, 0.0912, 0.3341],
        [0.5585, 0.8019, 0.6590, 0.2268]])
tensor([0, 1, 1])
torch.Size([3])

可以看见指定dim=1时，列的size没有了。