Custom Metric logging - Log the returned elements of tuple separately from inside MetricCollection

[rahulvigneswaran]

I have implemented a custom metric based on the existing torchmetrics.ConfusionMatrix which will take tuple of list of classes as input and returns a tuple of 3 during self.compute.

class ShotAcc(torchmetrics.ConfusionMatrix):
    def __init__(self, num_classes: int, classes_to_include, normalize: Optional[str] = None, threshold: float = 0.5, multilabel: bool = False, compute_on_step: bool = True, dist_sync_on_step: bool = False, process_group: Optional[Any] = None) -> None:
        super().__init__(num_classes, normalize, threshold, multilabel, compute_on_step, dist_sync_on_step, process_group)
        self.classes_to_include = classes_to_include
    
    def compute(self) -> Tensor:
        """Computes confusion matrix.
        Returns:
            If `multilabel=False` this will be a `[n_classes, n_classes]` tensor and if `multilabel=True`
            this will be a `[n_classes, 2, 2]` tensor
        """
        cm = _confusion_matrix_compute(self.confmat, self.normalize)
        cm_zero_corrected = cm.sum(dim=0)
        cm_zero_corrected[cm_zero_corrected==0] = 1

        many = (cm.diag()/cm_zero_corrected)[self.many_mid_low_tuple[0]].mean()
        mid = (cm.diag()/cm_zero_corrected)[self.many_mid_low_tuple[1]].mean()
        low = (cm.diag()/cm_zero_corrected)[self.many_mid_low_tuple[2]].mean()

        if many.isnan():
            many = torch.tensor(0.)
        if mid.isnan():
            mid = torch.tensor(0.)
        if low.isnan():
            low = torch.tensor(0.)

        return many, mid, low

Currently am logging the existing metrics like below inside the __init__ of the LightningModule,

# torchmetrics
metric_dict = {"Accuracy": Accuracy(num_classes=cfg.DATASET.NUM_CLS, average="micro"),
                "AccuracyMacro": Accuracy(num_classes=cfg.DATASET.NUM_CLS, average="macro"),
                "ConfusionMatrix": ConfusionMatrix(num_classes=cfg.DATASET.NUM_CLS, normalize="true")}
    
metrics = MetricCollection(metric_dict)
self.train_acc_mic = metrics.clone(prefix='train/') 
self.val_acc_mic = metrics.clone(prefix='val/')
self.test_acc_mic = metrics.clone(prefix='test/')

I am not sure how to add the custom metric defined above - ShotAcc to the metric_dict above and log the value of the tuple (many, mid, low) that it returns separately like, train/AccuracyMany, train/AccuracyMid, train/AccuracyLow, val/AccuracyMany, etc.