Pytorch dim

最近使用nn.softmax函数，会有dim= 这一个参数，这就是志明softmax作用的维度。
[batch_size, channel, h, w]
dim=0 batch_size, dim=1channel，以此类推。
这里有一个地方比较容易混，举个例子[h, w]，分别代表行和列，
dim=0表示作用在行维度上，即对同一列求softmax
dim=1表示作用在列维度上，即对同一行求softmax

x = Variable(torch.Tensor([[1, 2, 3],
                           [3, 4, 5]]))

y = Variable(torch.Tensor([1, 0]))
soft = nn.Softmax(dim=0)
print(soft(x))
# 得到结果tensor([[0.1192, 0.1192, 0.1192],
#               [0.8808, 0.8808, 0.8808]])
soft = nn.Softmax(dim=1)
print(soft(x))
# 得到结果tensor([[0.0900, 0.2447, 0.6652],
#              [0.0900, 0.2447, 0.6652]])