How to calculate AUC over an entire test set?

[EvanZ]

So I have code like this:

def test_step(self, batch, batch_idx):
    x = batch['src']
    y = batch['label']
    mask = batch['mask']

    x = self.base_model(x, mask)
    x = self.linear(x).mean(axis=1).squeeze(1)
    
    loss = F.binary_cross_entropy_with_logits(input=x,
                           target=y)
    try:
        auc = roc_auc_score(y_true=y.cpu().numpy(), y_score=torch.sigmoid(x).cpu().numpy())
    except ValueError:
        auc = 0
    metrics = {'test_loss': loss, 'test_auc': auc}
    self.log_dict(metrics)
    return metrics

I'm trying to calculate the AUC for the model over a test set, but I assume this function is only calculating the AUC for a single batch. So I guess my question is whether the test set is treated as a single batch and if not, how would I go about doing this? Would I have to average the AUC over all the batches? If so, I'm not sure how to do that. But more importantly, that doesn't seem like the ideal way to calculate this metric to me, given how much variance there would be over each batch.

[rohitgr7]

averaging auc won't give the correct metric over the entire dataset plus you might get some error in case you have all targets of the same label which is possible in case of a single batch. To avoid this you can checkout https://torchmetrics.readthedocs.io/en/stable/references/modules.html#auroc. It will calculate the metrics the right way for you and you don't have to take care of averaging/accumulation or anything. There are different subpackage for each metric: module and functional. For your use-case I'd suggest using module one.

[EvanZ]

Thanks @rohitgr7. I tried to implement this package as follows:

def validation_step(self, batch, batch_idx):
    x = batch['src']
    y = batch['label']
    mask = batch['mask']

    x = self.base_model(x, mask)
    x = self.linear(x).mean(axis=1).squeeze(1)
    
    loss = F.binary_cross_entropy_with_logits(input=x,
                           target=y)
    return {'loss': loss, 'preds': x, 'target': y}

def validation_step_end(self, outputs):
    self.valid_auc(torch.sigmoid(outputs['preds']), outputs['target'].int())
    self.log('valid_auc', self.valid_auc)
    self.log('valid_loss', outputs['loss'])

Do you think this is correct or do I also need to write explicit code for validation_epoch_end?

[rohitgr7]

yes, by default for validation/test it accumulates over the whole corresponding dataset automatically. In training, it does not by default.
you can control it too using self.log(on_step=True|False, on_epoch=True|False).
you can check the table here: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#log

[EvanZ]

Hmm so I tried moving that code all to validation_step:

` def validation_step(self, batch, batch_idx):
x = batch['src']
y = batch['label']
mask = batch['mask']

    x = self.base_model(x, mask)
    x = self.linear(x).mean(axis=1).squeeze(1)
    
    loss = F.binary_cross_entropy_with_logits(input=x,
                           target=y)
    self.log('valid_loss', loss, on_step=True, on_epoch=True)
    self.valid_auc(torch.sigmoid(x),y.int())
    self.log('valid_auc', self.valid_auc, on_epoch=True, on_step=True)
    return {'loss': loss, 'preds': x, 'target': y}

`

But now I'm getting the following error:

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking arugment for argument tensors in method wrapper__cat) File "/home/ubuntu/deep-behavior-embedding/src/model/lightning_model.py", line 147, in validation_step (Current frame) self.valid_auc(torch.sigmoid(x),y.int())

Any idea why this would be the case? It was working before.

[rohitgr7]

are you using DP or DDP2 mode?

[EvanZ]

"DP"...

[rohitgr7]

great.. then you have to use *_step_end hooks for metric computation.
Ref doc: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#training-step-end