Multiple Evaluator Datasets

composer/core/evaluator.py

@@ -0,0 +1,38 @@
+# Copyright 2021 MosaicML. All Rights Reserved.
+from __future__ import annotations
+
+import copy
+from typing import TYPE_CHECKING, Union
+
+from torchmetrics import Metric, MetricCollection
+
+from composer.core.data_spec import DataSpec as DataSpec
+
+if TYPE_CHECKING:
+    from composer.core.types import DataLoader, Metrics
+
+
+class Evaluator:
+    """Wrapper for a dataloader to include metrics that apply to a specific
+    dataset.
+
+    Attributes:
+        label (str): Name of the Evaluator
+        dataloader (Union[DataSpec, DataLoader]): Dataloader/DataSpec for evaluation data
+        metrics (Metrics): Metrics to log. The metrics will be deep-copied to ensure that
+            each evaluator updates only its metrics.
+    """
+
+    def __init__(self, *, label: str, dataloader: Union[DataSpec, DataLoader], metrics: Metrics):
+        self.label = label
+        if isinstance(dataloader, DataSpec):
+            self.dataloader = dataloader
+        else:
+            self.dataloader = DataSpec(dataloader)
+
+        # Forcing metrics to be a MetricCollection simplifies logging results
+        metrics = copy.deepcopy(metrics)
+        if isinstance(metrics, Metric):
+            self.metrics = MetricCollection([metrics])
+        else:
+            self.metrics = metrics

composer/core/state.py

@@ -59,7 +59,7 @@
     "batch_num_tokens",
     "outputs",
     "train_dataloader",
-    "eval_dataloader",
+    "evaluators",
     "_steps_per_epoch",
     "_precision_context",
     "profiler",
@@ -79,8 +79,8 @@ class State(Serializable):
         grad_accum (int): The number of gradient accumulation steps to use. The size of each microbatch is ``train_batch_size / num_gpus / grad_accum``.
         train_dataloader (types.DataLoader, types.DataSpec, or dict):
             The :class:`types.DataLoader`, :class:`types.DataSpec`, or dict of :class:`types.DataSpec` kwargs to used for training.
-        eval_dataloader (types.DataLoader, types.DataSpec, or dict):
-            The :class:`types.DataLoader`, :class:`types.DataSpec`, or dict of :class:`types.DataSpec` kwargs to used for evaluation.
+        evaluators (Evaluators):
+            The :class:`types.Evaluators` contain the evaluation datasets used for evaluation with specific metrics.
         max_duration (str or Time): The maximum duration to train for.
 
         precision (str | Precision): The numerical precision to use for training. Should be one of ``[fp32, amp]``.
@@ -122,7 +122,7 @@ def __init__(
             # data configurations
             grad_accum: int,
             train_dataloader: types.DataLoader,
-            eval_dataloader: types.DataLoader,
+            evaluators: types.Evaluators,
 
             # stopping conditions
             max_duration: Union[str, Time[int]],
@@ -148,7 +148,7 @@ def __init__(
         self.model = model
         self.grad_accum = grad_accum
         self.train_dataloader = train_dataloader
-        self.eval_dataloader = eval_dataloader
+        self.evaluators = list(ensure_tuple(evaluators))
         self.max_duration = max_duration
         self.steps_per_epoch = steps_per_epoch
 

composer/core/types.py

@@ -16,6 +16,7 @@
 
 from composer.core.algorithm import Algorithm as Algorithm
 from composer.core.data_spec import DataSpec as DataSpec
+from composer.core.evaluator import Evaluator as Evaluator
 from composer.core.event import Event as Event
 from composer.core.logging import Logger as Logger
 from composer.core.precision import Precision as Precision
@@ -141,8 +142,8 @@ def __len__(self) -> int:
         ...
 
 
+Evaluators = Union[Evaluator, List[Evaluator], Tuple[Evaluator, ...]]
 Metrics = Union[Metric, MetricCollection]
-
 Optimizer = torch.optim.Optimizer
 Optimizers = Union[Optimizer, Tuple[Optimizer, ...], List[Optimizer]]
 Scheduler = torch.optim.lr_scheduler._LRScheduler

composer/datasets/__init__.py

@@ -5,6 +5,8 @@
 from composer.datasets.cifar10 import CIFAR10DatasetHparams as CIFAR10DatasetHparams
 from composer.datasets.dataloader import DataloaderHparams as DataloaderHparams
 from composer.datasets.dataloader import WrappedDataLoader as WrappedDataLoader
+from composer.datasets.dataset_registry import get_dataset_registry as get_dataset_registry
+from composer.datasets.evaluator import EvaluatorHparams as EvaluatorHparams
 from composer.datasets.glue import GLUEHparams as GLUEHparams
 from composer.datasets.hparams import DatasetHparams as DatasetHparams
 from composer.datasets.hparams import SyntheticHparamsMixin as SyntheticHparamsMixin

composer/datasets/dataset_registry.py

@@ -0,0 +1,23 @@
+# Copyright 2021 MosaicML. All Rights Reserved.
+
+from composer.datasets.ade20k import ADE20kDatasetHparams
+from composer.datasets.brats import BratsDatasetHparams
+from composer.datasets.cifar10 import CIFAR10DatasetHparams
+from composer.datasets.glue import GLUEHparams
+from composer.datasets.imagenet import ImagenetDatasetHparams
+from composer.datasets.lm_datasets import LMDatasetHparams
+from composer.datasets.mnist import MNISTDatasetHparams
+
+registry = {
+    "ade20k": ADE20kDatasetHparams,
+    "brats": BratsDatasetHparams,
+    "imagenet": ImagenetDatasetHparams,
+    "cifar10": CIFAR10DatasetHparams,
+    "mnist": MNISTDatasetHparams,
+    "lm": LMDatasetHparams,
+    "glue": GLUEHparams
+}
+
+
+def get_dataset_registry():
+    return registry

composer/datasets/evaluator.py

@@ -0,0 +1,91 @@
+# Copyright 2021 MosaicML. All Rights Reserved.
+
+from __future__ import annotations
+
+import copy
+import logging
+import textwrap
+from dataclasses import dataclass
+from typing import List, Optional
+
+import yahp as hp
+from torchmetrics import Metric, MetricCollection
+
+from composer.core.types import Evaluator
+from composer.datasets import DataloaderHparams
+from composer.datasets.dataset_registry import get_dataset_registry
+from composer.datasets.hparams import DatasetHparams
+from composer.models.base import BaseMosaicModel
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class EvaluatorHparams(hp.Hparams):
+    """Params for the :class:`Evaluator`.
+
+    See the documentation for the :class:`Evaluator`.
+    """
+    hparams_registry = {  # type: ignore
+        "eval_dataset": get_dataset_registry(),
+    }
+
+    label: str = hp.required(doc="Name of the Evaluator object. Used for logging/reporting metrics")
+    eval_dataset: DatasetHparams = hp.required(doc="Evaluator dataset for the Evaluator")
+    metric_names: Optional[List[str]] = hp.optional(
+        doc=textwrap.dedent("""Name of the metrics for the evaluator. Can be a torchmetrics metric name or the
+        class name of a metric returned by model.metrics(). If None (the default), uses all metrics in the model"""),
+        default=None)
+    eval_batch_size: Optional[int] = hp.optional(
+        doc="batch size to use for each evaluation step",
+        default=None,
+    )
+
+    def initialize_object(self, model: BaseMosaicModel, batch_size: int, dataloader_hparams: DataloaderHparams):
+        """Initialize an :class:`Evaluator`
+
+        If the Evaluatormetric_names is empty or None is provided, the function returns
+        a copy of all the model's default evaluation metrics.
+
+        Args:
+            model (BaseMosaicModel): The model, which is used to retrieve metric names
+            batch_size (int): The device batch size to use for the evaluation dataset
+            dataloader_hparams (DataloaderHparams): The hparams to use to construct a dataloader for the evaluation dataset
+
+        Returns:
+            Evaluator: The evaluator
+        """
+        evaluator_batch_size = self.eval_batch_size if self.eval_batch_size is not None else batch_size
+        dataloader = self.eval_dataset.initialize_object(batch_size=evaluator_batch_size,
+                                                         dataloader_hparams=dataloader_hparams)
+
+        # Get and copy all the model's associated evaluation metrics
+        model_metrics = model.metrics(train=False)
+        if isinstance(model_metrics, Metric):
+            # Forcing metrics to be a MetricCollection simplifies logging results
+            model_metrics = MetricCollection([model_metrics])
+
+        # Use all the metrics from the model if no metric_names are specified
+        if self.metric_names is None:
+            evaluator_metrics = copy.deepcopy(model_metrics)
+        else:
+            evaluator_metrics = MetricCollection([])
+            for metric_name in self.metric_names:
+                try:
+                    metric = model_metrics[metric_name]
+                except KeyError as e:
+                    raise RuntimeError(
+                        textwrap.dedent(f"""No metric found with the name {metric_name}. Check if this"
+                                       "metric is compatible/listed in your model metrics.""")) from e
+                assert isinstance(metric, Metric), "all values of a MetricCollection.__getitem__ should be a metric"
+                evaluator_metrics.add_metrics(copy.deepcopy(metric))
+            if len(evaluator_metrics) == 0:
+                raise RuntimeError(
+                    textwrap.dedent(f"""No metrics compatible with your model were added to this evaluator.
+                    Check that the metrics you specified are compatible/listed in your model."""))
+
+        return Evaluator(
+            label=self.label,
+            dataloader=dataloader,
+            metrics=evaluator_metrics,
+        )

composer/loggers/logger_hparams.py

@@ -156,7 +156,10 @@ def get_flattened_dict(data: Dict[str, Any], _prefix: List[str] = []) -> Dict[st
                             if isinstance(item, dict):
                                 found_sub_dicts = True
                                 for sub_key, sub_val in item.items():
-                                    all_items.update(get_flattened_dict(sub_val, key_items + [sub_key]))
+                                    if isinstance(sub_val, dict):
+                                        all_items.update(get_flattened_dict(sub_val, key_items + [sub_key]))
+                                    else:
+                                        all_items.update({sub_key: sub_val})
                         if not found_sub_dicts:
                             all_items[key_name] = val
                     else:

composer/loggers/tqdm_logger.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import collections.abc
 import sys
 from dataclasses import asdict, dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
@@ -108,9 +109,15 @@ def _start(self, state: State):
         if dist.get_global_rank() != 0:
             return
         assert self.is_train is not None, "self.is_train should be set by the callback"
-        # TODO(anis) -- in #120, len(state.eval_dataloader) is inaccurate, as it does not incorporate
-        # trainer._eval_subset_num_batches. The evaluator spec should fix this.
-        total_steps = state.steps_per_epoch if self.is_train else len(state.eval_dataloader)
+        if self.is_train:
+            total_steps = state.steps_per_epoch
+        else:
+            total_steps = 0
+            for evaluator in state.evaluators:
+                dataloader_spec = evaluator.dataloader
+                assert isinstance(dataloader_spec.dataloader, collections.abc.Sized)
+                total_steps += len(dataloader_spec.dataloader)
+
         desc = f'Epoch {int(state.timer.epoch)}'
         position = 0 if self.is_train else 1
         if not self.is_train:

composer/profiler/dataloader_profiler.py

@@ -49,4 +49,6 @@ def init(self, state: State, logger: Logger):
                 textwrap.dedent("""To use the dataloader profiler, state.profiler must be set.
                 Make sure to run composer with the profiler -- i.e. with the `--profiler` CLI flag."""))
         state.train_dataloader = ProfiledDataLoader(state.profiler, state.train_dataloader, "train")
-        state.eval_dataloader = ProfiledDataLoader(state.profiler, state.eval_dataloader, "eval")
+        for evaluator in state.evaluators:
+            evaluator.dataloader.dataloader = ProfiledDataLoader(state.profiler, evaluator.dataloader.dataloader,
+                                                                 evaluator.label)