deepmodeling / deepmodeling/unimol_tools

[Code scan] Keep multilabel classification targets float and mask missing labels

Open
#27 0 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.