pytorchメモ

Tweet

pytorchのCrossEntropyLossについて

CrossEntropyLossを使う場合、kerasは正解ラベルを1-hotベクトルで与えるのに対し、pytorchでは正解ラベルの数字配列を与えるだけでよい。

# coding: UTF-8
import torch
import torch.nn as nn

criterion = nn.CrossEntropyLoss()
inputs = torch.randn(3, 5)  # バッチサイズ3, 5次元のデータ
print(inputs)
targets = torch.tensor([0, 1, 3])  # 正解ラベル
loss = criterion(inputs, targets)  
print(loss)

出力結果は以下である。

tensor([[-1.8435,  1.1621,  0.1474,  1.9798, -0.9921],
        [ 0.3649,  0.5836, -0.4989, -0.0762, -0.7389],
        [ 0.2992,  0.3589,  0.1440, -1.0190, -0.7391]])
tensor(2.6648)

CrossEntropyLossを使用する場合、内部でsoftmaxが計算されるため、modelクラスの最後の出力層にsoftmaxのレイヤーは不要である。