torch.nn.Softmax中dim参数的用途是什么

Question

我不知道torch.nn.Softmax中的dim参数适用于什么。有一个警告告诉我要使用它，并将其设置为1，但我不知道自己要设置什么。在公式中使用它的位置：

Softmax(xi​)=exp(xi)/∑j​exp(xj​)​

这里没有暗淡，那么它适用于什么？

Answer 1

torch.nn.Softmax上的Pytorch documentation状态：dim（int）–用来计算Softmax的尺寸（因此，沿着dim的每个切片的总和为1）。

例如，如果您有一个二维矩阵，则可以选择是否要将softmax应用于行或列：

import torch 
import numpy as np

softmax0 = torch.nn.Softmax(dim=0) # Applies along columns
softmax1 = torch.nn.Softmax(dim=1) # Applies along rows 

v = np.array([[1,2,3],
              [4,5,6]])
v =  torch.from_numpy(v).float()

softmax0(v)
# Returns
[[0.0474, 0.0474, 0.0474],
 [0.9526, 0.9526, 0.9526]])


softmax1(v)
# Returns
[[0.0900, 0.2447, 0.6652],
 [0.0900, 0.2447, 0.6652]]

请注意，对于softmax0，列如何添加到1，对于softmax1，行如何添加到1。