diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 00000000..7e14f855 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,24 @@ +--- +name: Bug Report +about: Report a bug in the system +title: "[Bug]: " +labels: bug +--- + +### Description +A clear and concise description of the bug. + +### Steps to Reproduce +Give a detailed description of a minimal example to reproduce the bug. + +### Expected Behavior +Explain what you expected to happen. + +### Screenshots +Add screenshots if applicable. + +### Environment +Please describe your environment, like for example the OS you're using, your python version and your package versions. + +### Additional Context +Add any other context about the problem here. \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 00000000..bca2bc53 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,17 @@ +blank_issues_enabled: false +issue_templates: + - name: "Bug Report" + description: "Report a bug in the system." + title: "[Bug]: " + labels: ["bug"] + body: "./ISSUE_TEMPLATE/bug_report.md" + - name: "Feature Request" + description: "Propose a new feature or improvement." + title: "[Feature]: " + labels: ["enhancement"] + body: "./ISSUE_TEMPLATE/feature_request.md" + - name: "Documentation" + description: "Suggest updates or additions to the documentation." + title: "[Docs]: " + labels: ["documentation"] + body: "./ISSUE_TEMPLATE/documentation.md" \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/documentation.md b/.github/ISSUE_TEMPLATE/documentation.md new file mode 100644 index 00000000..744d3614 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/documentation.md @@ -0,0 +1,18 @@ +--- +name: Documentation +about: Suggest updates or additions to documentation +title: "[Docs]: " +labels: documentation +--- + +### Documentation Update +What part of the documentation needs to be updated or added? + +### Why Is This Needed? +Explain the importance of this update. + +### Suggested Changes +Provide a detailed description of the changes. + +### Additional Context +Include any related resources. \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 00000000..9964bb79 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,18 @@ +--- +name: Feature Request +about: Suggest a new feature or improvement +title: "[Feature]: " +labels: enhancement +--- + +### Feature Description +What feature would you like to see? + +### Why Is This Needed? +Explain the problem or need for this feature. + +### Suggested Solutions +Describe how this feature could be implemented. + +### Additional Context +Add any relevant screenshots, links, or resources. \ No newline at end of file diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 00000000..6c947bae --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,19 @@ +* **Please check if the PR fulfills these requirements** +- [ ] New and existing unit tests pass locally with my changes +- [ ] Tests for the changes have been added (for bug fixes/features) +- [ ] Docs have been added / updated (for bug fixes / features) + + +* **What kind of change does this PR introduce?** (Bug fix, feature, docs update, ...) + + +* **What is the current behavior?** (You can also link to an open issue here) + + +* **What is the new behavior (if this is a feature change)?** + + +* **Does this PR introduce a breaking change?** (What changes might users need to make in their application due to this PR?) + + +* **Other information**: \ No newline at end of file diff --git a/conf/metrics/example.conf b/conf/metrics/example.conf new file mode 100644 index 00000000..4075dc1e --- /dev/null +++ b/conf/metrics/example.conf @@ -0,0 +1,17 @@ +train_metrics: + - name: AUROC + params: + thresholds: 10 + average: macro + +valid_metrics: + - name: AUROC + params: + thresholds: 10 + average: macro + +test_metrics: + - name: AUROC + params: + thresholds: 10 + average: macro diff --git a/conf/optuna_hypers-class-imbalance.conf b/conf/optuna/optuna_hypers-class-imbalance.conf similarity index 100% rename from conf/optuna_hypers-class-imbalance.conf rename to conf/optuna/optuna_hypers-class-imbalance.conf diff --git a/conf/optuna_hypers-difficult-datasets.conf b/conf/optuna/optuna_hypers-difficult-datasets.conf similarity index 100% rename from conf/optuna_hypers-difficult-datasets.conf rename to conf/optuna/optuna_hypers-difficult-datasets.conf diff --git a/conf/optuna_hypers-difficult-examples-prediction-errors.conf b/conf/optuna/optuna_hypers-difficult-examples-prediction-errors.conf similarity index 100% rename from conf/optuna_hypers-difficult-examples-prediction-errors.conf rename to conf/optuna/optuna_hypers-difficult-examples-prediction-errors.conf diff --git a/conf/optuna_hypers-dpdl-benchmark.conf b/conf/optuna/optuna_hypers-dpdl-benchmark.conf similarity index 100% rename from conf/optuna_hypers-dpdl-benchmark.conf rename to conf/optuna/optuna_hypers-dpdl-benchmark.conf diff --git a/conf/optuna_hypers-hpo-alternatives.conf b/conf/optuna/optuna_hypers-hpo-alternatives.conf similarity index 100% rename from conf/optuna_hypers-hpo-alternatives.conf rename to conf/optuna/optuna_hypers-hpo-alternatives.conf diff --git a/conf/optuna_hypers-privacy-assistant-real-world.conf b/conf/optuna/optuna_hypers-privacy-assistant-real-world.conf similarity index 100% rename from conf/optuna_hypers-privacy-assistant-real-world.conf rename to conf/optuna/optuna_hypers-privacy-assistant-real-world.conf diff --git a/conf/optuna_hypers-privacy-assistant.conf b/conf/optuna/optuna_hypers-privacy-assistant.conf similarity index 100% rename from conf/optuna_hypers-privacy-assistant.conf rename to conf/optuna/optuna_hypers-privacy-assistant.conf diff --git a/conf/optuna_hypers-shots1.conf b/conf/optuna/optuna_hypers-shots1.conf similarity index 100% rename from conf/optuna_hypers-shots1.conf rename to conf/optuna/optuna_hypers-shots1.conf diff --git a/conf/optuna_hypers-shots10.conf b/conf/optuna/optuna_hypers-shots10.conf similarity index 100% rename from conf/optuna_hypers-shots10.conf rename to conf/optuna/optuna_hypers-shots10.conf diff --git a/conf/optuna_hypers-shots100.conf b/conf/optuna/optuna_hypers-shots100.conf similarity index 100% rename from conf/optuna_hypers-shots100.conf rename to conf/optuna/optuna_hypers-shots100.conf diff --git a/conf/optuna_hypers-shots25.conf b/conf/optuna/optuna_hypers-shots25.conf similarity index 100% rename from conf/optuna_hypers-shots25.conf rename to conf/optuna/optuna_hypers-shots25.conf diff --git a/conf/optuna_hypers-shots250.conf b/conf/optuna/optuna_hypers-shots250.conf similarity index 100% rename from conf/optuna_hypers-shots250.conf rename to conf/optuna/optuna_hypers-shots250.conf diff --git a/conf/optuna_hypers-shots5.conf b/conf/optuna/optuna_hypers-shots5.conf similarity index 100% rename from conf/optuna_hypers-shots5.conf rename to conf/optuna/optuna_hypers-shots5.conf diff --git a/conf/optuna_hypers-shots50.conf b/conf/optuna/optuna_hypers-shots50.conf similarity index 100% rename from conf/optuna_hypers-shots50.conf rename to conf/optuna/optuna_hypers-shots50.conf diff --git a/conf/optuna_hypers-shots500.conf b/conf/optuna/optuna_hypers-shots500.conf similarity index 100% rename from conf/optuna_hypers-shots500.conf rename to conf/optuna/optuna_hypers-shots500.conf diff --git a/conf/optuna_hypers-subset0.1.conf b/conf/optuna/optuna_hypers-subset0.1.conf similarity index 100% rename from conf/optuna_hypers-subset0.1.conf rename to conf/optuna/optuna_hypers-subset0.1.conf diff --git a/conf/optuna_hypers-subset1.0.conf b/conf/optuna/optuna_hypers-subset1.0.conf similarity index 100% rename from conf/optuna_hypers-subset1.0.conf rename to conf/optuna/optuna_hypers-subset1.0.conf diff --git a/conf/optuna_hypers-us-vs-hyfreedp.conf b/conf/optuna/optuna_hypers-us-vs-hyfreedp.conf similarity index 100% rename from conf/optuna_hypers-us-vs-hyfreedp.conf rename to conf/optuna/optuna_hypers-us-vs-hyfreedp.conf diff --git a/conf/optuna_hypers-weight-perturbation.conf b/conf/optuna/optuna_hypers-weight-perturbation.conf similarity index 100% rename from conf/optuna_hypers-weight-perturbation.conf rename to conf/optuna/optuna_hypers-weight-perturbation.conf diff --git a/conf/optuna_hypers.conf b/conf/optuna/optuna_hypers.conf similarity index 100% rename from conf/optuna_hypers.conf rename to conf/optuna/optuna_hypers.conf diff --git a/conf/optuna_hypers_ordered.conf b/conf/optuna/optuna_hypers_ordered.conf similarity index 100% rename from conf/optuna_hypers_ordered.conf rename to conf/optuna/optuna_hypers_ordered.conf diff --git a/conf/optuna_trials-cifar100.conf b/conf/optuna/optuna_trials-cifar100.conf similarity index 100% rename from conf/optuna_trials-cifar100.conf rename to conf/optuna/optuna_trials-cifar100.conf diff --git a/conf/optuna_trials-cifar10_10pct_plus_cifar100_humans.conf b/conf/optuna/optuna_trials-cifar10_10pct_plus_cifar100_humans.conf similarity index 100% rename from conf/optuna_trials-cifar10_10pct_plus_cifar100_humans.conf rename to conf/optuna/optuna_trials-cifar10_10pct_plus_cifar100_humans.conf diff --git a/conf/optuna_trials-svhn_cropped_balanced-avg.conf b/conf/optuna/optuna_trials-svhn_cropped_balanced-avg.conf similarity index 100% rename from conf/optuna_trials-svhn_cropped_balanced-avg.conf rename to conf/optuna/optuna_trials-svhn_cropped_balanced-avg.conf diff --git a/conf/optuna_trials-svhn_cropped_balanced-bad.conf b/conf/optuna/optuna_trials-svhn_cropped_balanced-bad.conf similarity index 100% rename from conf/optuna_trials-svhn_cropped_balanced-bad.conf rename to conf/optuna/optuna_trials-svhn_cropped_balanced-bad.conf diff --git a/conf/optuna_trials-svhn_cropped_balanced-good.conf b/conf/optuna/optuna_trials-svhn_cropped_balanced-good.conf similarity index 100% rename from conf/optuna_trials-svhn_cropped_balanced-good.conf rename to conf/optuna/optuna_trials-svhn_cropped_balanced-good.conf diff --git a/conf/optuna_trials-svhn_cropped_balanced.conf b/conf/optuna/optuna_trials-svhn_cropped_balanced.conf similarity index 100% rename from conf/optuna_trials-svhn_cropped_balanced.conf rename to conf/optuna/optuna_trials-svhn_cropped_balanced.conf diff --git a/conf/optuna_trials.conf b/conf/optuna/optuna_trials.conf similarity index 100% rename from conf/optuna_trials.conf rename to conf/optuna/optuna_trials.conf diff --git a/docs/hyperparameter-optimization.md b/docs/hyperparameter-optimization.md index fd71ee40..6e335a0c 100644 --- a/docs/hyperparameter-optimization.md +++ b/docs/hyperparameter-optimization.md @@ -6,9 +6,9 @@ We aim to keep the document brief and point to code for deeper understanding. Relevant sources: - Optimizer logic: [`dpdl/hyperparameteroptimizer.py`](../dpdl/hyperparameteroptimizer.py) -- Search space example: [`conf/optuna_hypers.conf`](../conf/optuna_hypers.conf) -- Ordered search space example: [`conf/optuna_hypers_ordered.conf`](../conf/optuna_hypers_ordered.conf) -- Manual trials example: [`conf/optuna_trials.conf`](../conf/optuna_trials.conf) +- Search space example: [`conf/optuna_hypers.conf`](../conf/optuna/optuna_hypers.conf) +- Ordered search space example: [`conf/optuna_hypers_ordered.conf`](../conf/optuna/optuna_hypers_ordered.conf) +- Manual trials example: [`conf/optuna_trials.conf`](../conf/optuna/optuna_trials.conf) ## Overview @@ -31,14 +31,14 @@ Supported types: In the config, we can use `-1` to indicate the [maximum batch size](../dpdl/hyperparameteroptimizer.py#L146) (i.e. full dataset). -See [`conf/optuna_hypers.conf`](../conf/optuna_hypers.conf) for an example. +See [`conf/optuna_hypers.conf`](../conf/optuna/optuna_hypers.conf) for an example. ## Manual trials (`conf/optuna_trials.conf`) It is also possible to configure manual trials. These define hypers that Optuna will try before starting its own algorithm for suggesting the next hypers. -In [`conf/optuna_trials.conf`](../conf/optuna_trials.conf), each entry is a dictionary of hyperparameters and the respective values to try. +In [`conf/optuna_trials.conf`](../conf/optuna/optuna_trials.conf), each entry is a dictionary of hyperparameters and the respective values to try. ## Ordered discrete type (`ordered`) @@ -48,7 +48,7 @@ This enables, for example, using discrete learning rate values while still being Gotchas: - For convenience, for `batch_size` with `ordered` type, the values are auto‑generated as powers of two from `min` up to the dataset size, plus `-1` for full batch. -- For other ordered parameters (e.g., `learning_rate`), you must provide `options` explicitly (see [`conf/optuna_hypers_ordered.conf`](../conf/optuna_hypers_ordered.conf) for an example). +- For other ordered parameters (e.g., `learning_rate`), you must provide `options` explicitly (see [`conf/optuna_hypers_ordered.conf`](../conf/optuna/optuna_hypers_ordered.conf) for an example). ## CLI flags and behavior diff --git a/docs/trainer.md b/docs/trainer.md index 0f2e2ce9..ad548182 100644 --- a/docs/trainer.md +++ b/docs/trainer.md @@ -24,6 +24,32 @@ We support [torchmetrics](https://github.com/Lightning-AI/torchmetrics) through Metrics live in the [model](../dpdl/models/model_base.py) (`train_metrics`, `valid_metrics`, `test_metrics`) and the Trainer(s) update them during different training phases. +### Built-in metrics + +`MetricsFactory` always initialises a task-appropriate baseline set: + +- **Classification** (`ImageClassification`, `SequenceClassification`): macro and micro `MulticlassAccuracy` per class, and a `ConfusionMatrix` on the test split. +- **Language modelling** (`CausalLM`, `InstructLM`): micro `MulticlassAccuracy` and `Perplexity`. + +### Custom metrics + +Additional metrics can be declared in a YAML config file and passed via `--metric-config`. The file has three optional sections (`train_metrics`, `valid_metrics`, `test_metrics`), each a list of entries. An entry is either a plain metric class name or a mapping: + +```yaml +train_metrics: + - Accuracy # shorthand: just the class name + +valid_metrics: + - name: AUROC # torchmetrics class name + params: + thresholds: 10 + average: macro +``` + +See [`conf/metrics/example.conf`](../conf/metrics/example.conf) for a working example. + +`CustomMetricsFactory` resolves each name against the torchmetrics submodule hierarchy (`torchmetrics`, `torchmetrics.classification`, `torchmetrics.text`, `torchmetrics.regression`, `torchmetrics.image`, `torchmetrics.audio`). Task-level defaults such as `num_classes` and `task` are injected automatically where the metric accepts them; `sync_on_compute` is also set to match the train/eval setting. Params declared in the config always take precedence over these defaults. + ## Task adapters The trainers use **task adapters** to keep task-specific logic out of the training loop. diff --git a/dpdl/cli.py b/dpdl/cli.py index 969535a7..57592f5a 100644 --- a/dpdl/cli.py +++ b/dpdl/cli.py @@ -564,6 +564,13 @@ def cli( rich_help_panel='Bayesian optimization (Optuna) options', ) ] = 'conf/optuna_hypers.conf', + metric_config: Annotated[ + Optional[str], + typer.Option( + help='YAML file path containing custom metrics for train/validation/test', + rich_help_panel='Logging options', + ) + ] = None, optuna_manual_trials: Annotated[ Optional[str], typer.Option( diff --git a/dpdl/configurationmanager.py b/dpdl/configurationmanager.py index 153830c5..8b005881 100644 --- a/dpdl/configurationmanager.py +++ b/dpdl/configurationmanager.py @@ -103,6 +103,7 @@ class Configuration(BaseModel): optuna_target_metric: str = 'loss' optuna_direction: Literal['minimize', 'maximize'] = 'minimize' optuna_config: str = 'conf/optuna_hypers.conf' + metric_config: Optional[str] = None optuna_manual_trials: Optional[str] = None optuna_journal: str = 'optuna.journal' optuna_resume: bool = False @@ -254,6 +255,7 @@ def __str__(self): ('Enable callback debug logging', self.verbose_callback), ('Fairness-style subsampling class', self.fairness_imbalance_class), ('Random seed for creating dataset subsets', self.split_seed), + ('Metric configuration', self.metric_config), ('LLM use', self.llm), ('Task', self.task), ] diff --git a/dpdl/metrics_factory.py b/dpdl/metrics_factory.py index b6a9f355..84b68da1 100644 --- a/dpdl/metrics_factory.py +++ b/dpdl/metrics_factory.py @@ -1,52 +1,97 @@ -from dataclasses import dataclass -from typing import Optional, Dict -from torchmetrics.text import Perplexity - import logging +from typing import Any, Dict, Optional + import torch import torchmetrics +from torchmetrics.text import Perplexity +import yaml + log = logging.getLogger(__name__) -def _get_classification_metrics( - output_dim: int, - sync: bool, - with_confusion_matrix: bool, -) -> torchmetrics.MetricCollection: - # NB: If `sync_on_compute` is enabled, this breaks - # distributed training. If this needs to be enabled, - # then we also need to actually run the validation on - # all the GPUs. - metrics = { - 'MulticlassAccuracy': torchmetrics.classification.MulticlassAccuracy( - num_classes=output_dim, - average='macro', - sync_on_compute=sync, - ), - 'MulticlassAccuracyWithMicro': torchmetrics.classification.MulticlassAccuracy( - num_classes=output_dim, - average='micro', - sync_on_compute=sync, - ), - 'MulticlassAccuracyPerClass': torchmetrics.classification.MulticlassAccuracy( - num_classes=output_dim, - average='none', - sync_on_compute=sync, - ), - } - - if with_confusion_matrix: - metrics['ConfusionMatrix'] = torchmetrics.ConfusionMatrix( - task='multiclass' if output_dim > 2 else 'binary', - num_classes=output_dim, - sync_on_compute=sync, - ) - - return torchmetrics.MetricCollection(metrics) +_METRIC_MODULES = ( + torchmetrics, + getattr(torchmetrics, 'classification', None), + getattr(torchmetrics, 'text', None), + getattr(torchmetrics, 'regression', None), + getattr(torchmetrics, 'image', None), + getattr(torchmetrics, 'audio', None), +) + + +def _build_custom_metrics(metric_config, key, default_kwargs, sync_on_compute): + if not metric_config: + return None + return CustomMetricsFactory.build_metric_collection(metric_config[key], default_kwargs, sync_on_compute) + + +class ClassificationMetrics(torchmetrics.MetricCollection): + def __init__( + self, + output_dim: int, + sync: bool, + with_confusion_matrix: bool, + custom_metrics: Optional[Dict[str, torchmetrics.Metric]] = None, + ) -> None: + # NB: If `sync_on_compute` is enabled, this breaks + # distributed training. If this needs to be enabled, + # then we also need to actually run the validation on + # all the GPUs. + metrics = { + 'MulticlassAccuracy': torchmetrics.classification.MulticlassAccuracy( + num_classes=output_dim, + average='macro', + sync_on_compute=sync, + ), + 'MulticlassAccuracyWithMicro': torchmetrics.classification.MulticlassAccuracy( + num_classes=output_dim, + average='micro', + sync_on_compute=sync, + ), + 'MulticlassAccuracyPerClass': torchmetrics.classification.MulticlassAccuracy( + num_classes=output_dim, + average='none', + sync_on_compute=sync, + ), + } + + if with_confusion_matrix: + metrics['ConfusionMatrix'] = torchmetrics.ConfusionMatrix( + task='multiclass' if output_dim > 2 else 'binary', + num_classes=output_dim, + sync_on_compute=sync, + ) + + if custom_metrics: + metrics.update(custom_metrics) + + super().__init__(metrics) + + def update(self, preds, target) -> None: + # ConfusionMatrix expects class labels, not logits, unlike the + # other classification metrics, which apply argmax internally. + if 'ConfusionMatrix' not in self.keys(): + return super().update(preds, target) + + preds_labels = preds.argmax(dim=-1) if preds.ndim > target.ndim else preds + + for name, metric in self.items(): + if name == 'ConfusionMatrix': + metric.update(preds_labels, target) + else: + metric.update(preds, target) + + return None class LanguageModelMetrics(torchmetrics.MetricCollection): - def __init__(self, vocab_size: int, ignore_index: int, sync: bool) -> None: + def __init__( + self, + vocab_size: int, + ignore_index: int, + sync: bool, + custom_metrics: Optional[Dict[str, torchmetrics.Metric]] = None, + ) -> None: metrics = { 'MulticlassAccuracy': torchmetrics.classification.MulticlassAccuracy( num_classes=vocab_size, @@ -59,44 +104,120 @@ def __init__(self, vocab_size: int, ignore_index: int, sync: bool) -> None: sync_on_compute=sync, ), } + + if custom_metrics: + metrics.update(custom_metrics) + super().__init__(metrics) def update(self, preds, target) -> None: - # Accuracy metrics use standard flattened inputs - if not hasattr(preds, 'ndim'): + # Perplexity expects 3D logits and 2D labels, unlike the other + # metrics, which expect flattened inputs. + if not hasattr(preds, 'ndim') or preds.ndim != 3: return super().update(preds, target) - # We need to shape the data for perplexity that expects 3D logits and 2D labels - if preds.ndim == 3: - shift_logits = preds[:, :-1, :].contiguous() # (batch, seq_len-1, vocab) - shift_labels = target[:, 1:].contiguous() # (batch, seq_len-1) - shift_logits_flat = shift_logits.view(-1, shift_logits.size(-1)) # (batch*(seq_len-1), vocab) - shift_labels_flat = shift_labels.view(-1) # (batch*(seq_len-1)) + shift_logits = preds[:, :-1, :].contiguous() # (batch, seq_len-1, vocab) + shift_labels = target[:, 1:].contiguous() # (batch, seq_len-1) + shift_logits_flat = shift_logits.view(-1, shift_logits.size(-1)) # (batch*(seq_len-1), vocab) + shift_labels_flat = shift_labels.view(-1) # (batch*(seq_len-1)) - self['Perplexity'].update(shift_logits, shift_labels) + for name, metric in self.items(): + if name == 'Perplexity': + metric.update(shift_logits, shift_labels) + else: + metric.update(shift_logits_flat, shift_labels_flat) - for name, metric in self.items(): - if name == 'Perplexity': - continue + return None - metric.update(shift_logits_flat, shift_labels_flat) - return +class CustomMetricsFactory: + @staticmethod + def _resolve_metric_class(metric_name: str): + for module in _METRIC_MODULES: + if module and hasattr(module, metric_name): + return getattr(module, metric_name) + raise ValueError(f'Metric class "{metric_name}" not found in torchmetrics.') - return super().update(preds, target) + @staticmethod + def _build_metric(metric_spec: Dict[str, Any], default_kwargs: Dict[str, Any], sync_on_compute: bool): + metric_name = metric_spec['name'] + metric_cls = CustomMetricsFactory._resolve_metric_class(metric_name) + user_params = dict(metric_spec.get('params') or {}) + + params = default_kwargs.copy() + params.update(user_params) -def _get_language_model_metrics( - vocab_size: int, - ignore_index: int, - sync: bool, -) -> torchmetrics.MetricCollection: + if 'sync_on_compute' not in params: + params['sync_on_compute'] = sync_on_compute + + return metric_cls(**params) + + @staticmethod + def _normalize_metric_entries(metric_entries, section_name: str, metric_config: str) -> list[Dict[str, Any]]: + normalized: list[Dict[str, Any]] = [] + + for idx, entry in enumerate(metric_entries): + if isinstance(entry, str): + entry = {'name': entry, 'params': {}} + elif not isinstance(entry, dict): + raise ValueError( + f'Metric config file "{metric_config}" section "{section_name}" entry #{idx + 1} must be a mapping or string.' + ) + + metric_name = entry.get('name') + if not metric_name: + raise ValueError( + f'Metric config file "{metric_config}" section "{section_name}" entry #{idx + 1} is missing a metric name.' + ) + + params = entry.get('params') or {} + if not isinstance(params, dict): + raise ValueError( + f'Metric config file "{metric_config}" section "{section_name}" entry #{idx + 1} has invalid params.' + ) + + normalized.append({'name': metric_name, 'params': params}) + + return normalized + + @staticmethod + def read_metric_config(metric_config: Optional[str]) -> Dict[str, list[Dict[str, Any]]]: + if not metric_config: + return {'train_metrics': [], 'valid_metrics': [], 'test_metrics': []} + + with open(metric_config, 'rb') as fh: + raw_config = yaml.safe_load(fh) or {} + + if not isinstance(raw_config, dict): + raise ValueError(f'Metric config file "{metric_config}" must contain a mapping at the top level.') + + normalized_config = { + 'train_metrics': [], + 'valid_metrics': [], + 'test_metrics': [], + } + + for key in ('train_metrics', 'valid_metrics', 'test_metrics'): + entries = raw_config.get(key, None) + if entries is None: + continue + if not isinstance(entries, list): + raise ValueError( + f'Metric config file "{metric_config}" section "{key}" must contain a list of metrics.' + ) + + normalized_config[key].extend(CustomMetricsFactory._normalize_metric_entries(entries, key, metric_config)) + + return normalized_config + + @staticmethod + def build_metric_collection(metric_specs, default_kwargs: Dict[str, Any], sync_on_compute: bool) -> Dict[str, torchmetrics.Metric]: + return { + spec['name']: CustomMetricsFactory._build_metric(spec, default_kwargs, sync_on_compute) + for spec in metric_specs + } - return LanguageModelMetrics( - vocab_size=vocab_size, - ignore_index=ignore_index, - sync=sync, - ) class MetricsFactory: @@ -107,6 +228,8 @@ def get_metrics( output_dim: Optional[int] = None, ) -> Dict[str, torchmetrics.MetricCollection]: task = configuration.task + metric_config = getattr(configuration, 'metric_config', None) + metric_config = CustomMetricsFactory.read_metric_config(metric_config) if metric_config else None # we only validate on rank 0, so there's no need to # synchronize when calculating the metrics. @@ -119,20 +242,28 @@ def get_metrics( if not output_dim or output_dim < 1: raise ValueError('output_dim required for classification tasks') - train = _get_classification_metrics( + classification_defaults = { + 'num_classes': output_dim, + 'task': 'multiclass' if output_dim > 2 else 'binary', + } + + train = ClassificationMetrics( output_dim=output_dim, sync=train_sync, with_confusion_matrix=False, + custom_metrics=_build_custom_metrics(metric_config, 'train_metrics', classification_defaults, train_sync), ) - valid = _get_classification_metrics( + valid = ClassificationMetrics( output_dim=output_dim, sync=eval_sync, with_confusion_matrix=False, + custom_metrics=_build_custom_metrics(metric_config, 'valid_metrics', classification_defaults, eval_sync), ) - test = _get_classification_metrics( + test = ClassificationMetrics( output_dim=output_dim, sync=eval_sync, with_confusion_matrix=True, + custom_metrics=_build_custom_metrics(metric_config, 'test_metrics', classification_defaults, eval_sync), ) elif task in ('CausalLM', 'InstructLM'): @@ -142,20 +273,29 @@ def get_metrics( vocab_size = int(output_dim) ignore_index = -100 - train = _get_language_model_metrics( + language_defaults = { + 'num_classes': vocab_size, + 'ignore_index': ignore_index, + 'task': 'multiclass', + } + + train = LanguageModelMetrics( vocab_size=vocab_size, ignore_index=ignore_index, sync=train_sync, + custom_metrics=_build_custom_metrics(metric_config, 'train_metrics', language_defaults, train_sync), ) - valid = _get_language_model_metrics( + valid = LanguageModelMetrics( vocab_size=vocab_size, ignore_index=ignore_index, sync=eval_sync, + custom_metrics=_build_custom_metrics(metric_config, 'valid_metrics', language_defaults, eval_sync), ) - test = _get_language_model_metrics( + test = LanguageModelMetrics( vocab_size=vocab_size, ignore_index=ignore_index, sync=eval_sync, + custom_metrics=_build_custom_metrics(metric_config, 'test_metrics', language_defaults, eval_sync), ) else: diff --git a/dpdl/trainer.py b/dpdl/trainer.py index bd688f6f..9eb7b845 100644 --- a/dpdl/trainer.py +++ b/dpdl/trainer.py @@ -771,8 +771,7 @@ def update_metrics(self, model, batch, forward_output, metrics = None): else: metrics_to_update = model.train_metrics if model.training else model.valid_metrics - preds = torch.argmax(forward_output, dim=1) - metrics_to_update.update(preds, y) + metrics_to_update.update(forward_output, y) class LanguageModelAdapter(TaskAdapter): diff --git a/tests/test_metrics_factory.py b/tests/test_metrics_factory.py new file mode 100644 index 00000000..0e98e176 --- /dev/null +++ b/tests/test_metrics_factory.py @@ -0,0 +1,60 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip('torch') + +from dpdl.metrics_factory import ClassificationMetrics, MetricsFactory, CustomMetricsFactory + + +def _example_metric_config() -> Path: + return Path('conf/metrics/example.conf') + + +def test_read_metric_config_uses_example_conf() -> None: + config = CustomMetricsFactory.read_metric_config(str(_example_metric_config())) + + assert [metric['name'] for metric in config['train_metrics']] == ['AUROC'] + assert [metric['name'] for metric in config['valid_metrics']] == ['AUROC'] + assert [metric['name'] for metric in config['test_metrics']] == ['AUROC'] + + +def test_get_metrics_merges_custom_classification_metrics(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.distributed, 'get_rank', lambda: 0) + + configuration = SimpleNamespace(task='ImageClassification', metric_config=str(_example_metric_config())) + metrics = MetricsFactory.get_metrics(configuration, output_dim=3) + + assert 'AUROC' in metrics['train_metrics'].keys() + assert 'AUROC' in metrics['valid_metrics'].keys() + assert 'AUROC' in metrics['test_metrics'].keys() + assert 'MulticlassAccuracy' in metrics['train_metrics'].keys() + assert 'ConfusionMatrix' in metrics['test_metrics'].keys() + + +def test_classification_metrics_update_argmaxes_logits_for_confusion_matrix() -> None: + metrics = ClassificationMetrics(output_dim=3, sync=False, with_confusion_matrix=True) + + logits = torch.tensor([[2.0, 0.5, 0.1], [0.1, 0.2, 3.0]]) + target = torch.tensor([0, 2]) + + metrics.update(logits, target) + + confusion_matrix = metrics['ConfusionMatrix'].compute() + assert confusion_matrix.sum().item() == 2 + assert confusion_matrix[0, 0].item() == 1 + assert confusion_matrix[2, 2].item() == 1 + + +def test_get_metrics_merges_custom_language_metrics(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.distributed, 'get_rank', lambda: 0) + + configuration = SimpleNamespace(task='CausalLM', metric_config=str(_example_metric_config())) + metrics = MetricsFactory.get_metrics(configuration, output_dim=8) + + assert 'AUROC' in metrics['train_metrics'].keys() + assert 'AUROC' in metrics['valid_metrics'].keys() + assert 'AUROC' in metrics['test_metrics'].keys() + assert 'MulticlassAccuracy' in metrics['train_metrics'].keys() + assert 'Perplexity' in metrics['test_metrics'].keys()