CrossEntropyLoss equivalence to LogSoftmax + NLLLoss - deep-learning

Question

According to the docs, CrossEntropyLoss criterion combines LogSoftmax function and NLLLoss criterion.
That is all fine and well, but testing it doesn't seem to substantiate this claim (ie assertion fails):
model_nll = nn.Sequential(nn.Linear(3072, 1024),
nn.Tanh(),
nn.Linear(1024, 512),
nn.Tanh(),
nn.Linear(512, 128),
nn.Tanh(),
nn.Linear(128, 2),
nn.LogSoftmax(dim=1))
model_ce = nn.Sequential(nn.Linear(3072, 1024),
nn.Tanh(),
nn.Linear(1024, 512),
nn.Tanh(),
nn.Linear(512, 128),
nn.Tanh(),
nn.Linear(128, 2),
nn.LogSoftmax(dim=1))
loss_fn_ce = nn.CrossEntropyLoss()
loss_fn_nll = nn.NLLLoss()
t = torch.rand(1,3072)
target = torch.tensor([1])
with torch.no_grad():
loss_nll = loss_fn_nll(model_nll(t), target)
loss_ce = loss_fn_ce(model_ce(t), target)
assert torch.eq(loss_nll, loss_ce)
I'm obviously missing something basic here.

Answer 1

As you noticed, the weights are initialized randomly.
One way to get two modules sharing the same weights is to simply export with state_dict the state of one and set it on the other with load_state_dict.
This is a one-liner:
>>> model_ce.load_state_dict(model_nll.state_dict())

Answer 2

The following assertion passes:
model = nn.Sequential(
nn.Linear(3072, 1024),
nn.Tanh(),
nn.Linear(1024, 512),
nn.Tanh(),
nn.Linear(512, 128),
nn.Tanh(),
nn.Linear(128, 2),
)
loss_fn_nll = nn.NLLLoss()
loss_fn_ce = nn.CrossEntropyLoss()
t = torch.rand(1, 3072)
target = torch.tensor([1])
with torch.no_grad():
loss_nll = loss_fn_nll(nn.LogSoftmax(dim=1)(model(t)), target)
loss_ce = loss_fn_ce(model(t), target)
assert torch.eq(loss_nll, loss_ce)
I assume the weights are randomly generated in the two networks in the original question. Even with torch.manual_seed(0) it still doesn't equate.

Answer 3

import torch
import torch.nn as nn
model_nll = nn.Sequential(nn.Linear(3072, 1024),
nn.Tanh(),
nn.Linear(1024, 512),
nn.Tanh(),
nn.Linear(512, 128),
nn.Tanh(),
nn.Linear(128, 2),
nn.LogSoftmax(dim=1))
model_ce = nn.Sequential(nn.Linear(3072, 1024),
nn.Tanh(),
nn.Linear(1024, 512),
nn.Tanh(),
nn.Linear(512, 128),
nn.Tanh(),
nn.Linear(128, 2)
)
model_nll.load_state_dict(dict(model_ce.named_parameters()))
loss_fn_ce = nn.CrossEntropyLoss()
loss_fn_nll = nn.NLLLoss()
t = torch.rand(1,3072)
target = torch.tensor([1])
with torch.no_grad():
loss_nll = loss_fn_nll(model_nll(t), target)
loss_ce = loss_fn_ce(model_ce(t), target)
print(loss_nll, loss_ce)
assert torch.eq(loss_nll, loss_ce)
Your code has two problems:
The weights must be identical in both models. Initialization is always random, therefore you must force them to be the same.
You should not add the LogSoftmax in the model_ce, internally in the CrossEntropyLoss it is already computed. This allows it to be numerically more stable: it simplifies the derivative and allows the log-sum-exp trick to be applied.