Correct way to implement custom metrics - Tensors keep filling up

[cemde]

I wrote a custom BrierScore metric, but think I am missing the design idea of custom metrics. I store predictions and targets in preds and target. I expect these Tensors to empty after compute, but they keep filling up further and further. What did I do wrong?

class BrierScore(torchmetrics.Metric):
    def __init__(self, cfg, mode='logit'):
        super().__init__()
        
        self.cfg = cfg
        self.mode = mode
        if self.mode == 'logit':
            self.pred_transform = lambda x: torch.nn.functional.softmax(x, dim=1)
        else:
            self.pred_transform = lambda x: x
        
        self.add_state("preds", default=Tensor(), dist_reduce_fx="cat")
        self.add_state("target", default=Tensor(), dist_reduce_fx="cat")
    
    def update(self, preds: Tensor, target: Tensor) -> None:
        preds = self.pred_transform(preds.detach())
        self.preds = torch.cat([self.preds, preds])
        self.target = torch.cat([self.target, torch.nn.functional.one_hot(target, num_classes=self.cfg.DATA.NUM_CLASSES)])

    def compute(self):
        score = (self.preds - self.target)**2
        score = score.sum()/score.shape[0]
        return score

[justusschock]

Hi,

First, those are only emptied after a call to .reset() (which is automated within lightning), not after every call to compute.

Second: I would suggest to use lists for you state, since you are concatenating tensors which can be a memory allocating operation (not necessary at that time). So probably you could go with self.add_state("preds", default=[], dist_reduce_fx="cat") and self.preds.append(preds) instead.

[cemde]

Thank you so much for the reply! I assumed that .reset() would be called after every compute()call, because the data are no longer needed.

I am curious, what is the reason for not resetting them? After a few epochs these Tensors have almost a million rows, with e.g. 10 metrics for the train and validation split, that is a lot of memory...

I will just simply add .reset() manually within my compute() method.

And I will change to lists :) that makes a lot of sense

[justusschock]

In the original implementation, compute triggered reset automatically, but we changed this to not being done implicitly a while ago, because people needed to compute multiple times on the same data and this was not possible that way.