Skip to content

Support flexible EpochMetric compute_fn outputs#3802

Open
Abhisri436 wants to merge 3 commits into
pytorch:masterfrom
Abhisri436:fix-issue-1757
Open

Support flexible EpochMetric compute_fn outputs#3802
Abhisri436 wants to merge 3 commits into
pytorch:masterfrom
Abhisri436:fix-issue-1757

Conversation

@Abhisri436

@Abhisri436 Abhisri436 commented Jul 3, 2026

Copy link
Copy Markdown

Fixes #1757

Description:

EpochMetric is currently typed and documented around scalar outputs, while issue #1757 requests support for richer compute_fn return types. My investigation also showed that unsupported values such as strings could pass through without clear validation in single-process mode.

This PR extends EpochMetric.compute_fn output handling beyond scalar values. It adds validation for supported output types including int, float, torch.Tensor, sequences, and mappings with string keys. Unsupported output types now raise a clear TypeError.

It also adds tests for scalar tensor, vector tensor, tuple/list tensor outputs, mapping outputs, and unsupported output behavior.

For distributed mode, this implementation takes a conservative approach: scalar and tensor outputs can still be broadcast, while tuple/list/mapping outputs raise a clear NotImplementedError when world_size > 1 instead of attempting unverified recursive distributed container broadcasting.

Validation:

  • Before: unsupported output types such as str could return without a clear error in single-process mode.
  • After: supported tensor, sequence, and mapping outputs are validated and tested, while unsupported output types raise a clear TypeError.
  • pytest tests/ignite/metrics/test_epoch_metric.py -vv → 17 passed, 57 skipped
  • WORLD_SIZE=2 distributed tests with the gloo backend → 10 passed
  • The distributed tests verify that non-scalar tensor outputs are consistent across ranks and that unsupported container/type outputs fail symmetrically without mismatched collective calls.
  • ruff check ignite/metrics/epoch_metric.py tests/ignite/metrics/test_epoch_metric.py → All checks passed
  • ruff format --check ignite/metrics/epoch_metric.py tests/ignite/metrics/test_epoch_metric.py → 2 files already formatted

Note: I saw that PR #3789 also addresses this issue. I had already selected and started working on issue #1757 before discovering the overlapping PR. I am submitting my completed branch for review and am happy to adjust this PR if maintainers prefer.

@github-actions github-actions Bot added the module: metrics Metrics module label Jul 3, 2026
@Abhisri436

Copy link
Copy Markdown
Author

Hi maintainers and reviewers! I’ve added permanent WORLD_SIZE=2 distributed tests and updated the PR with the latest validation results.

I also compared the distributed behavior with the overlapping PR and found a meaningful difference in how multi-rank outputs are handled. I’ve documented the tested behavior and added regression coverage in commit fcaef4cf.

I would appreciate a review when someone has time, and I’m happy to adjust or consolidate the approach if preferred. Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

module: metrics Metrics module

Projects

None yet

Development

Successfully merging this pull request may close these issues.

More EpochMetric's compute_fn output types

1 participant