sigmoid_focal_loss gives NaN gradients for confident predictions when gamma < 1

Aperta Adatta ai principianti
  • Qualcosa non va?
#9,743 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
78/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
machine-learning

Direzione di ricerca

Inizia da torchvision/ops/focal_loss.py alle righe collegate nel report e riproduci i gradienti NaN con lo snippet fornito. Aggiungi un test di regressione per le predizioni corrette con elevata sicurezza quando gamma < 1, includendo la precisione ridotta se opportuno. Il lavoro è completato quando i gradienti rimangono finiti e il comportamento della loss è preservato.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

🐛 Describe the bug

With 0 < gamma < 1, torchvision.ops.sigmoid_focal_loss returns a finite loss but NaN gradients for any element whose prediction is confidently correct. One such element makes the gradients of the whole model NaN. The default gamma=2 is not affected, and neither is gamma=1.

import torch
from torchvision.ops import sigmoid_focal_loss

# Confidently correct predictions: a positive with logit 20 and a negative with logit -20
logits = torch.tensor([20.0, -20.0], requires_grad=True)
targets = torch.tensor([1.0, 0.0])
sigmoid_focal_loss(logits, targets, gamma=0.5, reduction="sum").backward()
print(logits.grad)  # tensor([nan, nan])

# In bfloat16 it already happens at logit 6.5
logits = torch.tensor([6.5, -6.5], dtype=torch.bfloat16, requires_grad=True)
sigmoid_focal_loss(logits, targets.bfloat16(), gamma=0.5, reduction="sum").backward()
print(logits.grad)  # tensor([nan, nan], dtype=torch.bfloat16)

# One such logit makes every gradient of the classifier NaN
torch.manual_seed(0)
head = torch.nn.Linear(16, 1)
features = torch.randn(8, 16)
with torch.no_grad():
    features[0] = -15 * head.weight[0].sign()  # one easy, very confident negative
logits = head(features).squeeze(1)
loss = sigmoid_focal_loss(logits, torch.zeros(8), gamma=0.5, reduction="sum")
loss.backward()
print(logits[0].item(), loss.item(), head.weight.grad.isnan().all().item())  # about -24, finite loss, True

The gradients become NaN once |logit| of a correctly classified element reaches 16.75 in float32, 8.5 in float16 and 6.25 in bfloat16 (the same for gamma 0.25, 0.5 and 0.75).

The cause is in focal_loss.py#L41-L43. p = torch.sigmoid(inputs) rounds to exactly 1 (or 1 - p to exactly 1), so 1 - p_t is exactly 0. The derivative of (1 - p_t) ** gamma at 0 is infinite for gamma < 1, and the chain rule multiplies it by p * (1 - p), which is also 0, giving NaN.

Two changes I tried, both of which give the same loss as the current code in float64 (relative difference 0 on 10,000 random elements, gamma 0.5 and 2):

  1. Compute 1 - p_t without the 1 - p cancellation: p * (1 - targets) + torch.sigmoid(-inputs) * targets. This moves the threshold out to |logit| = 88.75, where sigmoid(-x) underflows.
  2. Clamp it: (1 - p_t).clamp(min=torch.finfo(inputs.dtype).tiny) ** gamma. No NaN at any logit I tried.

I'm happy to open a PR with either one plus a test, if that sounds useful.

AI disclosure: I used Claude Code to write a script that runs torchvision.ops functions at edge-case inputs and checks values and gradients. It flagged this, then helped measure the thresholds and try the two fixes. I checked the cause in the code and ran the snippet above myself.

Versions
  • PyTorch version: 2.16.0.dev20261107+cpu
  • torchvision version: 0.30.0.dev20261107+cpu (the same code is on main at 426030113)
  • OS: Ubuntu 24.04.5 LTS (x86_64), WSL2
  • Python version: 3.12.3
Lingua principale
Python
Stelle
18k
Fork
7.3k
Merge medio
4h 56m
PR unite (30g)
9

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di pytorch/vision

Tutte le issue di pytorch/vision

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.