Accumulating data in a Callback

[lvoegtlin]

Hei everyone

I am currently working on a little framework for my experiments and I have a problem I can not solve. I would like to use Callbacks to create confusion matrixes and upload these as pictures as well as .tsv files onto wandb. The problem I have is that when I am training with DDP I get a warning which kills my main process: UserWarning: resource_tracker: There appear to be 6 leaked semaphore objects to clean up at shutdown. Here is my code:

class LogConfusionMatrixToWandbVal(Callback):
    def __init__(self):
        self.preds = []
        self.targets = []
        self.ready = True

    def on_sanity_check_start(self, trainer, pl_module) -> None:
        self.ready = False

    def on_sanity_check_end(self, trainer, pl_module):
        self.ready = True

    def on_validation_batch_end(
            self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx
    ):
        if self.ready:
            self.preds.append(outputs[OutputKeys.PREDICTION].detach().cpu().numpy())
            self.targets.append(outputs[OutputKeys.TARGET].detach().cpu().numpy())

    @rank_zero_only
    def on_validation_epoch_end(self, trainer, pl_module):
        if not self.ready:
            return

        conf_mat_name = f'CM_epoch_{trainer.current_epoch}'
        logger = get_wandb_logger(trainer)
        experiment = logger.experiment

        preds = []
        for step_pred in self.preds:
            preds.append(trainer.model.module.module.to_metrics_format(np.array(step_pred)))

        preds = np.concatenate(preds).flatten()
        targets = np.concatenate(np.array(self.targets)).flatten()

        num_classes = max(np.max(preds), np.max(targets)) + 1

        conf_mat = confusion_matrix(target=torch.tensor(targets),
                                                                    preds=torch.tensor(preds),
                                                                    num_classes=num_classes)

        # set figure size
        plt.figure(figsize=(14, 8))
        # set labels size
        sn.set(font_scale=1.4)
        # set font size
        fig = sn.heatmap(conf_mat, annot=True, annot_kws={"size": 8}, fmt="g")

        for i in range(conf_mat.shape[0]):
            fig.add_patch(Rectangle((i, i), 1, 1, fill=False, edgecolor='yellow', lw=3))
        plt.xlabel('Predictions')
        plt.ylabel('Targets')
        plt.title(conf_mat_name)

        conf_mat_path = Path(os.getcwd()) / 'conf_mats' / 'val'
        conf_mat_path.mkdir(parents=True, exist_ok=True)
        conf_mat_file_path = conf_mat_path / (conf_mat_name + '.txt')
        df = pd.DataFrame(conf_mat.detach().cpu().numpy())

        # save as csv or tsv to disc
        df.to_csv(path_or_buf=conf_mat_file_path, sep='\t')
        # save tsv to wandb
        experiment.save(glob_str=str(conf_mat_file_path), base_path=os.getcwd())
        # names should be uniqe or else charts from different experiments in wandb will overlap
        experiment.log({f"confusion_matrix_val_img/ep_{trainer.current_epoch}": wandb.Image(plt)},
                       commit=False)
        # according to wandb docs this should also work but it crashes
        # experiment.log(f{"confusion_matrix/{experiment.name}": plt})
        # reset plot
        plt.clf()
        self.preds.clear()
        self.targets.clear()

I think the problem is the two lists self.preds and self.targets but I don't know how I can make them synchronized to avoid concurrency problems.
Does somebody have any idea how to solve that?
Thanks in advance :)

[morganmcg1]

Hey @lvoegtlin , I'm not sure whats going on with the leaked semaphore issues, do you think its a wandb issue? If you're not, try make sure you're on the latest version of 0.12.6 in case its related to any bugs previously squashed in wandb.

Also, you can log confusion matrices with W&B's native confusion matrix plots too if you like: https://docs.wandb.ai/guides/track/log/plots#model-evaluation-charts

[lvoegtlin]

Hei,

Thanks for your input. We tried to solve the problem by decorating the functions with rank_zero_only but this should not be needed as the logger from PyTorch already just works on rank zero. Still, this did not solve our problem. I also think in the end wandb is the problem but we were not able to fix the issue in the Callback.

We solved it by moving everything into our LightningModule, which is less elegant but it works.
Also, we did not invest a lot of time into it as we had to get our experiments running :)

Thanks for your response :)