[PyTorch] NLLLoss と CrossEntropyLoss の違い

# この記事で説明すること
PyTorchのチュートリアルなどで，`torch.nn.NLLLoss` を交差エントロピーを計算するために使っている場面を見かけます．
私は初めて見た時，なぜ `torch.nn.CrossEntropyLoss` を使っていないのか疑問に感じました（こっちの方が関数名で何をするか想像しやすいし...）．
この記事では，この `NLLLoss` がどういう計算をしているのか，`CrossEntropyLoss` とどう違うのかについて説明します．

# `LogSoftmax` + `NLLLoss` と `CrossEntropyLoss` は同様の計算を行う
結論を述べると，`torch.nn.LogSoftmax` に続いて `torch.nn.NLLLoss` を適用するのと，`torch.nn.CrossEntropyLoss` を適用するのは同じことです．
以下のコードで実際に確認してみます．

```python:LogSoftmax+NLLLossとCrossEntropyLossの違い
import torch
import torch.nn as nn
torch.manual_seed(7777)

ce_loss = nn.CrossEntropyLoss()
logsoftmax = nn.LogSoftmax(dim=1)
nll_loss = nn.NLLLoss()

# モデルの最終的な全結合層の出力(という想定)
x = torch.randn(2, 5)
# 正解ラベルのリスト(という想定)
target = torch.tensor([2, 1])

# CrossEntropyLoss単体
print(ce_loss(x, target))
# LogSoftmax + NLLLoss
z =  logsoftmax(x)
print(nll_loss(z, target))

# 出力
# tensor(1.7886)
# tensor(1.7886)
```

# `LogSoftmax` + `NLLLoss` の動作説明
## `LogSoftmax`

`torch.nn.LogSoftmax` は，入力されたテンソル $x$ に対してsoftmaxの計算をしてからlogを取るだけです．
以下，[`torch.nn.LogSoftmax` の公式ドキュメント](https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html?highlight=logsoftmax)から式を引用します．

```math
\text{LogSoftmax}(x_i) = \log\left( \frac{\exp(x_i)}{\sum_j \exp(x_j)} \right)
```

ここで注意すべきは，**softmaxの式中の総和の部分をどの軸について取るか？**を指定する必要があることです．
これは引数 `dim` で指定します．同じテンソルでも，どの軸で総和を取るかによって，結果が異なってしまいます．その例を以下に示します．
以下を見ると，引数 `dim` が違うだけで計算結果が異なることがわかります．

```python:LogSoftmaxはdimに注意
x = torch.randn(2, 5)
z0 = nn.LogSoftmax(dim=0)(x)
z1 = nn.LogSoftmax(dim=1)(x)
print(z0)
print(z1)

# 出力
# tensor([[-0.0889, -1.0106, -0.4296, -0.7947, -0.8084],
#         [-2.4647, -0.4526, -1.0521, -0.6009, -0.5898]])
# tensor([[-0.2800, -2.7624, -2.4444, -2.7047, -3.5961],
#         [-1.5842, -1.1328, -1.9953, -1.4392, -2.3059]])

```

- 上の例では，`x` は (2,5) という形状のテンソルです．
- `z0` は0番目の軸について総和を取るため，ソフトマックスの式中の総和は2つの和となります．
    - 例えば，`z0[0][0]` を計算する際，*総和の部分* = `exp(x[0][0])+exp(x[1][0])` となります．
- `z1` は1番目の軸について総和を取るため，ソフトマックスの式中の総和は5つの和となります．
    - 例えば，`z1[0][0]` を計算する際，*総和の部分* = `exp(x[0][0])+exp(x[0][1])+...+exp(x[0][4])` となります．


### 実際に分類問題などで用いる場合
ニューラルネットワークを $C$ クラスの分類問題などで用いる際，ミニバッチサイズ $N$ として，この入力のテンソル $x$ が $(N,C)$ という形状の場合は，`dim=1` とする必要があります．
ミニバッチに含まれる各サンプルに対する計算結果（長さ $C$ のリスト）について，LogSoftmaxの総和の部分を計算する必要があるためです．


## `NLLLoss`

[`torch.nn.NLLLoss`の公式ドキュメント](https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html?highlight=nll#torch.nn.NLLLoss)に基づいて説明します．
まず，`NLLLoss` は Negative Log-Likelihood Loss を表すそうです．
しかし，実態を見ると，Log-Likelihood（対数尤度）の計算は特に担っておらず，**基本的に 'Negative' の部分しか担っていない**ことがわかりました．

### `reduction='mean'` 時の動作

**オプションについて:**

`NLLLoss` の引数 `reduction` によって計算する内容が異なります．
デフォルトは `reduction='mean'` なので，この時の動作について説明します．
また，引数 `weight` を設定することで，ロス計算時に各クラスに対する重みを設定できます（不均衡データなどの場合に有用らしい）が，デフォルトでは `weight=None` となっているのでこの点についても省略します．


**入力:**
最初のコードで示したように，`NLLLoss` には2つのテンソル $x,y$ を入力します．

- $x$ : 形状 $(N, C)$ の2次元のテンソル
- $y$ : 形状 $(N,)$ の1次元のテンソル

$x$ の方は，ミニバッチの各サンプルについて，各クラスに所属する確率のlogをとったものが格納された形状 $(N,C)$ のテンソルが入力されることが想定されています．
確率に変換し，logを取るという操作は `LogSoftmax` によって行えます．よって，この関数は`LogSoftmax` の後に適用することが想定されています．
一方で，$y$ の方は，ミニバッチの各サンプルについて，正解ラベルが $[0, C-1]$ の整数で格納された形状 $(N,)$ のテンソルが入力されることが想定されています．

**出力:**
そして，以下の値 $\ell (x,y)$ を計算して返します．

```math
\ell (x,y) = \frac{1}{N} \sum_{n=1}^N (-x_{n, y_n})
```

上式で使われている文字の説明は以下の通りです：

- $N$ : ミニバッチのサイズ
- $y_n$ : ミニバッチの $n$ 番目のサンプルの正解ラベル（$[0, C-1]$ の整数）
- $x_{n, y_n}$ : ミニバッチの $n$ 番目のサンプルに対する，クラス $y_n$（=正解のクラス）の予測値

つまり，ミニバッチの各サンプルについて，正解のラベルへの予測値にマイナスつけたものを平均しているだけです．
これが，基本的に 'Negative' の部分しか担っていない，という意味でした．

`reduction` が `'sum'`や`'none'` の場合の動作については，公式ドキュメントを見てください．
しかし，この `'mean'` の場合の動作が大体理解できれば他の場合も理解しやすいと思います．

**計算例**
以下に `NLLLoss` の計算例を示します．
ミニバッチサイズ $N=2$ ，クラス数 $C=5$ の場合です．
$\frac{1}{2} (-x_{0,4}-x_{1,1}) = \frac{1}{2}(-0.5-0.1) = -0.3$ と計算されています．

```python:NLLossの計算例
nll_loss = nn.NLLLoss()
xx = torch.tensor([ [0.1, 0.1, 0.1, 0.2, 0.5],
                    [0.1, 0.1, 0.1, 0.2, 0.5]
                  ])
yy = torch.tensor([4, 1])
nll_loss(xx, yy)

# 出力
# tensor(-0.3000)
```

# なぜ `LogSoftmax` と `NLLLoss` で `CrossEntropyLoss` と同じになるのか

[`torch.nn.CrossEntropyLoss` の公式ドキュメント](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)を見ると，以下の計算式によって計算した値を返す旨が書かれています．
この式は交差エントロピーの定義を少し変形すればわかります．

```math
\text{loss}(x, class) = -\log{\left( \frac{\exp({x[class]})}{\sum_j \exp(x[j])} \right)}
```

この式の右辺は，`LogSoftmax` で計算する式と似ています．違うのは，

- 先頭にマイナスがついている．
- $x$ の全ての要素ではなく，正解ラベル$class$に対応する要素だけを取り出している．

上記の2つの操作はちょうど `NLLLoss` を使うことで実現できます．
そのため，`LogSoftmax` と `NLLLoss` を使うと `CrossEntropyLoss` と同じになります．

### 実際の利用時
以上のことから，

- 訓練時に `NLLLoss` を利用する場合は，モデルの最後の全結合層の後に，`LogSoftmax` の層を入れておく必要があります．
- 訓練時に `CrossEntropyLoss` を利用する場合は，そのような追加の層を入れる必要はありません．

# 参考
[1] [`torch.nn.LogSoftmax` の公式ドキュメント](https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html?highlight=logsoftmax)
[2] [`torch.nn.NLLLoss`の公式ドキュメント](https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html?highlight=nll#torch.nn.NLLLoss)
[3] [`torch.nn.CrossEntropyLoss` の公式ドキュメント](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)
[4] [交差エントロピー - Wikipedia](https://ja.wikipedia.org/wiki/%E4%BA%A4%E5%B7%AE%E3%82%A8%E3%83%B3%E3%83%88%E3%83%AD%E3%83%94%E3%83%BC)
[5] [pytorch の NLLLoss の挙動 - メモ](https://tatsukawa.hatenablog.com/entry/2020/04/06/054700)