torch Softmax 차원에 따른 출력 값 차이

softmax 는 dim 을 지정해서 apply 할 수 있다. 

test = torch.randn(2,3,2)
print(F.softmax(test, dim = 0))
print(F.softmax(test, dim = 1))
print(F.softmax(test, dim = 2))

3D tensor에 0차원 1차원 2차원에 softmax를 apply 하고 결과를 보자. 

dim = 2 결과 부터 보면 2차원 원소들에 softmax 가 취해진 것을 볼 수 있다. 

0.7388 +  0.2612 = 1

0.4158 + 0.5842 = 1

 

dim = 1 결과를 보면,  (3) 

0.5528  + 0.3504 + 0.0968 = 1 

 

dim =2 결과를 보면, (2)

0.4200 + 0.5800 = 1

0.0447 + 0.9553 = 1

... 

 

착각하기 좋은 건 dim = 2 일 때, dim = 0 , dim = 1, dim = 2까지 모두 포함해서 softmax를 취한다고 이해하는 것이다. 

유의하자. 

tensor([[[0.4200, 0.0447],
         [0.7285, 0.2803],
         [0.3622, 0.8369]],

        [[0.5800, 0.9553],
         [0.2715, 0.7197],
         [0.6378, 0.1631]]])
tensor([[[0.5528, 0.1213],
         [0.3504, 0.3055],
         [0.0968, 0.5732]],

        [[0.7172, 0.7430],
         [0.1227, 0.2250],
         [0.1601, 0.0320]]])
tensor([[[0.7388, 0.2612],
         [0.4158, 0.5842],
         [0.0948, 0.9052]],

        [[0.1546, 0.8454],
         [0.0936, 0.9064],
         [0.4862, 0.5138]]])​