最近使用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]])