TorchMetrics: NaN-compatible variants for common Metrics functions [nan_(mean, sum, max)]

[danieleschmidt]

🚀 Feature

Now that callables are no longer supported (see Lightning-AI/pytorch-lightning#8894), here's a feature request for TorchMetric support for these metrics.

pytorch_lightning.utilities.exceptions.MisconfigurationException: Only self.log(..., reduce_fx={min,max,mean,sum})are currently supported. Please, open an issue inhttps://github.com/PyTorchLightning/pytorch-lightning/issues`. Found: <function nanmean at 0x7fa4b394a9e0>
`

Motivation

I'm currently working on a multitask learning problem, where some tasks include missing labels for some of the tasks. As a result, some batches contain nans for specific tasks. I thus require a custom NaNMean torchmetrics class to aggregate batches properly - which I think many in the community will find useful.

Alternatives

Another solution would be to support callables for reduce_fx functions, as were supported previously (preferable), or provide support for skip_nan=bool when calling self.log(...)

[github-actions]

Hi! thanks for your contribution!, great first issue!

[SkafteNicki]

Hi @danieleschmidt,
Thanks for raising this issue. I created PR #506 that introduces a couple of aggregation metrics (sum, mean, max, min, cat) that can all deal with with nan input in multiple ways. Please take a look and comment if anything is missing.