PyTorch で Softmax するとき計算する方向に関する注意点

機械学習フレームワーク PyTorch を使ってモデルを作成する際、softmax レイヤーを使う場合は注意が必要

softmax を計算する次元(軸)は

PyTorch で input データを作成するときは、以下のように配列の次元が増えていく

例えば、raw input のデータ１つが１次元データだった場合(時系列データなど)

1. raw inputが１次元のデータの場合 [[data1], [data2], [data3]]
   (0次元, 1次元) -> (データすべてでsoftmaxする方向, データの中身でsoftmaxする方向)
2. ミニバッチ学習させる場合 [[[data1], [data2], [data3]], [[data4], [data5], [data6]], ...]
3. (0次元, 1次元) -> (1つのバッチ内のすべてのデータでsoftmaxする方向, データの中身でsoftmaxする方向)
4. Convolution layer の input channel だけチャンネル数がある場合
   (0次元, 1次元, 2次元) -> (1つのバッチ内のすべてのデータでsoftmaxする方向, 各チャンネルごとにsoftmaxする方向(channel num: 1 バッチ内のすべてのデータでsoftmaxする方向と同じ), データの中身でsoftmaxする方向)

torch_tensor_test.py

test_tensor = torch.tensor([[[1.0,0,0, -1]],
                            [[1,0,0, -1]],
                            [[-1,-1,-1, -1]],
                            [[1,0,0, -1]],
                            [[2,0,0, -1]],
                            ])

#test_tensor shape: 5, 1, 4
m = nn.LogSoftmax(dim=-1)
output1 = m(test_tensor)
print(output1)
print(output1.shape)
'''
output1:
tensor([[[-0.6265, -1.6265, -1.6265, -2.6265]],

        [[-0.6265, -1.6265, -1.6265, -2.6265]],

        [[-1.3863, -1.3863, -1.3863, -1.3863]],

        [[-0.6265, -1.6265, -1.6265, -2.6265]],

        [[-0.2780, -2.2780, -2.2780, -3.2780]]])
torch.Size([5, 1, 4])
'''
m = nn.LogSoftmax(dim=2)
output2 = m(test_tensor)
print(output2)
print(output2.shape)
'''
output1:
tensor([[[-0.6265, -1.6265, -1.6265, -2.6265]],

        [[-0.6265, -1.6265, -1.6265, -2.6265]],

        [[-1.3863, -1.3863, -1.3863, -1.3863]],

        [[-0.6265, -1.6265, -1.6265, -2.6265]],

        [[-0.2780, -2.2780, -2.2780, -3.2780]]])
torch.Size([5, 1, 4])
'''

# if you want to apply softmax calculation on each input data, then dim = -1 or 2 would be good choice!