pytorch中dim的理解

首先先说个结论，就是dim等于那个维度，就是把那个维度给消除了，比如说shape=(2,3,4),如果dim=0,最后的结果的shape=(3,4),如果dim=1,最后的结果的shape=(2,4)，如果dim=2的话，最后的结果的shape=(2,3)

首先我们看个例子吧：

import numpy as np
import torch
x = torch.tensor([
        [1,2,3],
        [4,5,6]
    ])
# 我们可以看到"行"是dim=0, "列"是dim=1
print(x.shape)
'''输出
torch.Size([2, 3])
'''

# 但是我们按照dim=0求和, 是按照列相加，最后的shape=(3)
torch.sum(x, dim=0)
'''输出
tensor([5, 7, 9])
'''

# 但是我们按照dim=1求和, 是按照行相加1，最后的shape=(2)
torch.sum(x, dim=1)
'''输出
tensor([ 6, 15])
'''

下面也是类似：

# 看一下三维的
x = torch.tensor([
        [
         [1,2,3],
         [4,5,6]
        ],
        [
         [1,2,3],
         [4,5,6]
        ],
        [
         [1,2,3],
         [4,5,6]
        ]
    ])
# 我们可以看到第三维是dim=0, "行"是dim=1, 列是dim=2
print(x.shape)
'''输出
torch.Size([3, 2, 3])
'''

然后对于softmax:

import torch
import numpy as np
import torch.nn.functional as F
 
data = np.array([[0.1, 0.3, 0.6], [1.5,2.1 ,0.55]])
t_data = torch.from_numpy(data)
print(t_data)
print(t_data.shape)
#print(t_data.type())
print("**************************************************************************")
prob = F.softmax(t_data,dim=0) # dim = 0,在列上进行Softmax;dim=1,在行上进行Softmax
print(prob)
print(prob.shape)
#print(prob.type())
print("**************************************************************************")
prob = F.softmax(t_data,dim=1) # dim = 0,在列上进行Softmax;dim=1,在行上进行Softmax
print(prob)
print(prob.shape)
#print(prob.type())