Skip to content

[torchlib] Support label_smoothing in aten::cross_entropy_loss - #3069

Open
Raashish Aggarwal (raashish1601) wants to merge 1 commit into
microsoft:mainfrom
raashish1601:fix/cross-entropy-label-smoothing
Open

Raashish Aggarwal (raashish1601) wants to merge 1 commit into
microsoft:mainfrom
raashish1601:fix/cross-entropy-label-smoothing

Conversation

@raashish1601

Copy link
Copy Markdown

aten_cross_entropy_loss accepted label_smoothing but ignored it, so models that use F.cross_entropy(..., label_smoothing=...) (or nn.CrossEntropyLoss(label_smoothing=...)) export without errors but compute a different loss:

x = torch.tensor([[2.0, 0.0, 0.0]]); t = torch.tensor([0])
F.cross_entropy(x, t, label_smoothing=0.1)  # 0.3729
# exported model returns 0.2395 (the loss without smoothing)

This PR follows PyTorch's cross_entropy_loss_label_smoothing. It takes the log_prob output of SoftmaxCrossEntropyLoss, sums -log_prob (weighted when weight is given) over the class dim, zeros it where target == ignore_index, and reduces it the same way as the nll term (for mean, divided by the summed weights or the count of non-ignored targets). The result is (1 - eps) * nll + eps / C * smooth. When label_smoothing == 0 the graph is unchanged.

Testing:

  • The nn.functional.cross_entropy OpInfo samples in PyTorch don't use label_smoothing, so I added an ops.aten.cross_entropy_loss.label_smoothing OpInfo with index targets, all three reductions, with and without weight, and with an ignore_index that hits some targets (2D and 4D inputs). Its 24 subtests fail on main and pass with the change.
  • pytest tests/function_libs/torch_lib/ops_test.py -k cross_entropy passes locally.
  • I also compared torch.onnx.export(..., dynamo=True) + onnxruntime against eager PyTorch for 72 combinations (2D/3D/4D inputs, reductions, weight, ignore_index, label_smoothing 0 and 0.2). All match within 1e-5.
  • ruff check and format are clean on the changed files.

label_smoothing was silently ignored, so exported models gave a different loss than PyTorch. Compute the smoothing term from the log_prob output of SoftmaxCrossEntropyLoss, following PyTorch's cross_entropy_loss_label_smoothing (weight, ignore_index and all reductions).

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

1 participant