deepmodeling / deepmodeling/unimol_tools
[Code scan] Keep multilabel classification targets float and mask missing labels
- Dominant language
- Python
- Stars
- 34
- Forks
- 9
- PR merge metrics
- No merged PRs in 30d
Description
This issue is a result of a Codex global repository scan.
## Summary
Trainer.decorate_* casts multilabel_classification targets to long together with single-label classification tasks. The registered multilabel losses include BCEWithLogitsLoss, GHMC_Loss, and FocalLossWithLogits, which require float targets. Casting also destroys NaN missing-label masks before the focal loss can filter them.
## Code references
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/tasks/trainer.py#L136-L164
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/models/nnmodel.py#L30-L33
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/models/loss.py#L106-L117
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/models/loss.py#L218-L234
## Impact
Multilabel classification can fail at runtime with BCE/GHM losses because float logits are compared against long targets. Missing labels can also be treated as real labels after integer casting, corrupting the loss.
## Suggested fix
Cast only classification and multiclass targets to long. Keep multilabel_classification targets as float tensors, and add a shared masked multilabel loss path that ignores NaN and any supported missing-label sentinel such as -1.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.