docs: fix subtle math bug in cifar10_cnn_pytorch example
Previously, the final layer of the network output a softmax
distribution, but was missing a log() before feeding that to the NLLLoss
function. This moves to removing the softmax layer and using the PyTorch
cross entropy loss, which includes the log-softmax operation.
@@ -68,7 +68,6 @@ class CIFARTrial(PyTorchTrial):
nn.ReLU(),
nn.Dropout2d(self.context.get_hparam("layer3_dropout")),
nn.Linear(512, NUM_CLASSES),
- nn.Softmax(dim=0),
)
# If loading backbone weights, do not call reset_parameters() or
@@ -91,7 +90,7 @@ class CIFARTrial(PyTorchTrial):
data, labels = batch
output = model(data)
- loss = torch.nn.functional.nll_loss(output, labels)
+ loss = torch.nn.functional.cross_entropy(output, labels)
error = error_rate(output, labels)
return {"loss": loss, "train_error": error}
