Named dataloaders for MetricLogger

[arseniybelkov]

We want dataloader_idx to be name and not simply an id of loader.

[AnihilatorGun]

import warnings
from collections import defaultdict

import numpy as np
import torch
from lightning import Callback
from more_itertools import zip_equal
from toolz import valmap
from lightning_utilities import apply_to_collection


def mb_asarray(x):
    try:
        return np.asarray(x)
    except ValueError:
        return x


def to_np(*x):
    return squeeze_first(apply_to_collection(x, torch.Tensor, tensor2np))


def squeeze_first(inputs):
    if len(inputs) == 1:
        inputs = inputs[0]
    return inputs


def tensor2np(x):
    return x.detach().cpu().numpy()


class ComputeMetrics(Callback):
    def __init__(self, val_dataloader_dict, test_dataloader_dict, single_metrics_dict, group_metrics_dict):
        self.val_dataloader_dict = val_dataloader_dict
        self.test_dataloader_dict = test_dataloader_dict
        self.single_metrics_dict = single_metrics_dict
        self.group_metrics_dict = group_metrics_dict

        self._train_losses = []
        # {dl_name: {metric_name: [metric_value...]}}
        self._single_metrics_accum = defaultdict(lambda: defaultdict(list))
        # {dl_name: {preprocessor: [list of preprocessed]}}
        self._group_metrics_accum = defaultdict(lambda: defaultdict(list))

    def _accumulate_metrics(self, pred, target, dl_name):
        # одиночные метрики
        for single_processor, metric_name2metric_func in self.single_metrics_dict.items():
            preprocessed = single_processor(target, pred)
            # если препроцессинг вернул одно значение — превращаем в кортеж
            if not isinstance(preprocessed, tuple):
                preprocessed = (preprocessed,)

            for metric_name, metric_func in metric_name2metric_func.items():
                single_metric_value = metric_func(*preprocessed)
                self._single_metrics_accum[dl_name][metric_name].append(single_metric_value)

        # групповые метрики
        for group_processor in self.group_metrics_dict:
            processed = group_processor(target, pred)
            if not isinstance(processed, tuple):
                processed = (processed,)
            self._group_metrics_accum[dl_name][group_processor].append(processed)

    def _finalize_metrics(self, pl_module, phase):
        for dl_name, metric_dict in self._single_metrics_accum.items():
            for metric_name, metric_values in metric_dict.items():
                pl_module.log(f"{phase}/{dl_name}/{metric_name}", np.nanmean(metric_values))

        for dl_name, group_processor2processed in self._group_metrics_accum.items():
            for group_processor, processed in group_processor2processed.items():
                #  [(a1, b1), (a2, b2), ...] -> ([a1,a2,...], [b1,b2,...]) - grouped and stacked args for group_metric
                processed = list(zip(*processed))
                processed = [mb_asarray(x) for x in processed]

                for metric_name, metric_func in self.group_metrics_dict[group_processor].items():
                    metric_value = metric_func(*processed)
                    pl_module.log(f"{phase}/{dl_name}/{metric_name}", metric_value)

        self._single_metrics_accum.clear()
        self._group_metrics_accum.clear()

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        dl_name = list(self.val_dataloader_dict.keys())[dataloader_idx]
        pred, target = outputs
        self._accumulate_metrics(pred, target, dl_name)

    def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
        dl_name = list(self.test_dataloader_dict.keys())[dataloader_idx]
        pred, target = outputs
        self._accumulate_metrics(pred, target, dl_name)

    def on_validation_epoch_end(self, trainer, pl_module):
        self._finalize_metrics(pl_module, "val")

    def on_test_epoch_end(self, trainer, pl_module):
        self._finalize_metrics(pl_module, "test")

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx: int) -> None:
        if outputs is None:
            return

        if isinstance(outputs, torch.Tensor):
            outputs = {'loss': to_np(outputs)}
        if isinstance(outputs, dict):
            outputs = valmap(to_np, outputs)
        elif isinstance(outputs, (list, tuple)):
            outputs = dict(zip_equal(map(str, range(len(outputs))), map(to_np, outputs)))
        else:
            raise TypeError(f"Unknown type of outputs: {type(outputs[0])}")

        self._train_losses.append(outputs)

    def on_train_epoch_end(self, trainer, pl_module) -> None:
        group = defaultdict(list)
        for loss in self._train_losses:
            for k, v in loss.items():
                group[k].append(v)

        n_entries = set(map(len, group.values()))
        if len(n_entries) != 1:
            warnings.warn(
                "Losses are inconsistent, number of entries for each loss: "
                f"{valmap(len, group)}. "
                "This can also happen due to rerun experiment, "
                "however please validate your loss function."
            )

        for k, vs in group.items():
            pl_module.log(f'train/{k}', np.mean(vs))

        self._train_losses.clear()