How do I log a confusion matrix into Wanddb?

Viewed 109

I'm using pytorch lightning, and at the end of each epoch, I create a confusion matrix from torchmetrics.ConfusionMatrix (see code below). I would like to log this into Wandb, but the Wandb confusion matrix logger only accepts y_targets and y_predictions. Does anyone know how to extract the updated confusion matrix y_targets and y_predictions from a confusion matrix, or alternatively give Wandb my updated confusion matrix in a way that it can be processed into eg a heatmap within wandb?

class ClassificationTask(pl.LightningModule):
    def __init__(self, model, lr=1e-4, augmentor=augmentor):
        super().__init__()
        self.model = model
        self.lr = lr
        self.save_hyperparameters() #not being used at the moment, good to have ther in the future
        self.augmentor=augmentor
        
        self.matrix = torchmetrics.ConfusionMatrix(num_classes=9)
        
        self.y_trues=[]
        self.y_preds=[]
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        x=self.augmentor(x)#.to('cuda')
        y_pred = self.model(x)
        loss = F.cross_entropy(y_pred, y,)  #weights=class_weights_tensor
        

        acc = accuracy(y_pred, y)
        metrics = {"train_acc": acc, "train_loss": loss}
        self.log_dict(metrics)
        return loss
    
    def validation_step(self, batch, batch_idx):
        loss, acc = self._shared_eval_step(batch, batch_idx)
        metrics = {"val_acc": acc, "val_loss": loss, }
        self.log_dict(metrics)
        return metrics
    
    def _shared_eval_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = F.cross_entropy(y_hat, y)
        acc = accuracy(y_hat, y)
        self.matrix.update(y_hat,y)
        return loss, acc
    
    def validation_epoch_end(self, outputs):
        confusion_matrix = self.matrix.compute()
        wandb.log({"my_conf_mat_id" : confusion_matrix})
        
    def configure_optimizers(self):
        return torch.optim.Adam((self.model.parameters()), lr=self.lr)
1 Answers

I'm actually working on the same issue currently. I found this great PR Feature request for PyTorch lightning. Perhaps this could be of help. I think a possible solution is utilizing torch metrics confusion matrix and then incorporating that into your train/val/test steps and logging them.

https://github.com/Lightning-AI/metrics/issues/880

Related