Torchvisionのmodelの最後にsoftmaxがいらない件について。

背景

以下のコードを見てくれればわかるが、最終層はLinerでsoftmax層が入っていない。
https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py

「え、これで大丈夫なの？」と思ったので、

理由

以下に全て書いてあった。
https://discuss.pytorch.org/t/torchvision-models-dont-have-softmax-layer/18071

学習時には、nn.CrossEntoropyLoss()を使うが、nn.LogSoftmaxとnn.NLLLoss.から成り立ってるので、必要ないとのこと。

推論する際も、各クラスの確率が欲しいならnn.functional.softmax()が必要だが、予測ならtorch.maxなどで最も大きい値のidxを抜いてこればいいとのこと。

クラス分類時は脳死で最終層はsoftmaxと考えていたけど、よくよく考えればloss計算に必要なだけなので、「なるほどな〜」と思った。