How to implement a custom Metric based on sklearn's functions

[sandylaker]

❓ Questions and Help

What is your question?

How to implement a custom Metric based on sklearn's functions?

Code

I tried to implement a ROCAUC Metric with help of sklearn's roc_auc_score, because it supports multilabel classification.

from pytorch_lightning.metrics import Metric
from sklearn.metrics import roc_auc_score
from typing import Optional, Any
import torch


class ROCAUC(Metric):
    def __init__(
        self,
        average='macro',
        compute_on_step: bool = True,
        dist_sync_on_step: bool = False,
        process_group: Optional[Any] = None,
    ):
        super().__init__(
            compute_on_step=compute_on_step,
            dist_sync_on_step=dist_sync_on_step,
            process_group=process_group,
        )

        self.average = average
        self.add_state("preds", default=[], dist_reduce_fx=None)
        self.add_state("target", default=[], dist_reduce_fx=None)

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        """
        Update state with predictions and targets.

        Args:
            preds: Predictions from model
            target: Ground truth values
        """
        self.preds.append(preds)
        self.target.append(target)

    def compute(self):
        preds = torch.cat(self.preds, dim=0).cpu().numpy()
        target = torch.cat(self.target, dim=0).cpu().numpy().astype(int)
        score = roc_auc_score(target, preds, average=self.average)
        return torch.tensor([score], dtype=torch.float32)

In the training procedure, I log the ROCAUC score along with other metrics at the end of validation epoch.

# this is the code in LightningModule
def validation_epoch_end(self, outputs):
    for name, metric in self.metrics.items():
        value = metric.compute()
        self.log(
            name,
            value,
            prog_bar=True,
            logger=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True)

It worked fine with single GPU training (without ddp).
However, when using ddp accelerator and setting move_metrics_to_cpu = True . It threw the following error:

Traceback (most recent call last):
  File "train.py", line 75, in <module>
    main()
  File "train.py", line 71, in main
    train_classifier(cfg, train_loader, val_loader, work_dir, seed, **args)
  File "/home/sandylaker/PycharmProjects/xray-classification/xray_cls/apis/train.py", line 40, in train_classifier
    trainer.fit(lit_model, train_loader, val_loader)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 470, in fit
    results = self.accelerator_backend.train()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 152, in train
    results = self.ddp_train(process_idx=self.task_idx, model=model)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 307, in ddp_train
    results = self.train_or_test()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 69, in train_or_test
    results = self.trainer.train()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 492, in train
    self.run_sanity_check(self.get_model())
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 690, in run_sanity_check
    _, eval_results = self.run_evaluation(test_mode=False, max_batches=self.num_sanity_val_batches)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 622, in run_evaluation
    deprecated_eval_results = self.evaluation_loop.evaluation_epoch_end()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 208, in evaluation_epoch_end
    deprecated_results = self.__run_eval_epoch_end(self.num_dataloaders, using_eval_result)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 246, in __run_eval_epoch_end
    eval_results = model.validation_epoch_end(eval_results)
  File "/home/sandylaker/PycharmProjects/xray-classification/xray_cls/apis/lit_model.py", line 97, in validation_epoch_end
    self.log(
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 266, in log
    self._results.log(
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 146, in log
    value = sync_fn(value, group=sync_dist_group, reduce_op=sync_dist_op)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 344, in sync_tensor
    return sync_ddp_if_available(tensor, group, reduce_op)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py", line 124, in sync_ddp_if_available
    return sync_ddp(result, group=group, reduce_op=reduce_op)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py", line 155, in sync_ddp
    torch.distributed.all_reduce(result, op=op, group=group, async_op=False)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 954, in all_reduce
    work = _default_pg.allreduce([tensor], opts)
RuntimeError: Tensors must be CUDA and dense

It seemed to be caused by sync_dist=True in the log function. When I modified the last line in ROCAUC's compute function:

return torch.tensor([score], dtype=torch.float32, device=self.preds[0].device)

the error is solved.

When we wrap the sklearn's metric function into pytorch-lightning's Metric, we have to first convert Tensor to numpy.ndarray, then we have to convert numpy.ndarray to Tensor after computing the metric value.

So my question is: Is there a better way than manually moving the Tensor to GPU like this?

What's your environment?

• OS: Ubuntu 20.04
• Packaging: pip
• Version: 1.1.2

[SkafteNicki]

Calling metric.compute() will sync the metric states from the different devises, so there is no need for calling self.log with sync_dist=True as value is already synced.
Also, to my understading roc_auc_score returns a native float which can directly be logged using self.log, so there is no need for the last conversion to torch.tensor.

[sandylaker]

Calling metric.compute() will sync the metric states from the different devises, so there is no need for calling self.log with sync_dist=True as value is already synced.
Also, to my understading roc_auc_score returns a native float which can directly be logged using self.log, so there is no need for the last conversion to torch.tensor.

Thanks for your suggestions. It really helped me a lot. By the way, is the syn_dist operation applied at the first step in compute function? I.e., first sync all the cached states, then compute the final metric?

[SkafteNicki]

Yes, first sync states from all processes and then compute the value.
Internally, we wrap the users compute, such that we can do the sync before calling the users compute method, it happens here in the code:
https://github.com/PyTorchLightning/pytorch-lightning/blob/019e4ff8cddadfde4d3bd48b4e8b8d294950ca2a/pytorch_lightning/metrics/metric.py#L203-L225