From f6dc05d96823b41c92c11d37edeb5821066e9c01 Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Mon, 11 May 2026 14:10:04 +0300 Subject: [PATCH 1/9] feat(metrics): add WelfordVariance and WelfordCovariance helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces ignite/metrics/_running_stats.py with two numerically stable running-statistics primitives that variance- and covariance-bearing metrics can share, instead of each one rolling its own naive Σx² − (Σx)²/n implementation. WelfordVariance mean, variance, std for a single variable. WelfordCovariance variance_x, variance_y, covariance, and Pearson correlation for a paired (x, y) stream. Both classes: - keep internal state in float64 regardless of input dtype, so the classic E[X²] − E[X]² cancellation does not bite at large means; - update incrementally via Welford's online algorithm; - expose merge() implementing the Chan / Welford parallel formula, suitable for cross-rank distributed reduction or any other case where two accumulators need to be combined without re-iterating the raw data. This is PR 1 of the plan in #3748. Follow-ups: PR 2 will port R2Score (#3662-style regression test attached). PR 3 will refactor #3741 to consume WelfordCovariance instead of its current inline Welford state. Tests (20 total, all passing): - per-class correctness vs numpy mean / var / cov / corrcoef - multi-batch update matches single-batch update - merge matches concatenated update - merge with empty accumulators on either side - numerical-stability regression (mean=1e6 in float32) for both classes, with an assertion that the naive float32 formula actually does fail on the same data so the test documents what we're protecting against - shape-mismatch raises ValueError - empty-batch update is a no-op - reset clears state - input dtypes (int32) upcast to float64 correctly - cross-class sanity: WelfordCovariance.variance_x matches WelfordVariance.variance fed the same x --- ignite/metrics/_running_stats.py | 273 ++++++++++++++++++++ tests/ignite/metrics/test_running_stats.py | 277 +++++++++++++++++++++ 2 files changed, 550 insertions(+) create mode 100644 ignite/metrics/_running_stats.py create mode 100644 tests/ignite/metrics/test_running_stats.py diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py new file mode 100644 index 000000000000..206353b9cabd --- /dev/null +++ b/ignite/metrics/_running_stats.py @@ -0,0 +1,273 @@ +"""Numerically stable running variance and covariance helpers. + +Shared by metrics that need to accumulate variance / covariance from +streaming batches without falling into the catastrophic-cancellation +trap of the naive ``E[X^2] - E[X]^2`` formula. Used by +:class:`~ignite.metrics.regression.PearsonCorrelation` and +:class:`~ignite.metrics.regression.R2Score`; new metrics with the same +need should consume these helpers rather than rolling their own. + +Both classes keep internal state in ``float64`` regardless of the input +dtype, follow Welford's online algorithm for incremental updates, and +fold accumulators together with the Chan / Welford parallel formula +(used both for batch-wise updates and for cross-rank merges in +distributed settings). + +References: + Welford, B. P. (1962). Note on a method for calculating corrected + sums of squares and products. Technometrics 4 (3), 419 to 420. + Chan, T. F., Golub, G. H., LeVeque, R. J. (1979). Updating formulae + and a pairwise algorithm for computing sample variances. + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm +""" + +from typing import Union + +import torch + + +class WelfordVariance: + """Numerically stable running mean and variance via Welford's algorithm. + + Accumulates samples in batches via :meth:`update` and reads off the + mean, variance, or standard deviation through the corresponding + properties. Two accumulators can be combined with :meth:`merge`, + which uses the Chan / Welford parallel formula and is the basis for + distributed reductions. + + State is kept in ``float64`` regardless of input dtype so that the + classic ``E[X^2] - E[X]^2`` cancellation does not bite for inputs + with large means. + + Args: + device: device on which to keep the running-state tensors. + Default: ``"cpu"``. + + Example: + + .. code-block:: python + + ws = WelfordVariance() + for batch in stream: + ws.update(batch) + print(ws.mean.item(), ws.variance.item(), ws.std.item()) + """ + + n_samples: int + mean: torch.Tensor + sum_sq_dev_from_mean: torch.Tensor + + def __init__(self, device: Union[str, torch.device] = "cpu") -> None: + self._device = torch.device(device) + self.reset() + + def reset(self) -> None: + """Drop all accumulated state.""" + self.n_samples = 0 + self.mean = torch.tensor(0.0, dtype=torch.float64, device=self._device) + self.sum_sq_dev_from_mean = torch.tensor(0.0, dtype=torch.float64, device=self._device) + + @torch.no_grad() + def update(self, batch: torch.Tensor) -> None: + """Fold a batch of samples into the running state. + + Empty batches are silently ignored. Inputs of any dtype are + upcast to ``float64`` for the internal computation. + """ + if batch.numel() == 0: + return + batch64 = batch.detach().to(dtype=torch.float64).flatten() + n_b = batch64.shape[0] + + mean_b = batch64.mean().to(self._device) + m2_b = (batch64 - batch64.mean()).square().sum().to(self._device) + + if self.n_samples == 0: + self.mean = mean_b + self.sum_sq_dev_from_mean = m2_b + self.n_samples = n_b + return + + n_a = self.n_samples + n_ab = n_a + n_b + delta = mean_b - self.mean + self.mean = self.mean + delta * n_b / n_ab + self.sum_sq_dev_from_mean = self.sum_sq_dev_from_mean + m2_b + delta * delta * n_a * n_b / n_ab + self.n_samples = n_ab + + def merge(self, other: "WelfordVariance") -> None: + """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. + + Used both for cross-rank reduction in distributed settings and + for any other case where two accumulators need to be combined + into one without re-iterating the raw data. + """ + if other.n_samples == 0: + return + if self.n_samples == 0: + self.n_samples = other.n_samples + self.mean = other.mean.detach().clone().to(self._device) + self.sum_sq_dev_from_mean = other.sum_sq_dev_from_mean.detach().clone().to(self._device) + return + n_a = self.n_samples + n_b = other.n_samples + n_ab = n_a + n_b + delta = other.mean.to(self._device) - self.mean + self.mean = self.mean + delta * n_b / n_ab + self.sum_sq_dev_from_mean = ( + self.sum_sq_dev_from_mean + other.sum_sq_dev_from_mean.to(self._device) + delta * delta * n_a * n_b / n_ab + ) + self.n_samples = n_ab + + @property + def variance(self) -> torch.Tensor: + """Population variance (divisor ``n``). Returns ``0.0`` when empty.""" + if self.n_samples == 0: + return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return torch.clamp(self.sum_sq_dev_from_mean / self.n_samples, min=0.0) + + @property + def std(self) -> torch.Tensor: + """Population standard deviation (divisor ``n``).""" + return self.variance.sqrt() + + +class WelfordCovariance: + """Numerically stable running covariance for a pair of variables (x, y). + + Exposes :attr:`variance_x`, :attr:`variance_y`, :attr:`covariance`, + and :meth:`correlation` (Pearson) through the same Welford-style + online update + Chan / Welford parallel merge as + :class:`WelfordVariance`. + + Args: + device: device on which to keep the running-state tensors. + Default: ``"cpu"``. + """ + + n_samples: int + mean_x: torch.Tensor + mean_y: torch.Tensor + sum_sq_dev_x: torch.Tensor + sum_sq_dev_y: torch.Tensor + sum_product_of_devs: torch.Tensor + + def __init__(self, device: Union[str, torch.device] = "cpu") -> None: + self._device = torch.device(device) + self.reset() + + def reset(self) -> None: + """Drop all accumulated state.""" + self.n_samples = 0 + self.mean_x = torch.tensor(0.0, dtype=torch.float64, device=self._device) + self.mean_y = torch.tensor(0.0, dtype=torch.float64, device=self._device) + self.sum_sq_dev_x = torch.tensor(0.0, dtype=torch.float64, device=self._device) + self.sum_sq_dev_y = torch.tensor(0.0, dtype=torch.float64, device=self._device) + self.sum_product_of_devs = torch.tensor(0.0, dtype=torch.float64, device=self._device) + + @torch.no_grad() + def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: + """Fold a paired batch ``(x_i, y_i)`` into the running state. + + ``batch_x`` and ``batch_y`` must have the same shape; both are + flattened internally and upcast to ``float64``. + """ + if batch_x.shape != batch_y.shape: + raise ValueError( + f"batch_x and batch_y must have the same shape, got {tuple(batch_x.shape)} and {tuple(batch_y.shape)}." + ) + if batch_x.numel() == 0: + return + + x64 = batch_x.detach().to(dtype=torch.float64).flatten() + y64 = batch_y.detach().to(dtype=torch.float64).flatten() + n_b = x64.shape[0] + + mean_x_b = x64.mean().to(self._device) + mean_y_b = y64.mean().to(self._device) + dx_b = x64 - mean_x_b + dy_b = y64 - mean_y_b + m2_x_b = dx_b.square().sum().to(self._device) + m2_y_b = dy_b.square().sum().to(self._device) + cxy_b = (dx_b * dy_b).sum().to(self._device) + + if self.n_samples == 0: + self.mean_x = mean_x_b + self.mean_y = mean_y_b + self.sum_sq_dev_x = m2_x_b + self.sum_sq_dev_y = m2_y_b + self.sum_product_of_devs = cxy_b + self.n_samples = n_b + return + + n_a = self.n_samples + n_ab = n_a + n_b + cross = n_a * n_b / n_ab + delta_x = mean_x_b - self.mean_x + delta_y = mean_y_b - self.mean_y + + self.mean_x = self.mean_x + delta_x * n_b / n_ab + self.mean_y = self.mean_y + delta_y * n_b / n_ab + self.sum_sq_dev_x = self.sum_sq_dev_x + m2_x_b + delta_x * delta_x * cross + self.sum_sq_dev_y = self.sum_sq_dev_y + m2_y_b + delta_y * delta_y * cross + self.sum_product_of_devs = self.sum_product_of_devs + cxy_b + delta_x * delta_y * cross + self.n_samples = n_ab + + def merge(self, other: "WelfordCovariance") -> None: + """Combine ``other`` into ``self`` using the Chan / Welford parallel formula.""" + if other.n_samples == 0: + return + if self.n_samples == 0: + self.n_samples = other.n_samples + self.mean_x = other.mean_x.detach().clone().to(self._device) + self.mean_y = other.mean_y.detach().clone().to(self._device) + self.sum_sq_dev_x = other.sum_sq_dev_x.detach().clone().to(self._device) + self.sum_sq_dev_y = other.sum_sq_dev_y.detach().clone().to(self._device) + self.sum_product_of_devs = other.sum_product_of_devs.detach().clone().to(self._device) + return + n_a = self.n_samples + n_b = other.n_samples + n_ab = n_a + n_b + cross = n_a * n_b / n_ab + delta_x = other.mean_x.to(self._device) - self.mean_x + delta_y = other.mean_y.to(self._device) - self.mean_y + + self.mean_x = self.mean_x + delta_x * n_b / n_ab + self.mean_y = self.mean_y + delta_y * n_b / n_ab + self.sum_sq_dev_x = self.sum_sq_dev_x + other.sum_sq_dev_x.to(self._device) + delta_x * delta_x * cross + self.sum_sq_dev_y = self.sum_sq_dev_y + other.sum_sq_dev_y.to(self._device) + delta_y * delta_y * cross + self.sum_product_of_devs = ( + self.sum_product_of_devs + other.sum_product_of_devs.to(self._device) + delta_x * delta_y * cross + ) + self.n_samples = n_ab + + @property + def variance_x(self) -> torch.Tensor: + """Population variance of ``x`` (divisor ``n``).""" + if self.n_samples == 0: + return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return torch.clamp(self.sum_sq_dev_x / self.n_samples, min=0.0) + + @property + def variance_y(self) -> torch.Tensor: + """Population variance of ``y`` (divisor ``n``).""" + if self.n_samples == 0: + return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return torch.clamp(self.sum_sq_dev_y / self.n_samples, min=0.0) + + @property + def covariance(self) -> torch.Tensor: + """Population covariance of ``(x, y)`` (divisor ``n``).""" + if self.n_samples == 0: + return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return self.sum_product_of_devs / self.n_samples + + def correlation(self, eps: float = 1e-8) -> torch.Tensor: + """Pearson correlation coefficient with a small clamp for safety. + + Args: + eps: floor on the denominator to avoid division by zero when + one of the variables is constant. + """ + denom = torch.clamp(self.variance_x.sqrt() * self.variance_y.sqrt(), min=eps) + return self.covariance / denom diff --git a/tests/ignite/metrics/test_running_stats.py b/tests/ignite/metrics/test_running_stats.py new file mode 100644 index 000000000000..a147a5a9badd --- /dev/null +++ b/tests/ignite/metrics/test_running_stats.py @@ -0,0 +1,277 @@ +import numpy as np +import pytest +import torch + +from ignite.metrics._running_stats import WelfordCovariance, WelfordVariance + + +# --------------------------------------------------------------------------- +# WelfordVariance +# --------------------------------------------------------------------------- + + +class TestWelfordVariance: + def test_empty_accumulator(self): + ws = WelfordVariance() + assert ws.n_samples == 0 + assert ws.variance.item() == 0.0 + assert ws.std.item() == 0.0 + + def test_update_then_compute_matches_numpy(self): + rng = np.random.default_rng(0) + # Use float64 throughout so we compare apples to apples; Welford + # upcasts internally, so feeding float32 inputs and comparing against + # float32 numpy stats would understate the helper's precision. + data = rng.standard_normal(1000) + + ws = WelfordVariance() + ws.update(torch.from_numpy(data)) + + assert ws.n_samples == 1000 + assert ws.mean.item() == pytest.approx(float(data.mean()), abs=1e-12) + assert ws.variance.item() == pytest.approx(float(data.var()), rel=1e-12) + + def test_multi_batch_matches_single_batch(self): + rng = np.random.default_rng(1) + data = rng.standard_normal(1000).astype(np.float32) + data_t = torch.from_numpy(data) + + single = WelfordVariance() + single.update(data_t) + + multi = WelfordVariance() + for start in range(0, 1000, 37): + multi.update(data_t[start : start + 37]) + + assert multi.n_samples == single.n_samples + assert multi.mean.item() == pytest.approx(single.mean.item(), abs=1e-12) + assert multi.variance.item() == pytest.approx(single.variance.item(), rel=1e-12) + + def test_merge_matches_concatenated_update(self): + rng = np.random.default_rng(2) + a = torch.from_numpy(rng.standard_normal(400).astype(np.float64)) + b = torch.from_numpy(rng.standard_normal(600).astype(np.float64)) + + merged = WelfordVariance() + merged.update(a) + right = WelfordVariance() + right.update(b) + merged.merge(right) + + baseline = WelfordVariance() + baseline.update(torch.cat([a, b])) + + assert merged.n_samples == baseline.n_samples + assert merged.mean.item() == pytest.approx(baseline.mean.item(), abs=1e-12) + assert merged.variance.item() == pytest.approx(baseline.variance.item(), rel=1e-12) + + def test_merge_with_empty_accumulators(self): + rng = np.random.default_rng(3) + data = torch.from_numpy(rng.standard_normal(100)) + + # Empty merged into populated -> unchanged. + a = WelfordVariance() + a.update(data) + before_mean = a.mean.item() + a.merge(WelfordVariance()) + assert a.mean.item() == pytest.approx(before_mean, abs=1e-12) + + # Populated merged into empty -> takes the other's state. + b = WelfordVariance() + b.merge(a) + assert b.n_samples == a.n_samples + assert b.mean.item() == pytest.approx(a.mean.item(), abs=1e-12) + + def test_numerical_stability_large_mean_float32(self): + # The whole point of this helper: naive Σx^2 - (Σx)^2/n computed in + # float32 catastrophically cancels at mean=1e6, returning ~0 variance + # (or even negative). Welford in float64 stays exact. + rng = np.random.default_rng(4) + true_std = 1.0 + data = rng.standard_normal(10_000).astype(np.float32) * true_std + 1e6 + data_t = torch.from_numpy(data) + + # Naive formula in float32 collapses. + sum_x_f32 = data_t.sum() + sum_x2_f32 = (data_t * data_t).sum() + naive_var_f32 = (sum_x2_f32 - sum_x_f32 * sum_x_f32 / len(data_t)) / len(data_t) + # Use float64 ground truth so the assertion isn't measuring our own bug. + true_var = float(np.var(data.astype(np.float64))) + + # Welford in float64 should recover the true variance. + ws = WelfordVariance() + ws.update(data_t) + assert ws.variance.item() == pytest.approx(true_var, rel=1e-6) + + # And the naive float32 formula must demonstrably fail on the same + # data so the test documents what we're protecting against. + assert abs(float(naive_var_f32) - true_var) > 0.1, ( + "naive float32 formula did NOT cancel; test is no longer exercising the failure mode it claims to." + ) + + def test_single_sample(self): + ws = WelfordVariance() + ws.update(torch.tensor([42.0])) + assert ws.n_samples == 1 + assert ws.mean.item() == 42.0 + assert ws.variance.item() == 0.0 + + def test_empty_batch_is_noop(self): + ws = WelfordVariance() + ws.update(torch.tensor([1.0, 2.0, 3.0])) + before = (ws.n_samples, ws.mean.item(), ws.variance.item()) + ws.update(torch.tensor([])) + after = (ws.n_samples, ws.mean.item(), ws.variance.item()) + assert before == after + + def test_reset(self): + ws = WelfordVariance() + ws.update(torch.randn(100)) + ws.reset() + assert ws.n_samples == 0 + assert ws.mean.item() == 0.0 + assert ws.sum_sq_dev_from_mean.item() == 0.0 + + def test_input_dtype_upcast_to_float64(self): + ws = WelfordVariance() + ws.update(torch.tensor([1, 2, 3, 4], dtype=torch.int32)) + assert ws.mean.dtype == torch.float64 + assert ws.mean.item() == pytest.approx(2.5) + assert ws.variance.item() == pytest.approx(1.25) + + +# --------------------------------------------------------------------------- +# WelfordCovariance +# --------------------------------------------------------------------------- + + +class TestWelfordCovariance: + def test_empty_accumulator(self): + wc = WelfordCovariance() + assert wc.n_samples == 0 + assert wc.variance_x.item() == 0.0 + assert wc.variance_y.item() == 0.0 + assert wc.covariance.item() == 0.0 + + def test_update_matches_numpy_corrcoef(self): + rng = np.random.default_rng(5) + n = 1000 + x = rng.standard_normal(n) + y = 0.7 * x + rng.standard_normal(n) * 0.3 + + wc = WelfordCovariance() + wc.update(torch.from_numpy(x), torch.from_numpy(y)) + + np_var_x = float(np.var(x)) + np_var_y = float(np.var(y)) + np_cov = float(np.cov(x, y, bias=True)[0, 1]) + np_r = float(np.corrcoef(x, y)[0, 1]) + + assert wc.variance_x.item() == pytest.approx(np_var_x, rel=1e-12) + assert wc.variance_y.item() == pytest.approx(np_var_y, rel=1e-12) + assert wc.covariance.item() == pytest.approx(np_cov, rel=1e-12) + assert wc.correlation().item() == pytest.approx(np_r, rel=1e-10) + + def test_multi_batch_matches_single_batch(self): + rng = np.random.default_rng(6) + x = torch.from_numpy(rng.standard_normal(900)) + y = torch.from_numpy(rng.standard_normal(900)) + + single = WelfordCovariance() + single.update(x, y) + + multi = WelfordCovariance() + for start in range(0, 900, 31): + multi.update(x[start : start + 31], y[start : start + 31]) + + assert multi.mean_x.item() == pytest.approx(single.mean_x.item(), abs=1e-12) + assert multi.mean_y.item() == pytest.approx(single.mean_y.item(), abs=1e-12) + assert multi.covariance.item() == pytest.approx(single.covariance.item(), rel=1e-12) + assert multi.correlation().item() == pytest.approx(single.correlation().item(), rel=1e-12) + + def test_merge_matches_concatenated_update(self): + rng = np.random.default_rng(7) + x1 = torch.from_numpy(rng.standard_normal(300)) + y1 = torch.from_numpy(rng.standard_normal(300)) + x2 = torch.from_numpy(rng.standard_normal(500)) + y2 = torch.from_numpy(rng.standard_normal(500)) + + merged = WelfordCovariance() + merged.update(x1, y1) + right = WelfordCovariance() + right.update(x2, y2) + merged.merge(right) + + baseline = WelfordCovariance() + baseline.update(torch.cat([x1, x2]), torch.cat([y1, y2])) + + assert merged.covariance.item() == pytest.approx(baseline.covariance.item(), rel=1e-12) + assert merged.correlation().item() == pytest.approx(baseline.correlation().item(), rel=1e-12) + + def test_numerical_stability_large_mean(self): + # The Pearson-correlation regression case from issue #3662: mean=1e6, + # std=1 in float32 makes the naive E[X^2] - E[X]^2 formula return + # garbage. Welford with float64 internals recovers the true r. + rng = np.random.default_rng(8) + n = 10_000 + x = rng.standard_normal(n).astype(np.float32) + 1e6 + y = (0.99 * x + rng.standard_normal(n).astype(np.float32) * 0.1).astype(np.float32) + + true_r = float(np.corrcoef(x.astype(np.float64), y.astype(np.float64))[0, 1]) + # Sanity: the constructed series really is highly correlated. + assert true_r > 0.99 + + wc = WelfordCovariance() + wc.update(torch.from_numpy(x), torch.from_numpy(y)) + assert wc.correlation().item() == pytest.approx(true_r, rel=1e-4) + + def test_shape_mismatch_raises(self): + wc = WelfordCovariance() + with pytest.raises(ValueError, match="same shape"): + wc.update(torch.zeros(5), torch.zeros(6)) + + def test_empty_batch_is_noop(self): + wc = WelfordCovariance() + wc.update(torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])) + before = (wc.n_samples, wc.covariance.item()) + wc.update(torch.tensor([]), torch.tensor([])) + after = (wc.n_samples, wc.covariance.item()) + assert before == after + + def test_constant_variable_correlation_safe(self): + # When one series is constant the denominator of Pearson r is zero; + # the eps clamp keeps us from returning NaN / inf. + wc = WelfordCovariance() + wc.update(torch.tensor([1.0, 2.0, 3.0, 4.0]), torch.tensor([5.0, 5.0, 5.0, 5.0])) + r = wc.correlation().item() + assert r == 0.0 + assert not (r != r) # not NaN + + def test_reset(self): + wc = WelfordCovariance() + wc.update(torch.randn(50), torch.randn(50)) + wc.reset() + assert wc.n_samples == 0 + assert wc.covariance.item() == 0.0 + + +# --------------------------------------------------------------------------- +# Cross-class sanity: variance_x of WelfordCovariance == variance of +# WelfordVariance fed the same x. Catches drift between the two +# implementations. +# --------------------------------------------------------------------------- + + +def test_variance_x_matches_welford_variance(): + rng = np.random.default_rng(9) + x = torch.from_numpy(rng.standard_normal(1000)) + y = torch.from_numpy(rng.standard_normal(1000)) + + wv = WelfordVariance() + wv.update(x) + + wc = WelfordCovariance() + wc.update(x, y) + + assert wc.variance_x.item() == pytest.approx(wv.variance.item(), rel=1e-12) + assert wc.mean_x.item() == pytest.approx(wv.mean.item(), abs=1e-12) From 1b8244244b162223ccaf881c8a4f37c0e8cd85d1 Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Mon, 11 May 2026 20:11:07 +0300 Subject: [PATCH 2/9] review: drop device/dtype handling, switch to dataclass, no flatten Addressing @vfdev-5's inline review on #3750: - Drop the device and dtype constructor args. The helper now leaves placement and precision to the caller; state takes the dtype and device of the first batch passed to update(). PearsonCorrelation and R2Score already do their own float64 upcast before handing inputs to the helper, so this is a no-op for the planned consumers. - Switch both classes to @dataclass with field(default_factory=...) for the tensor fields. Drops the manual __init__ / reset() plumbing; "reset" is now reconstruction (m.welford = WelfordVariance()), which is the natural fit for how the consumer Metric.reset() methods already work. - Drop the explicit .flatten() on update inputs. batch.mean() and batch.numel() both reduce over the full tensor regardless of shape, so behavior for the current scalar-reduction consumers is unchanged, and the code reads more naturally for any shape. Tests adjusted accordingly: - test_reset replaced by test_fresh_instance_has_zero_state, which documents the dataclass default-factory behavior. - test_input_dtype_upcast_to_float64 replaced by test_state_dtype_follows_first_batch, which verifies dtype is preserved (the design change). - Stability tests upcast inputs caller-side before handing to the helper, matching how the metric classes will use it. - test_multi_batch_matches_single_batch switched to float64 inputs so it exercises the algorithm rather than float32 noise. All 20 tests still pass, ruff format / check clean. The question about axis-aware reduction is deferred to the review thread; I'll follow it once @vfdev-5 confirms whether it lands here or as a follow-up. --- ignite/metrics/_running_stats.py | 175 +++++++++------------ tests/ignite/metrics/test_running_stats.py | 49 +++--- 2 files changed, 106 insertions(+), 118 deletions(-) diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py index 206353b9cabd..7816080d32eb 100644 --- a/ignite/metrics/_running_stats.py +++ b/ignite/metrics/_running_stats.py @@ -7,11 +7,16 @@ :class:`~ignite.metrics.regression.R2Score`; new metrics with the same need should consume these helpers rather than rolling their own. -Both classes keep internal state in ``float64`` regardless of the input -dtype, follow Welford's online algorithm for incremental updates, and -fold accumulators together with the Chan / Welford parallel formula -(used both for batch-wise updates and for cross-rank merges in -distributed settings). +Both classes are tensor-type-agnostic dataclasses: callers supply +tensors in whatever dtype and device they want, and the helper +preserves both. For numerical stability under large means, callers +should pre-cast inputs to ``float64`` (the consumer metric classes +already do this in their own ``update`` methods). + +Updates follow Welford's online algorithm. Two accumulators can be +combined into one via :meth:`merge`, which implements the Chan / +Welford parallel formula and is the basis for cross-rank distributed +reductions. References: Welford, B. P. (1962). Note on a method for calculating corrected @@ -21,11 +26,16 @@ https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm """ -from typing import Union +from dataclasses import dataclass, field import torch +def _zero() -> torch.Tensor: + return torch.tensor(0.0) + + +@dataclass class WelfordVariance: """Numerically stable running mean and variance via Welford's algorithm. @@ -35,52 +45,38 @@ class WelfordVariance: which uses the Chan / Welford parallel formula and is the basis for distributed reductions. - State is kept in ``float64`` regardless of input dtype so that the - classic ``E[X^2] - E[X]^2`` cancellation does not bite for inputs - with large means. + No dtype or device handling is performed inside the class: the + state takes the dtype and device of the first batch passed to + :meth:`update`. For numerical stability under large means, callers + should hand in ``float64`` tensors. - Args: - device: device on which to keep the running-state tensors. - Default: ``"cpu"``. + Example:: - Example: - - .. code-block:: python - - ws = WelfordVariance() - for batch in stream: - ws.update(batch) - print(ws.mean.item(), ws.variance.item(), ws.std.item()) + ws = WelfordVariance() + for batch in stream: + ws.update(batch.to(torch.float64)) + print(ws.mean.item(), ws.variance.item(), ws.std.item()) """ - n_samples: int - mean: torch.Tensor - sum_sq_dev_from_mean: torch.Tensor - - def __init__(self, device: Union[str, torch.device] = "cpu") -> None: - self._device = torch.device(device) - self.reset() - - def reset(self) -> None: - """Drop all accumulated state.""" - self.n_samples = 0 - self.mean = torch.tensor(0.0, dtype=torch.float64, device=self._device) - self.sum_sq_dev_from_mean = torch.tensor(0.0, dtype=torch.float64, device=self._device) + n_samples: int = 0 + mean: torch.Tensor = field(default_factory=_zero) + sum_sq_dev_from_mean: torch.Tensor = field(default_factory=_zero) @torch.no_grad() def update(self, batch: torch.Tensor) -> None: """Fold a batch of samples into the running state. - Empty batches are silently ignored. Inputs of any dtype are - upcast to ``float64`` for the internal computation. + Empty batches are silently ignored. ``batch.mean()`` and + ``batch.numel()`` perform a full reduction over the input, so + any shape is accepted and treated as ``numel`` scalar samples. """ if batch.numel() == 0: return - batch64 = batch.detach().to(dtype=torch.float64).flatten() - n_b = batch64.shape[0] + batch = batch.detach() + n_b = batch.numel() - mean_b = batch64.mean().to(self._device) - m2_b = (batch64 - batch64.mean()).square().sum().to(self._device) + mean_b = batch.mean() + m2_b = (batch - mean_b).square().sum() if self.n_samples == 0: self.mean = mean_b @@ -96,26 +92,21 @@ def update(self, batch: torch.Tensor) -> None: self.n_samples = n_ab def merge(self, other: "WelfordVariance") -> None: - """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. - - Used both for cross-rank reduction in distributed settings and - for any other case where two accumulators need to be combined - into one without re-iterating the raw data. - """ + """Combine ``other`` into ``self`` using the Chan / Welford parallel formula.""" if other.n_samples == 0: return if self.n_samples == 0: self.n_samples = other.n_samples - self.mean = other.mean.detach().clone().to(self._device) - self.sum_sq_dev_from_mean = other.sum_sq_dev_from_mean.detach().clone().to(self._device) + self.mean = other.mean.detach().clone() + self.sum_sq_dev_from_mean = other.sum_sq_dev_from_mean.detach().clone() return n_a = self.n_samples n_b = other.n_samples n_ab = n_a + n_b - delta = other.mean.to(self._device) - self.mean + delta = other.mean - self.mean self.mean = self.mean + delta * n_b / n_ab self.sum_sq_dev_from_mean = ( - self.sum_sq_dev_from_mean + other.sum_sq_dev_from_mean.to(self._device) + delta * delta * n_a * n_b / n_ab + self.sum_sq_dev_from_mean + other.sum_sq_dev_from_mean + delta * delta * n_a * n_b / n_ab ) self.n_samples = n_ab @@ -123,7 +114,7 @@ def merge(self, other: "WelfordVariance") -> None: def variance(self) -> torch.Tensor: """Population variance (divisor ``n``). Returns ``0.0`` when empty.""" if self.n_samples == 0: - return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return torch.tensor(0.0) return torch.clamp(self.sum_sq_dev_from_mean / self.n_samples, min=0.0) @property @@ -132,6 +123,7 @@ def std(self) -> torch.Tensor: return self.variance.sqrt() +@dataclass class WelfordCovariance: """Numerically stable running covariance for a pair of variables (x, y). @@ -140,37 +132,24 @@ class WelfordCovariance: online update + Chan / Welford parallel merge as :class:`WelfordVariance`. - Args: - device: device on which to keep the running-state tensors. - Default: ``"cpu"``. + Like :class:`WelfordVariance`, the class is dtype and device + agnostic: state takes the dtype and device of the first batch + passed to :meth:`update`. """ - n_samples: int - mean_x: torch.Tensor - mean_y: torch.Tensor - sum_sq_dev_x: torch.Tensor - sum_sq_dev_y: torch.Tensor - sum_product_of_devs: torch.Tensor - - def __init__(self, device: Union[str, torch.device] = "cpu") -> None: - self._device = torch.device(device) - self.reset() - - def reset(self) -> None: - """Drop all accumulated state.""" - self.n_samples = 0 - self.mean_x = torch.tensor(0.0, dtype=torch.float64, device=self._device) - self.mean_y = torch.tensor(0.0, dtype=torch.float64, device=self._device) - self.sum_sq_dev_x = torch.tensor(0.0, dtype=torch.float64, device=self._device) - self.sum_sq_dev_y = torch.tensor(0.0, dtype=torch.float64, device=self._device) - self.sum_product_of_devs = torch.tensor(0.0, dtype=torch.float64, device=self._device) + n_samples: int = 0 + mean_x: torch.Tensor = field(default_factory=_zero) + mean_y: torch.Tensor = field(default_factory=_zero) + sum_sq_dev_x: torch.Tensor = field(default_factory=_zero) + sum_sq_dev_y: torch.Tensor = field(default_factory=_zero) + sum_product_of_devs: torch.Tensor = field(default_factory=_zero) @torch.no_grad() def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: """Fold a paired batch ``(x_i, y_i)`` into the running state. - ``batch_x`` and ``batch_y`` must have the same shape; both are - flattened internally and upcast to ``float64``. + ``batch_x`` and ``batch_y`` must have the same shape; the full + tensor is reduced as ``numel`` scalar samples. """ if batch_x.shape != batch_y.shape: raise ValueError( @@ -179,17 +158,17 @@ def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: if batch_x.numel() == 0: return - x64 = batch_x.detach().to(dtype=torch.float64).flatten() - y64 = batch_y.detach().to(dtype=torch.float64).flatten() - n_b = x64.shape[0] + x = batch_x.detach() + y = batch_y.detach() + n_b = x.numel() - mean_x_b = x64.mean().to(self._device) - mean_y_b = y64.mean().to(self._device) - dx_b = x64 - mean_x_b - dy_b = y64 - mean_y_b - m2_x_b = dx_b.square().sum().to(self._device) - m2_y_b = dy_b.square().sum().to(self._device) - cxy_b = (dx_b * dy_b).sum().to(self._device) + mean_x_b = x.mean() + mean_y_b = y.mean() + dx_b = x - mean_x_b + dy_b = y - mean_y_b + m2_x_b = dx_b.square().sum() + m2_y_b = dy_b.square().sum() + cxy_b = (dx_b * dy_b).sum() if self.n_samples == 0: self.mean_x = mean_x_b @@ -219,47 +198,45 @@ def merge(self, other: "WelfordCovariance") -> None: return if self.n_samples == 0: self.n_samples = other.n_samples - self.mean_x = other.mean_x.detach().clone().to(self._device) - self.mean_y = other.mean_y.detach().clone().to(self._device) - self.sum_sq_dev_x = other.sum_sq_dev_x.detach().clone().to(self._device) - self.sum_sq_dev_y = other.sum_sq_dev_y.detach().clone().to(self._device) - self.sum_product_of_devs = other.sum_product_of_devs.detach().clone().to(self._device) + self.mean_x = other.mean_x.detach().clone() + self.mean_y = other.mean_y.detach().clone() + self.sum_sq_dev_x = other.sum_sq_dev_x.detach().clone() + self.sum_sq_dev_y = other.sum_sq_dev_y.detach().clone() + self.sum_product_of_devs = other.sum_product_of_devs.detach().clone() return n_a = self.n_samples n_b = other.n_samples n_ab = n_a + n_b cross = n_a * n_b / n_ab - delta_x = other.mean_x.to(self._device) - self.mean_x - delta_y = other.mean_y.to(self._device) - self.mean_y + delta_x = other.mean_x - self.mean_x + delta_y = other.mean_y - self.mean_y self.mean_x = self.mean_x + delta_x * n_b / n_ab self.mean_y = self.mean_y + delta_y * n_b / n_ab - self.sum_sq_dev_x = self.sum_sq_dev_x + other.sum_sq_dev_x.to(self._device) + delta_x * delta_x * cross - self.sum_sq_dev_y = self.sum_sq_dev_y + other.sum_sq_dev_y.to(self._device) + delta_y * delta_y * cross - self.sum_product_of_devs = ( - self.sum_product_of_devs + other.sum_product_of_devs.to(self._device) + delta_x * delta_y * cross - ) + self.sum_sq_dev_x = self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * cross + self.sum_sq_dev_y = self.sum_sq_dev_y + other.sum_sq_dev_y + delta_y * delta_y * cross + self.sum_product_of_devs = self.sum_product_of_devs + other.sum_product_of_devs + delta_x * delta_y * cross self.n_samples = n_ab @property def variance_x(self) -> torch.Tensor: """Population variance of ``x`` (divisor ``n``).""" if self.n_samples == 0: - return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return torch.tensor(0.0) return torch.clamp(self.sum_sq_dev_x / self.n_samples, min=0.0) @property def variance_y(self) -> torch.Tensor: """Population variance of ``y`` (divisor ``n``).""" if self.n_samples == 0: - return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return torch.tensor(0.0) return torch.clamp(self.sum_sq_dev_y / self.n_samples, min=0.0) @property def covariance(self) -> torch.Tensor: """Population covariance of ``(x, y)`` (divisor ``n``).""" if self.n_samples == 0: - return torch.tensor(0.0, dtype=torch.float64, device=self._device) + return torch.tensor(0.0) return self.sum_product_of_devs / self.n_samples def correlation(self, eps: float = 1e-8) -> torch.Tensor: diff --git a/tests/ignite/metrics/test_running_stats.py b/tests/ignite/metrics/test_running_stats.py index a147a5a9badd..f2209873afc5 100644 --- a/tests/ignite/metrics/test_running_stats.py +++ b/tests/ignite/metrics/test_running_stats.py @@ -32,8 +32,10 @@ def test_update_then_compute_matches_numpy(self): assert ws.variance.item() == pytest.approx(float(data.var()), rel=1e-12) def test_multi_batch_matches_single_batch(self): + # Use float64 so the test exercises the algorithm rather than float32 + # accumulation noise. rng = np.random.default_rng(1) - data = rng.standard_normal(1000).astype(np.float32) + data = rng.standard_normal(1000) data_t = torch.from_numpy(data) single = WelfordVariance() @@ -85,7 +87,7 @@ def test_merge_with_empty_accumulators(self): def test_numerical_stability_large_mean_float32(self): # The whole point of this helper: naive Σx^2 - (Σx)^2/n computed in # float32 catastrophically cancels at mean=1e6, returning ~0 variance - # (or even negative). Welford in float64 stays exact. + # (or even negative). Welford fed float64 inputs stays exact. rng = np.random.default_rng(4) true_std = 1.0 data = rng.standard_normal(10_000).astype(np.float32) * true_std + 1e6 @@ -98,9 +100,10 @@ def test_numerical_stability_large_mean_float32(self): # Use float64 ground truth so the assertion isn't measuring our own bug. true_var = float(np.var(data.astype(np.float64))) - # Welford in float64 should recover the true variance. + # The helper is dtype-agnostic; the caller is responsible for the + # float64 upcast. Verify the upcast path recovers the true variance. ws = WelfordVariance() - ws.update(data_t) + ws.update(data_t.to(torch.float64)) assert ws.variance.item() == pytest.approx(true_var, rel=1e-6) # And the naive float32 formula must demonstrably fail on the same @@ -124,20 +127,28 @@ def test_empty_batch_is_noop(self): after = (ws.n_samples, ws.mean.item(), ws.variance.item()) assert before == after - def test_reset(self): + def test_fresh_instance_has_zero_state(self): + # The dataclass starts empty; "reset" is just reconstruction. Verifies + # that the default factories produce an empty accumulator. ws = WelfordVariance() - ws.update(torch.randn(100)) - ws.reset() assert ws.n_samples == 0 assert ws.mean.item() == 0.0 assert ws.sum_sq_dev_from_mean.item() == 0.0 - def test_input_dtype_upcast_to_float64(self): - ws = WelfordVariance() - ws.update(torch.tensor([1, 2, 3, 4], dtype=torch.int32)) - assert ws.mean.dtype == torch.float64 - assert ws.mean.item() == pytest.approx(2.5) - assert ws.variance.item() == pytest.approx(1.25) + def test_state_dtype_follows_first_batch(self): + # The helper does not handle dtype itself; it takes whatever dtype + # the first batch arrives in and preserves it. The caller chooses. + ws_f32 = WelfordVariance() + ws_f32.update(torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)) + assert ws_f32.mean.dtype == torch.float32 + assert ws_f32.mean.item() == pytest.approx(2.5) + assert ws_f32.variance.item() == pytest.approx(1.25) + + ws_f64 = WelfordVariance() + ws_f64.update(torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64)) + assert ws_f64.mean.dtype == torch.float64 + assert ws_f64.mean.item() == pytest.approx(2.5) + assert ws_f64.variance.item() == pytest.approx(1.25) # --------------------------------------------------------------------------- @@ -210,8 +221,8 @@ def test_merge_matches_concatenated_update(self): def test_numerical_stability_large_mean(self): # The Pearson-correlation regression case from issue #3662: mean=1e6, - # std=1 in float32 makes the naive E[X^2] - E[X]^2 formula return - # garbage. Welford with float64 internals recovers the true r. + # std=1 makes the naive E[X^2] - E[X]^2 formula return garbage in + # float32. Welford fed float64 inputs recovers the true r. rng = np.random.default_rng(8) n = 10_000 x = rng.standard_normal(n).astype(np.float32) + 1e6 @@ -221,8 +232,9 @@ def test_numerical_stability_large_mean(self): # Sanity: the constructed series really is highly correlated. assert true_r > 0.99 + # Caller-side upcast to float64 (the helper preserves whatever it gets). wc = WelfordCovariance() - wc.update(torch.from_numpy(x), torch.from_numpy(y)) + wc.update(torch.from_numpy(x).to(torch.float64), torch.from_numpy(y).to(torch.float64)) assert wc.correlation().item() == pytest.approx(true_r, rel=1e-4) def test_shape_mismatch_raises(self): @@ -247,10 +259,9 @@ def test_constant_variable_correlation_safe(self): assert r == 0.0 assert not (r != r) # not NaN - def test_reset(self): + def test_fresh_instance_has_zero_state(self): + # "Reset" is just reconstruction with this dataclass. wc = WelfordCovariance() - wc.update(torch.randn(50), torch.randn(50)) - wc.reset() assert wc.n_samples == 0 assert wc.covariance.item() == 0.0 From e153f57d72101d08e003cb9790368dd2696b38fa Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Wed, 13 May 2026 20:21:27 +0300 Subject: [PATCH 3/9] review: collapse update into merge, inline _zero, expand merge docstring Three readability changes responding to @aaishwarymishra's inline review: 1. _zero() helper removed; each tensor field uses field(default_factory=lambda: torch.tensor(0.0)) directly. 2. update() is now the degenerate case of merge() where "other" is a freshly-built single-batch accumulator. The Chan / Welford parallel formula lives in exactly one place. Same refactor applied to WelfordCovariance.update. 3. merge() docstring grew a paragraph explaining the distributed- reduction motivation -- without an explicit merge, cross-rank reduction has to re-iterate the raw data, which defeats the point of an online algorithm. Behavior is bit-equivalent: the only delta is an extra detach/clone on the first-batch path (via merge's first-time-absorb branch), which is a no-op for correctness. --- ignite/metrics/_running_stats.py | 175 +++++++++++++++++++------------ 1 file changed, 107 insertions(+), 68 deletions(-) diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py index 7816080d32eb..88168f2545bb 100644 --- a/ignite/metrics/_running_stats.py +++ b/ignite/metrics/_running_stats.py @@ -13,10 +13,13 @@ should pre-cast inputs to ``float64`` (the consumer metric classes already do this in their own ``update`` methods). -Updates follow Welford's online algorithm. Two accumulators can be -combined into one via :meth:`merge`, which implements the Chan / -Welford parallel formula and is the basis for cross-rank distributed -reductions. +Two operations matter: :meth:`merge`, which combines two accumulators +into one via the Chan / Welford parallel formula (also the basis for +cross-rank distributed reductions), and :meth:`update`, which folds a +new batch into the running state. ``update`` is the degenerate case of +``merge`` where ``other`` is a freshly-built single-batch accumulator, +and the implementation reflects that: ``update`` builds the batch +accumulator and delegates to ``merge``. There is one formula, not two. References: Welford, B. P. (1962). Note on a method for calculating corrected @@ -31,10 +34,6 @@ import torch -def _zero() -> torch.Tensor: - return torch.tensor(0.0) - - @dataclass class WelfordVariance: """Numerically stable running mean and variance via Welford's algorithm. @@ -58,14 +57,24 @@ class WelfordVariance: print(ws.mean.item(), ws.variance.item(), ws.std.item()) """ + # n_samples: count of samples folded in. + # mean: running sample mean (Welford state). + # sum_sq_dev_from_mean: Σ (x_i − mean)^2, the second central moment + # numerator, conventionally called "M2" in the Welford literature. n_samples: int = 0 - mean: torch.Tensor = field(default_factory=_zero) - sum_sq_dev_from_mean: torch.Tensor = field(default_factory=_zero) + mean: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) + sum_sq_dev_from_mean: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) @torch.no_grad() def update(self, batch: torch.Tensor) -> None: """Fold a batch of samples into the running state. + Implementation: build a single-batch accumulator from ``batch`` + and merge it into ``self``. ``update`` is the degenerate case + of :meth:`merge` where the right-hand side has just been built + from one batch; sharing the parallel formula keeps the two + paths in lock-step. + Empty batches are silently ignored. ``batch.mean()`` and ``batch.numel()`` perform a full reduction over the input, so any shape is accepted and treated as ``numel`` scalar samples. @@ -73,38 +82,60 @@ def update(self, batch: torch.Tensor) -> None: if batch.numel() == 0: return batch = batch.detach() - n_b = batch.numel() - - mean_b = batch.mean() - m2_b = (batch - mean_b).square().sum() - - if self.n_samples == 0: - self.mean = mean_b - self.sum_sq_dev_from_mean = m2_b - self.n_samples = n_b - return - - n_a = self.n_samples - n_ab = n_a + n_b - delta = mean_b - self.mean - self.mean = self.mean + delta * n_b / n_ab - self.sum_sq_dev_from_mean = self.sum_sq_dev_from_mean + m2_b + delta * delta * n_a * n_b / n_ab - self.n_samples = n_ab + batch_mean = batch.mean() + batch_acc = WelfordVariance( + n_samples=batch.numel(), + mean=batch_mean, + sum_sq_dev_from_mean=(batch - batch_mean).square().sum(), + ) + self.merge(batch_acc) def merge(self, other: "WelfordVariance") -> None: - """Combine ``other`` into ``self`` using the Chan / Welford parallel formula.""" + """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. + + Used in two places: by :meth:`update` (where ``other`` is a + freshly-built single-batch accumulator), and by callers that + need to combine independently-accumulated state from elsewhere. + The motivating second case is distributed training: each rank + accumulates its own ``WelfordVariance`` over its local samples, + then at eval time the ranks merge their accumulators rank-by-rank + to produce the population statistic. Without :meth:`merge` that + cross-rank reduction would have to re-iterate the raw data, + which defeats the whole point of an online algorithm. + + Given two accumulators ``A`` and ``B`` with sample counts + ``n_a, n_b`` and second-central-moment sums ``M2_a, M2_b``, the + combined ``M2`` over the concatenated stream is:: + + M2 = M2_a + M2_b + (mean_b - mean_a)^2 * n_a * n_b / (n_a + n_b) + + The third term is the *correction*: simply adding ``M2_a + M2_b`` + would under-count the variance whenever the two batches have + different sample means, because each ``M2`` is measured relative + to its own local mean. The correction folds in the spread of + the two local means about the combined mean. + """ if other.n_samples == 0: return if self.n_samples == 0: + # First-time absorb. Copy state so callers cannot mutate + # ``other`` and silently affect ``self``. self.n_samples = other.n_samples self.mean = other.mean.detach().clone() self.sum_sq_dev_from_mean = other.sum_sq_dev_from_mean.detach().clone() return + n_a = self.n_samples n_b = other.n_samples n_ab = n_a + n_b delta = other.mean - self.mean + + # Standard Welford incremental-mean update, weighted by the + # fraction of the combined sample size that ``other`` contributes. self.mean = self.mean + delta * n_b / n_ab + + # Parallel-formula combined M2. The (delta * delta * ...) term + # is the correction described in the docstring above. self.sum_sq_dev_from_mean = ( self.sum_sq_dev_from_mean + other.sum_sq_dev_from_mean + delta * delta * n_a * n_b / n_ab ) @@ -130,24 +161,32 @@ class WelfordCovariance: Exposes :attr:`variance_x`, :attr:`variance_y`, :attr:`covariance`, and :meth:`correlation` (Pearson) through the same Welford-style online update + Chan / Welford parallel merge as - :class:`WelfordVariance`. + :class:`WelfordVariance`. The only extension over the univariate + case is the cross-product accumulator + :attr:`sum_product_of_devs` = Σ (x_i - mean_x) (y_i - mean_y). Like :class:`WelfordVariance`, the class is dtype and device agnostic: state takes the dtype and device of the first batch passed to :meth:`update`. """ + # Two univariate Welford accumulators worth of state, plus the + # cross-product term that turns them into a covariance. n_samples: int = 0 - mean_x: torch.Tensor = field(default_factory=_zero) - mean_y: torch.Tensor = field(default_factory=_zero) - sum_sq_dev_x: torch.Tensor = field(default_factory=_zero) - sum_sq_dev_y: torch.Tensor = field(default_factory=_zero) - sum_product_of_devs: torch.Tensor = field(default_factory=_zero) + mean_x: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) + mean_y: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) + sum_sq_dev_x: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) + sum_sq_dev_y: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) + sum_product_of_devs: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) @torch.no_grad() def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: """Fold a paired batch ``(x_i, y_i)`` into the running state. + Same trick as :meth:`WelfordVariance.update`: build a single-batch + accumulator from ``(batch_x, batch_y)`` and merge it. One formula, + applied twice; see :meth:`merge` for the math. + ``batch_x`` and ``batch_y`` must have the same shape; the full tensor is reduced as ``numel`` scalar samples. """ @@ -160,40 +199,29 @@ def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: x = batch_x.detach() y = batch_y.detach() - n_b = x.numel() - mean_x_b = x.mean() mean_y_b = y.mean() - dx_b = x - mean_x_b - dy_b = y - mean_y_b - m2_x_b = dx_b.square().sum() - m2_y_b = dy_b.square().sum() - cxy_b = (dx_b * dy_b).sum() - - if self.n_samples == 0: - self.mean_x = mean_x_b - self.mean_y = mean_y_b - self.sum_sq_dev_x = m2_x_b - self.sum_sq_dev_y = m2_y_b - self.sum_product_of_devs = cxy_b - self.n_samples = n_b - return - - n_a = self.n_samples - n_ab = n_a + n_b - cross = n_a * n_b / n_ab - delta_x = mean_x_b - self.mean_x - delta_y = mean_y_b - self.mean_y - - self.mean_x = self.mean_x + delta_x * n_b / n_ab - self.mean_y = self.mean_y + delta_y * n_b / n_ab - self.sum_sq_dev_x = self.sum_sq_dev_x + m2_x_b + delta_x * delta_x * cross - self.sum_sq_dev_y = self.sum_sq_dev_y + m2_y_b + delta_y * delta_y * cross - self.sum_product_of_devs = self.sum_product_of_devs + cxy_b + delta_x * delta_y * cross - self.n_samples = n_ab + dx = x - mean_x_b + dy = y - mean_y_b + batch_acc = WelfordCovariance( + n_samples=x.numel(), + mean_x=mean_x_b, + mean_y=mean_y_b, + sum_sq_dev_x=dx.square().sum(), + sum_sq_dev_y=dy.square().sum(), + sum_product_of_devs=(dx * dy).sum(), + ) + self.merge(batch_acc) def merge(self, other: "WelfordCovariance") -> None: - """Combine ``other`` into ``self`` using the Chan / Welford parallel formula.""" + """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. + + Same correction term as the univariate version, applied three + times: once for ``sum_sq_dev_x``, once for ``sum_sq_dev_y``, and + once for ``sum_product_of_devs`` (using ``delta_x * delta_y`` + instead of ``delta * delta``). See + :meth:`WelfordVariance.merge` for the derivation. + """ if other.n_samples == 0: return if self.n_samples == 0: @@ -204,18 +232,29 @@ def merge(self, other: "WelfordCovariance") -> None: self.sum_sq_dev_y = other.sum_sq_dev_y.detach().clone() self.sum_product_of_devs = other.sum_product_of_devs.detach().clone() return + n_a = self.n_samples n_b = other.n_samples n_ab = n_a + n_b - cross = n_a * n_b / n_ab + # The correction-term coefficient n_a * n_b / n_ab shows up in + # every parallel-formula line below; compute it once. + cross_coef = n_a * n_b / n_ab delta_x = other.mean_x - self.mean_x delta_y = other.mean_y - self.mean_y + # Incremental means, weighted by other's share of the new total. self.mean_x = self.mean_x + delta_x * n_b / n_ab self.mean_y = self.mean_y + delta_y * n_b / n_ab - self.sum_sq_dev_x = self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * cross - self.sum_sq_dev_y = self.sum_sq_dev_y + other.sum_sq_dev_y + delta_y * delta_y * cross - self.sum_product_of_devs = self.sum_product_of_devs + other.sum_product_of_devs + delta_x * delta_y * cross + + # Three parallel-formula combinations: variance of x, variance + # of y, and covariance of (x, y). Each is M_self + M_other plus + # a correction for the fact that the two batches had different + # local means. + self.sum_sq_dev_x = self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * cross_coef + self.sum_sq_dev_y = self.sum_sq_dev_y + other.sum_sq_dev_y + delta_y * delta_y * cross_coef + self.sum_product_of_devs = ( + self.sum_product_of_devs + other.sum_product_of_devs + delta_x * delta_y * cross_coef + ) self.n_samples = n_ab @property From 89d5b3619604311b21faf7f65f3c5ee4bbeb9740 Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Thu, 14 May 2026 19:41:12 +0300 Subject: [PATCH 4/9] =?UTF-8?q?review:=20address=20Copilot=20pass=20?= =?UTF-8?q?=E2=80=94=20docstring=20drift,=20no=5Fgrad=20on=20merge,=20inli?= =?UTF-8?q?ned=20cross-coef?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four points from the Copilot auto-review on #3750: 1. Module docstring claimed the helpers are *used by* PearsonCorrelation and R2Score, but those metrics are wired up in follow-up PRs and still use their own running-sum state at HEAD. Reworded to "intended consumers (in follow-up PRs of #3748)" so the doc matches reality at this commit. 2. Same docstring claimed internal state is kept in float64 regardless of input dtype, contradicting the actual implementation (dtype- and device-agnostic; caller supplies the float64 cast when stability matters). Vfdev-5's earlier review specifically asked for this contract — aligning the prose with the code so users don't get a false sense of safety on float32 inputs. 3. WelfordVariance.merge and WelfordCovariance.merge are now both wrapped in `@torch.no_grad()` (mirroring update). Without it, a caller that merges an accumulator whose tensors still require grad would build an autograd graph and leak memory across the lifetime of the metric. Belt to the existing detach/clone-on-first-time- absorb suspenders. 4. WelfordCovariance.merge precomputed `cross_coef = n_a * n_b / n_ab` as a Python float and reused it three times. Dropped the temporary and inlined `n_a * n_b / n_ab` directly into each parallel-formula line — mirrors WelfordVariance.merge's existing style and keeps the arithmetic on the same dtype/device as the tensor operands rather than promoting through Python scalar land. 20 / 20 tests still passing. --- ignite/metrics/_running_stats.py | 49 +++++++++++++++++++------------- 1 file changed, 30 insertions(+), 19 deletions(-) diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py index 88168f2545bb..f33dda10ac91 100644 --- a/ignite/metrics/_running_stats.py +++ b/ignite/metrics/_running_stats.py @@ -1,17 +1,20 @@ """Numerically stable running variance and covariance helpers. -Shared by metrics that need to accumulate variance / covariance from -streaming batches without falling into the catastrophic-cancellation -trap of the naive ``E[X^2] - E[X]^2`` formula. Used by -:class:`~ignite.metrics.regression.PearsonCorrelation` and -:class:`~ignite.metrics.regression.R2Score`; new metrics with the same -need should consume these helpers rather than rolling their own. - -Both classes are tensor-type-agnostic dataclasses: callers supply -tensors in whatever dtype and device they want, and the helper -preserves both. For numerical stability under large means, callers -should pre-cast inputs to ``float64`` (the consumer metric classes -already do this in their own ``update`` methods). +Shared internal primitives for metrics that need to accumulate +variance / covariance from streaming batches without falling into +the catastrophic-cancellation trap of the naive ``E[X^2] - E[X]^2`` +formula. Intended consumers (in follow-up PRs of #3748): the +``R2Score`` denominator and ``PearsonCorrelation`` cross-product — +new metrics with the same need should consume these helpers rather +than rolling their own. + +Both classes are dtype and device agnostic dataclasses: state takes +the dtype and device of the first batch passed to :meth:`update`, +and is preserved across subsequent updates and merges. The caller +is responsible for upcasting inputs to ``float64`` when numerical +stability under large means matters — the helper does not silently +promote, because doing so would surprise callers that already work +in the higher-precision dtype. Two operations matter: :meth:`merge`, which combines two accumulators into one via the Chan / Welford parallel formula (also the basis for @@ -90,6 +93,7 @@ def update(self, batch: torch.Tensor) -> None: ) self.merge(batch_acc) + @torch.no_grad() def merge(self, other: "WelfordVariance") -> None: """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. @@ -213,6 +217,7 @@ def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: ) self.merge(batch_acc) + @torch.no_grad() def merge(self, other: "WelfordCovariance") -> None: """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. @@ -236,9 +241,6 @@ def merge(self, other: "WelfordCovariance") -> None: n_a = self.n_samples n_b = other.n_samples n_ab = n_a + n_b - # The correction-term coefficient n_a * n_b / n_ab shows up in - # every parallel-formula line below; compute it once. - cross_coef = n_a * n_b / n_ab delta_x = other.mean_x - self.mean_x delta_y = other.mean_y - self.mean_y @@ -249,11 +251,20 @@ def merge(self, other: "WelfordCovariance") -> None: # Three parallel-formula combinations: variance of x, variance # of y, and covariance of (x, y). Each is M_self + M_other plus # a correction for the fact that the two batches had different - # local means. - self.sum_sq_dev_x = self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * cross_coef - self.sum_sq_dev_y = self.sum_sq_dev_y + other.sum_sq_dev_y + delta_y * delta_y * cross_coef + # local means. The (n_a * n_b / n_ab) coefficient is inlined + # rather than precomputed as a Python float so the arithmetic + # stays on the same dtype/device as the M2 / cross-product + # tensors — mirrors WelfordVariance.merge. + self.sum_sq_dev_x = ( + self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * n_a * n_b / n_ab + ) + self.sum_sq_dev_y = ( + self.sum_sq_dev_y + other.sum_sq_dev_y + delta_y * delta_y * n_a * n_b / n_ab + ) self.sum_product_of_devs = ( - self.sum_product_of_devs + other.sum_product_of_devs + delta_x * delta_y * cross_coef + self.sum_product_of_devs + + other.sum_product_of_devs + + delta_x * delta_y * n_a * n_b / n_ab ) self.n_samples = n_ab From 1878aebc92d53217eb4d455f662602cb942d2382 Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Fri, 15 May 2026 14:36:45 +0300 Subject: [PATCH 5/9] review: trim docstrings, function-style tests, prefix-step numpy check Addresses @aaishwarymishra's second pass: 1. Trimmed verbose docstrings on _running_stats.py. Module + class + method docstrings cut by roughly 60 lines without losing the math or the rationale for the parallel formula. The 33-line module docstring is now 12 lines plus a sync_all_reduce code block. 2. Added a sync_all_reduce example to the module docstring. The default dist.all_reduce(SUM) does not work for Welford state (the parallel formula is not a sum of per-rank means); the example shows the all_gather-then-merge pattern that does. 3. Refactored tests from TestWelfordVariance / TestWelfordCovariance classes to standalone test_welford_*_* functions, matching the pure-function style used elsewhere in tests/ignite/metrics/. 4. Added test_welford_variance_matches_numpy_at_each_step and the covariance equivalent. After every individual update, the running stats are compared against numpy's mean/var/cov/corrcoef on the cumulative prefix, not just at the end of the stream. 5. Documented why WelfordCovariance.covariance does not use torch.clamp(min=0): covariance is legitimately signed and clamping would silently bias negatively-correlated pairs toward zero. Added test_welford_covariance_negative_correlation_not_clamped that would fail if a future hand introduced the clamp. 23 tests passing locally. --- ignite/metrics/_running_stats.py | 250 ++++----- tests/ignite/metrics/test_running_stats.py | 557 ++++++++++++--------- 2 files changed, 404 insertions(+), 403 deletions(-) diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py index f33dda10ac91..d7eab6278684 100644 --- a/ignite/metrics/_running_stats.py +++ b/ignite/metrics/_running_stats.py @@ -1,35 +1,42 @@ -"""Numerically stable running variance and covariance helpers. - -Shared internal primitives for metrics that need to accumulate -variance / covariance from streaming batches without falling into -the catastrophic-cancellation trap of the naive ``E[X^2] - E[X]^2`` -formula. Intended consumers (in follow-up PRs of #3748): the -``R2Score`` denominator and ``PearsonCorrelation`` cross-product — -new metrics with the same need should consume these helpers rather -than rolling their own. - -Both classes are dtype and device agnostic dataclasses: state takes -the dtype and device of the first batch passed to :meth:`update`, -and is preserved across subsequent updates and merges. The caller -is responsible for upcasting inputs to ``float64`` when numerical -stability under large means matters — the helper does not silently -promote, because doing so would surprise callers that already work -in the higher-precision dtype. - -Two operations matter: :meth:`merge`, which combines two accumulators -into one via the Chan / Welford parallel formula (also the basis for -cross-rank distributed reductions), and :meth:`update`, which folds a -new batch into the running state. ``update`` is the degenerate case of -``merge`` where ``other`` is a freshly-built single-batch accumulator, -and the implementation reflects that: ``update`` builds the batch -accumulator and delegates to ``merge``. There is one formula, not two. +"""Numerically stable running variance and covariance via Welford's algorithm. + +Shared internals for metrics that accumulate variance or covariance from +streaming batches without the catastrophic cancellation of the naive +``E[X^2] - E[X]^2`` formula. Intended consumers in follow-up PRs of +#3748: :class:`R2Score` denominator and :class:`PearsonCorrelation` +cross-product. + +State is dtype/device agnostic and takes the dtype/device of the first +batch. Cast to ``float64`` caller-side when stability under large means +matters; the helper does not silently promote. + +:meth:`update` and :meth:`merge` share one formula: ``update`` builds +a single-batch accumulator and calls ``merge``. + +Distributed reduction +--------------------- +``sync_all_reduce`` defaults to ``dist.all_reduce(SUM)``, which is not +the right operation for Welford state (the parallel formula is not a +sum of the per-rank means). The pattern is to gather each rank's +accumulator state and merge pairwise:: + + import ignite.distributed as idist + + def compute(self): + ws = self.welford + if idist.get_world_size() > 1: + n = idist.all_gather(torch.tensor([ws.n_samples])) + m = idist.all_gather(ws.mean.reshape(1)) + s = idist.all_gather(ws.sum_sq_dev_from_mean.reshape(1)) + ws = WelfordVariance() + for i in range(len(n)): + ws.merge(WelfordVariance(int(n[i]), m[i], s[i])) + return ws.variance References: - Welford, B. P. (1962). Note on a method for calculating corrected - sums of squares and products. Technometrics 4 (3), 419 to 420. + Welford, B. P. (1962). Technometrics 4(3), 419-420. Chan, T. F., Golub, G. H., LeVeque, R. J. (1979). Updating formulae - and a pairwise algorithm for computing sample variances. - https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm + and a pairwise algorithm for computing sample variances. """ from dataclasses import dataclass, field @@ -39,107 +46,63 @@ @dataclass class WelfordVariance: - """Numerically stable running mean and variance via Welford's algorithm. + """Running mean and population variance via Welford's online algorithm. - Accumulates samples in batches via :meth:`update` and reads off the - mean, variance, or standard deviation through the corresponding - properties. Two accumulators can be combined with :meth:`merge`, - which uses the Chan / Welford parallel formula and is the basis for - distributed reductions. - - No dtype or device handling is performed inside the class: the - state takes the dtype and device of the first batch passed to - :meth:`update`. For numerical stability under large means, callers - should hand in ``float64`` tensors. - - Example:: - - ws = WelfordVariance() - for batch in stream: - ws.update(batch.to(torch.float64)) - print(ws.mean.item(), ws.variance.item(), ws.std.item()) + Fold batches in with :meth:`update`. Read off via :attr:`mean`, + :attr:`variance`, :attr:`std`. Combine two accumulators with + :meth:`merge` (Chan parallel formula). """ - # n_samples: count of samples folded in. - # mean: running sample mean (Welford state). - # sum_sq_dev_from_mean: Σ (x_i − mean)^2, the second central moment - # numerator, conventionally called "M2" in the Welford literature. + # mean: running sample mean. + # sum_sq_dev_from_mean: Σ (x_i - mean)^2, conventionally called M2. n_samples: int = 0 mean: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) sum_sq_dev_from_mean: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) @torch.no_grad() def update(self, batch: torch.Tensor) -> None: - """Fold a batch of samples into the running state. + """Fold ``batch`` into the running state. Empty batches are a no-op. - Implementation: build a single-batch accumulator from ``batch`` - and merge it into ``self``. ``update`` is the degenerate case - of :meth:`merge` where the right-hand side has just been built - from one batch; sharing the parallel formula keeps the two - paths in lock-step. - - Empty batches are silently ignored. ``batch.mean()`` and - ``batch.numel()`` perform a full reduction over the input, so - any shape is accepted and treated as ``numel`` scalar samples. + Any tensor shape is accepted and treated as ``numel`` scalar samples. """ if batch.numel() == 0: return batch = batch.detach() batch_mean = batch.mean() - batch_acc = WelfordVariance( - n_samples=batch.numel(), - mean=batch_mean, - sum_sq_dev_from_mean=(batch - batch_mean).square().sum(), + self.merge( + WelfordVariance( + n_samples=batch.numel(), + mean=batch_mean, + sum_sq_dev_from_mean=(batch - batch_mean).square().sum(), + ) ) - self.merge(batch_acc) @torch.no_grad() def merge(self, other: "WelfordVariance") -> None: - """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. - - Used in two places: by :meth:`update` (where ``other`` is a - freshly-built single-batch accumulator), and by callers that - need to combine independently-accumulated state from elsewhere. - The motivating second case is distributed training: each rank - accumulates its own ``WelfordVariance`` over its local samples, - then at eval time the ranks merge their accumulators rank-by-rank - to produce the population statistic. Without :meth:`merge` that - cross-rank reduction would have to re-iterate the raw data, - which defeats the whole point of an online algorithm. - - Given two accumulators ``A`` and ``B`` with sample counts - ``n_a, n_b`` and second-central-moment sums ``M2_a, M2_b``, the - combined ``M2`` over the concatenated stream is:: + """Combine ``other`` into ``self`` via the Chan parallel formula. + + For two accumulators with sample counts ``n_a, n_b`` and M2 sums + ``M2_a, M2_b``:: M2 = M2_a + M2_b + (mean_b - mean_a)^2 * n_a * n_b / (n_a + n_b) - The third term is the *correction*: simply adding ``M2_a + M2_b`` - would under-count the variance whenever the two batches have - different sample means, because each ``M2`` is measured relative - to its own local mean. The correction folds in the spread of - the two local means about the combined mean. + The third term corrects for the spread of the two local means + about the combined mean. """ if other.n_samples == 0: return if self.n_samples == 0: - # First-time absorb. Copy state so callers cannot mutate - # ``other`` and silently affect ``self``. + # Copy so callers cannot mutate ``other`` and silently affect self. self.n_samples = other.n_samples self.mean = other.mean.detach().clone() self.sum_sq_dev_from_mean = other.sum_sq_dev_from_mean.detach().clone() return - n_a = self.n_samples - n_b = other.n_samples + n_a, n_b = self.n_samples, other.n_samples n_ab = n_a + n_b delta = other.mean - self.mean - # Standard Welford incremental-mean update, weighted by the - # fraction of the combined sample size that ``other`` contributes. self.mean = self.mean + delta * n_b / n_ab - - # Parallel-formula combined M2. The (delta * delta * ...) term - # is the correction described in the docstring above. self.sum_sq_dev_from_mean = ( self.sum_sq_dev_from_mean + other.sum_sq_dev_from_mean + delta * delta * n_a * n_b / n_ab ) @@ -147,9 +110,11 @@ def merge(self, other: "WelfordVariance") -> None: @property def variance(self) -> torch.Tensor: - """Population variance (divisor ``n``). Returns ``0.0`` when empty.""" + """Population variance (divisor ``n``). Zero on an empty accumulator.""" if self.n_samples == 0: return torch.tensor(0.0) + # Variance is non-negative by definition; clamp guards against float + # rounding producing a tiny negative value when all samples are equal. return torch.clamp(self.sum_sq_dev_from_mean / self.n_samples, min=0.0) @property @@ -160,22 +125,14 @@ def std(self) -> torch.Tensor: @dataclass class WelfordCovariance: - """Numerically stable running covariance for a pair of variables (x, y). - - Exposes :attr:`variance_x`, :attr:`variance_y`, :attr:`covariance`, - and :meth:`correlation` (Pearson) through the same Welford-style - online update + Chan / Welford parallel merge as - :class:`WelfordVariance`. The only extension over the univariate - case is the cross-product accumulator - :attr:`sum_product_of_devs` = Σ (x_i - mean_x) (y_i - mean_y). - - Like :class:`WelfordVariance`, the class is dtype and device - agnostic: state takes the dtype and device of the first batch - passed to :meth:`update`. + """Running covariance for a pair ``(x, y)`` via Welford + Chan. + + Same online algorithm as :class:`WelfordVariance`, extended with the + cross-product accumulator ``sum_product_of_devs = Σ (x_i - mean_x)(y_i - mean_y)``. + Read off via :attr:`variance_x`, :attr:`variance_y`, :attr:`covariance`, + :meth:`correlation`. """ - # Two univariate Welford accumulators worth of state, plus the - # cross-product term that turns them into a covariance. n_samples: int = 0 mean_x: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) mean_y: torch.Tensor = field(default_factory=lambda: torch.tensor(0.0)) @@ -185,15 +142,8 @@ class WelfordCovariance: @torch.no_grad() def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: - """Fold a paired batch ``(x_i, y_i)`` into the running state. - - Same trick as :meth:`WelfordVariance.update`: build a single-batch - accumulator from ``(batch_x, batch_y)`` and merge it. One formula, - applied twice; see :meth:`merge` for the math. - - ``batch_x`` and ``batch_y`` must have the same shape; the full - tensor is reduced as ``numel`` scalar samples. - """ + """Fold a paired batch into the running state. ``batch_x`` and + ``batch_y`` must have the same shape.""" if batch_x.shape != batch_y.shape: raise ValueError( f"batch_x and batch_y must have the same shape, got {tuple(batch_x.shape)} and {tuple(batch_y.shape)}." @@ -207,26 +157,22 @@ def update(self, batch_x: torch.Tensor, batch_y: torch.Tensor) -> None: mean_y_b = y.mean() dx = x - mean_x_b dy = y - mean_y_b - batch_acc = WelfordCovariance( - n_samples=x.numel(), - mean_x=mean_x_b, - mean_y=mean_y_b, - sum_sq_dev_x=dx.square().sum(), - sum_sq_dev_y=dy.square().sum(), - sum_product_of_devs=(dx * dy).sum(), + self.merge( + WelfordCovariance( + n_samples=x.numel(), + mean_x=mean_x_b, + mean_y=mean_y_b, + sum_sq_dev_x=dx.square().sum(), + sum_sq_dev_y=dy.square().sum(), + sum_product_of_devs=(dx * dy).sum(), + ) ) - self.merge(batch_acc) @torch.no_grad() def merge(self, other: "WelfordCovariance") -> None: - """Combine ``other`` into ``self`` using the Chan / Welford parallel formula. - - Same correction term as the univariate version, applied three - times: once for ``sum_sq_dev_x``, once for ``sum_sq_dev_y``, and - once for ``sum_product_of_devs`` (using ``delta_x * delta_y`` - instead of ``delta * delta``). See - :meth:`WelfordVariance.merge` for the derivation. - """ + """Combine ``other`` into ``self``. Same correction term as + :meth:`WelfordVariance.merge`, applied once per second moment + (``sum_sq_dev_x``, ``sum_sq_dev_y``, ``sum_product_of_devs``).""" if other.n_samples == 0: return if self.n_samples == 0: @@ -238,23 +184,16 @@ def merge(self, other: "WelfordCovariance") -> None: self.sum_product_of_devs = other.sum_product_of_devs.detach().clone() return - n_a = self.n_samples - n_b = other.n_samples + n_a, n_b = self.n_samples, other.n_samples n_ab = n_a + n_b delta_x = other.mean_x - self.mean_x delta_y = other.mean_y - self.mean_y - # Incremental means, weighted by other's share of the new total. self.mean_x = self.mean_x + delta_x * n_b / n_ab self.mean_y = self.mean_y + delta_y * n_b / n_ab - # Three parallel-formula combinations: variance of x, variance - # of y, and covariance of (x, y). Each is M_self + M_other plus - # a correction for the fact that the two batches had different - # local means. The (n_a * n_b / n_ab) coefficient is inlined - # rather than precomputed as a Python float so the arithmetic - # stays on the same dtype/device as the M2 / cross-product - # tensors — mirrors WelfordVariance.merge. + # Three parallel-formula combinations. Coefficient ``n_a * n_b / n_ab`` + # is inlined per term so arithmetic stays on the operand dtype/device. self.sum_sq_dev_x = ( self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * n_a * n_b / n_ab ) @@ -270,31 +209,34 @@ def merge(self, other: "WelfordCovariance") -> None: @property def variance_x(self) -> torch.Tensor: - """Population variance of ``x`` (divisor ``n``).""" + """Population variance of ``x``.""" if self.n_samples == 0: return torch.tensor(0.0) return torch.clamp(self.sum_sq_dev_x / self.n_samples, min=0.0) @property def variance_y(self) -> torch.Tensor: - """Population variance of ``y`` (divisor ``n``).""" + """Population variance of ``y``.""" if self.n_samples == 0: return torch.tensor(0.0) return torch.clamp(self.sum_sq_dev_y / self.n_samples, min=0.0) @property def covariance(self) -> torch.Tensor: - """Population covariance of ``(x, y)`` (divisor ``n``).""" + """Population covariance of ``(x, y)``. + + No ``torch.clamp`` here because covariance is legitimately signed + (negative correlation gives negative covariance). The variance + properties clamp at zero to guard against float rounding only; + applying the same clamp to covariance would silently bias + negatively-correlated pairs toward zero. + """ if self.n_samples == 0: return torch.tensor(0.0) return self.sum_product_of_devs / self.n_samples def correlation(self, eps: float = 1e-8) -> torch.Tensor: - """Pearson correlation coefficient with a small clamp for safety. - - Args: - eps: floor on the denominator to avoid division by zero when - one of the variables is constant. - """ + """Pearson correlation. ``eps`` floors the denominator so a + constant-variable input returns ``0`` instead of ``NaN``.""" denom = torch.clamp(self.variance_x.sqrt() * self.variance_y.sqrt(), min=eps) return self.covariance / denom diff --git a/tests/ignite/metrics/test_running_stats.py b/tests/ignite/metrics/test_running_stats.py index f2209873afc5..0c6cfbb1b09b 100644 --- a/tests/ignite/metrics/test_running_stats.py +++ b/tests/ignite/metrics/test_running_stats.py @@ -10,145 +10,160 @@ # --------------------------------------------------------------------------- -class TestWelfordVariance: - def test_empty_accumulator(self): - ws = WelfordVariance() - assert ws.n_samples == 0 - assert ws.variance.item() == 0.0 - assert ws.std.item() == 0.0 - - def test_update_then_compute_matches_numpy(self): - rng = np.random.default_rng(0) - # Use float64 throughout so we compare apples to apples; Welford - # upcasts internally, so feeding float32 inputs and comparing against - # float32 numpy stats would understate the helper's precision. - data = rng.standard_normal(1000) - - ws = WelfordVariance() - ws.update(torch.from_numpy(data)) - - assert ws.n_samples == 1000 - assert ws.mean.item() == pytest.approx(float(data.mean()), abs=1e-12) - assert ws.variance.item() == pytest.approx(float(data.var()), rel=1e-12) - - def test_multi_batch_matches_single_batch(self): - # Use float64 so the test exercises the algorithm rather than float32 - # accumulation noise. - rng = np.random.default_rng(1) - data = rng.standard_normal(1000) - data_t = torch.from_numpy(data) - - single = WelfordVariance() - single.update(data_t) - - multi = WelfordVariance() - for start in range(0, 1000, 37): - multi.update(data_t[start : start + 37]) - - assert multi.n_samples == single.n_samples - assert multi.mean.item() == pytest.approx(single.mean.item(), abs=1e-12) - assert multi.variance.item() == pytest.approx(single.variance.item(), rel=1e-12) - - def test_merge_matches_concatenated_update(self): - rng = np.random.default_rng(2) - a = torch.from_numpy(rng.standard_normal(400).astype(np.float64)) - b = torch.from_numpy(rng.standard_normal(600).astype(np.float64)) - - merged = WelfordVariance() - merged.update(a) - right = WelfordVariance() - right.update(b) - merged.merge(right) - - baseline = WelfordVariance() - baseline.update(torch.cat([a, b])) - - assert merged.n_samples == baseline.n_samples - assert merged.mean.item() == pytest.approx(baseline.mean.item(), abs=1e-12) - assert merged.variance.item() == pytest.approx(baseline.variance.item(), rel=1e-12) - - def test_merge_with_empty_accumulators(self): - rng = np.random.default_rng(3) - data = torch.from_numpy(rng.standard_normal(100)) - - # Empty merged into populated -> unchanged. - a = WelfordVariance() - a.update(data) - before_mean = a.mean.item() - a.merge(WelfordVariance()) - assert a.mean.item() == pytest.approx(before_mean, abs=1e-12) - - # Populated merged into empty -> takes the other's state. - b = WelfordVariance() - b.merge(a) - assert b.n_samples == a.n_samples - assert b.mean.item() == pytest.approx(a.mean.item(), abs=1e-12) - - def test_numerical_stability_large_mean_float32(self): - # The whole point of this helper: naive Σx^2 - (Σx)^2/n computed in - # float32 catastrophically cancels at mean=1e6, returning ~0 variance - # (or even negative). Welford fed float64 inputs stays exact. - rng = np.random.default_rng(4) - true_std = 1.0 - data = rng.standard_normal(10_000).astype(np.float32) * true_std + 1e6 - data_t = torch.from_numpy(data) - - # Naive formula in float32 collapses. - sum_x_f32 = data_t.sum() - sum_x2_f32 = (data_t * data_t).sum() - naive_var_f32 = (sum_x2_f32 - sum_x_f32 * sum_x_f32 / len(data_t)) / len(data_t) - # Use float64 ground truth so the assertion isn't measuring our own bug. - true_var = float(np.var(data.astype(np.float64))) - - # The helper is dtype-agnostic; the caller is responsible for the - # float64 upcast. Verify the upcast path recovers the true variance. - ws = WelfordVariance() - ws.update(data_t.to(torch.float64)) - assert ws.variance.item() == pytest.approx(true_var, rel=1e-6) - - # And the naive float32 formula must demonstrably fail on the same - # data so the test documents what we're protecting against. - assert abs(float(naive_var_f32) - true_var) > 0.1, ( - "naive float32 formula did NOT cancel; test is no longer exercising the failure mode it claims to." - ) +def test_welford_variance_empty_accumulator(): + ws = WelfordVariance() + assert ws.n_samples == 0 + assert ws.variance.item() == 0.0 + assert ws.std.item() == 0.0 + + +def test_welford_variance_single_update_matches_numpy(): + rng = np.random.default_rng(0) + data = rng.standard_normal(1000) + + ws = WelfordVariance() + ws.update(torch.from_numpy(data)) + + assert ws.n_samples == 1000 + assert ws.mean.item() == pytest.approx(float(data.mean()), abs=1e-12) + assert ws.variance.item() == pytest.approx(float(data.var()), rel=1e-12) + + +def test_welford_variance_matches_numpy_at_each_step(): + # After every individual update, the running mean/variance must equal + # numpy's mean/variance computed on the cumulative prefix. Catches drift + # in the incremental formula that a single end-of-stream assert would miss. + rng = np.random.default_rng(10) + data = rng.standard_normal(500) + batch_size = 17 + + ws = WelfordVariance() + seen = 0 + for start in range(0, len(data), batch_size): + chunk = data[start : start + batch_size] + ws.update(torch.from_numpy(chunk)) + seen += len(chunk) + prefix = data[:seen] + assert ws.n_samples == seen + assert ws.mean.item() == pytest.approx(float(prefix.mean()), abs=1e-12) + assert ws.variance.item() == pytest.approx(float(prefix.var()), rel=1e-12) + + +def test_welford_variance_multi_batch_matches_single_batch(): + rng = np.random.default_rng(1) + data = rng.standard_normal(1000) + data_t = torch.from_numpy(data) + + single = WelfordVariance() + single.update(data_t) + + multi = WelfordVariance() + for start in range(0, 1000, 37): + multi.update(data_t[start : start + 37]) + + assert multi.n_samples == single.n_samples + assert multi.mean.item() == pytest.approx(single.mean.item(), abs=1e-12) + assert multi.variance.item() == pytest.approx(single.variance.item(), rel=1e-12) + + +def test_welford_variance_merge_matches_concatenated_update(): + rng = np.random.default_rng(2) + a = torch.from_numpy(rng.standard_normal(400).astype(np.float64)) + b = torch.from_numpy(rng.standard_normal(600).astype(np.float64)) + + merged = WelfordVariance() + merged.update(a) + right = WelfordVariance() + right.update(b) + merged.merge(right) + + baseline = WelfordVariance() + baseline.update(torch.cat([a, b])) + + assert merged.n_samples == baseline.n_samples + assert merged.mean.item() == pytest.approx(baseline.mean.item(), abs=1e-12) + assert merged.variance.item() == pytest.approx(baseline.variance.item(), rel=1e-12) + + +def test_welford_variance_merge_with_empty_accumulators(): + rng = np.random.default_rng(3) + data = torch.from_numpy(rng.standard_normal(100)) + + # Empty merged into populated: unchanged. + a = WelfordVariance() + a.update(data) + before_mean = a.mean.item() + a.merge(WelfordVariance()) + assert a.mean.item() == pytest.approx(before_mean, abs=1e-12) + + # Populated merged into empty: takes the other's state. + b = WelfordVariance() + b.merge(a) + assert b.n_samples == a.n_samples + assert b.mean.item() == pytest.approx(a.mean.item(), abs=1e-12) + + +def test_welford_variance_numerical_stability_large_mean_float32(): + # The whole point of this helper: naive Σx^2 - (Σx)^2 / n in float32 + # catastrophically cancels at mean=1e6, returning ~0 variance (or even + # negative). Welford fed float64 inputs stays exact. + rng = np.random.default_rng(4) + true_std = 1.0 + data = rng.standard_normal(10_000).astype(np.float32) * true_std + 1e6 + data_t = torch.from_numpy(data) + + sum_x_f32 = data_t.sum() + sum_x2_f32 = (data_t * data_t).sum() + naive_var_f32 = (sum_x2_f32 - sum_x_f32 * sum_x_f32 / len(data_t)) / len(data_t) + true_var = float(np.var(data.astype(np.float64))) + + ws = WelfordVariance() + ws.update(data_t.to(torch.float64)) + assert ws.variance.item() == pytest.approx(true_var, rel=1e-6) + + # And the naive float32 formula must demonstrably fail on the same data so + # the test documents the failure mode it claims to protect against. + assert abs(float(naive_var_f32) - true_var) > 0.1, ( + "naive float32 formula did NOT cancel; test is no longer exercising the failure mode it claims to." + ) + + +def test_welford_variance_single_sample(): + ws = WelfordVariance() + ws.update(torch.tensor([42.0])) + assert ws.n_samples == 1 + assert ws.mean.item() == 42.0 + assert ws.variance.item() == 0.0 + + +def test_welford_variance_empty_batch_is_noop(): + ws = WelfordVariance() + ws.update(torch.tensor([1.0, 2.0, 3.0])) + before = (ws.n_samples, ws.mean.item(), ws.variance.item()) + ws.update(torch.tensor([])) + after = (ws.n_samples, ws.mean.item(), ws.variance.item()) + assert before == after + + +def test_welford_variance_fresh_instance_has_zero_state(): + ws = WelfordVariance() + assert ws.n_samples == 0 + assert ws.mean.item() == 0.0 + assert ws.sum_sq_dev_from_mean.item() == 0.0 + - def test_single_sample(self): - ws = WelfordVariance() - ws.update(torch.tensor([42.0])) - assert ws.n_samples == 1 - assert ws.mean.item() == 42.0 - assert ws.variance.item() == 0.0 - - def test_empty_batch_is_noop(self): - ws = WelfordVariance() - ws.update(torch.tensor([1.0, 2.0, 3.0])) - before = (ws.n_samples, ws.mean.item(), ws.variance.item()) - ws.update(torch.tensor([])) - after = (ws.n_samples, ws.mean.item(), ws.variance.item()) - assert before == after - - def test_fresh_instance_has_zero_state(self): - # The dataclass starts empty; "reset" is just reconstruction. Verifies - # that the default factories produce an empty accumulator. - ws = WelfordVariance() - assert ws.n_samples == 0 - assert ws.mean.item() == 0.0 - assert ws.sum_sq_dev_from_mean.item() == 0.0 - - def test_state_dtype_follows_first_batch(self): - # The helper does not handle dtype itself; it takes whatever dtype - # the first batch arrives in and preserves it. The caller chooses. - ws_f32 = WelfordVariance() - ws_f32.update(torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)) - assert ws_f32.mean.dtype == torch.float32 - assert ws_f32.mean.item() == pytest.approx(2.5) - assert ws_f32.variance.item() == pytest.approx(1.25) - - ws_f64 = WelfordVariance() - ws_f64.update(torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64)) - assert ws_f64.mean.dtype == torch.float64 - assert ws_f64.mean.item() == pytest.approx(2.5) - assert ws_f64.variance.item() == pytest.approx(1.25) +def test_welford_variance_state_dtype_follows_first_batch(): + ws_f32 = WelfordVariance() + ws_f32.update(torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)) + assert ws_f32.mean.dtype == torch.float32 + assert ws_f32.mean.item() == pytest.approx(2.5) + assert ws_f32.variance.item() == pytest.approx(1.25) + + ws_f64 = WelfordVariance() + ws_f64.update(torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64)) + assert ws_f64.mean.dtype == torch.float64 + assert ws_f64.mean.item() == pytest.approx(2.5) + assert ws_f64.variance.item() == pytest.approx(1.25) # --------------------------------------------------------------------------- @@ -156,120 +171,164 @@ def test_state_dtype_follows_first_batch(self): # --------------------------------------------------------------------------- -class TestWelfordCovariance: - def test_empty_accumulator(self): - wc = WelfordCovariance() - assert wc.n_samples == 0 - assert wc.variance_x.item() == 0.0 - assert wc.variance_y.item() == 0.0 - assert wc.covariance.item() == 0.0 - - def test_update_matches_numpy_corrcoef(self): - rng = np.random.default_rng(5) - n = 1000 - x = rng.standard_normal(n) - y = 0.7 * x + rng.standard_normal(n) * 0.3 - - wc = WelfordCovariance() - wc.update(torch.from_numpy(x), torch.from_numpy(y)) - - np_var_x = float(np.var(x)) - np_var_y = float(np.var(y)) - np_cov = float(np.cov(x, y, bias=True)[0, 1]) - np_r = float(np.corrcoef(x, y)[0, 1]) - - assert wc.variance_x.item() == pytest.approx(np_var_x, rel=1e-12) - assert wc.variance_y.item() == pytest.approx(np_var_y, rel=1e-12) - assert wc.covariance.item() == pytest.approx(np_cov, rel=1e-12) - assert wc.correlation().item() == pytest.approx(np_r, rel=1e-10) - - def test_multi_batch_matches_single_batch(self): - rng = np.random.default_rng(6) - x = torch.from_numpy(rng.standard_normal(900)) - y = torch.from_numpy(rng.standard_normal(900)) - - single = WelfordCovariance() - single.update(x, y) - - multi = WelfordCovariance() - for start in range(0, 900, 31): - multi.update(x[start : start + 31], y[start : start + 31]) - - assert multi.mean_x.item() == pytest.approx(single.mean_x.item(), abs=1e-12) - assert multi.mean_y.item() == pytest.approx(single.mean_y.item(), abs=1e-12) - assert multi.covariance.item() == pytest.approx(single.covariance.item(), rel=1e-12) - assert multi.correlation().item() == pytest.approx(single.correlation().item(), rel=1e-12) - - def test_merge_matches_concatenated_update(self): - rng = np.random.default_rng(7) - x1 = torch.from_numpy(rng.standard_normal(300)) - y1 = torch.from_numpy(rng.standard_normal(300)) - x2 = torch.from_numpy(rng.standard_normal(500)) - y2 = torch.from_numpy(rng.standard_normal(500)) - - merged = WelfordCovariance() - merged.update(x1, y1) - right = WelfordCovariance() - right.update(x2, y2) - merged.merge(right) - - baseline = WelfordCovariance() - baseline.update(torch.cat([x1, x2]), torch.cat([y1, y2])) - - assert merged.covariance.item() == pytest.approx(baseline.covariance.item(), rel=1e-12) - assert merged.correlation().item() == pytest.approx(baseline.correlation().item(), rel=1e-12) - - def test_numerical_stability_large_mean(self): - # The Pearson-correlation regression case from issue #3662: mean=1e6, - # std=1 makes the naive E[X^2] - E[X]^2 formula return garbage in - # float32. Welford fed float64 inputs recovers the true r. - rng = np.random.default_rng(8) - n = 10_000 - x = rng.standard_normal(n).astype(np.float32) + 1e6 - y = (0.99 * x + rng.standard_normal(n).astype(np.float32) * 0.1).astype(np.float32) - - true_r = float(np.corrcoef(x.astype(np.float64), y.astype(np.float64))[0, 1]) - # Sanity: the constructed series really is highly correlated. - assert true_r > 0.99 - - # Caller-side upcast to float64 (the helper preserves whatever it gets). - wc = WelfordCovariance() - wc.update(torch.from_numpy(x).to(torch.float64), torch.from_numpy(y).to(torch.float64)) - assert wc.correlation().item() == pytest.approx(true_r, rel=1e-4) - - def test_shape_mismatch_raises(self): - wc = WelfordCovariance() - with pytest.raises(ValueError, match="same shape"): - wc.update(torch.zeros(5), torch.zeros(6)) - - def test_empty_batch_is_noop(self): - wc = WelfordCovariance() - wc.update(torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])) - before = (wc.n_samples, wc.covariance.item()) - wc.update(torch.tensor([]), torch.tensor([])) - after = (wc.n_samples, wc.covariance.item()) - assert before == after - - def test_constant_variable_correlation_safe(self): - # When one series is constant the denominator of Pearson r is zero; - # the eps clamp keeps us from returning NaN / inf. - wc = WelfordCovariance() - wc.update(torch.tensor([1.0, 2.0, 3.0, 4.0]), torch.tensor([5.0, 5.0, 5.0, 5.0])) - r = wc.correlation().item() - assert r == 0.0 - assert not (r != r) # not NaN - - def test_fresh_instance_has_zero_state(self): - # "Reset" is just reconstruction with this dataclass. - wc = WelfordCovariance() - assert wc.n_samples == 0 - assert wc.covariance.item() == 0.0 +def test_welford_covariance_empty_accumulator(): + wc = WelfordCovariance() + assert wc.n_samples == 0 + assert wc.variance_x.item() == 0.0 + assert wc.variance_y.item() == 0.0 + assert wc.covariance.item() == 0.0 + + +def test_welford_covariance_single_update_matches_numpy(): + rng = np.random.default_rng(5) + n = 1000 + x = rng.standard_normal(n) + y = 0.7 * x + rng.standard_normal(n) * 0.3 + + wc = WelfordCovariance() + wc.update(torch.from_numpy(x), torch.from_numpy(y)) + + np_var_x = float(np.var(x)) + np_var_y = float(np.var(y)) + np_cov = float(np.cov(x, y, bias=True)[0, 1]) + np_r = float(np.corrcoef(x, y)[0, 1]) + + assert wc.variance_x.item() == pytest.approx(np_var_x, rel=1e-12) + assert wc.variance_y.item() == pytest.approx(np_var_y, rel=1e-12) + assert wc.covariance.item() == pytest.approx(np_cov, rel=1e-12) + assert wc.correlation().item() == pytest.approx(np_r, rel=1e-10) + + +def test_welford_covariance_matches_numpy_at_each_step(): + # Same shape of check as the univariate version: after every update, + # running variances + covariance + correlation must match numpy on the + # cumulative prefix. + rng = np.random.default_rng(11) + x = rng.standard_normal(500) + y = 0.5 * x + rng.standard_normal(500) * 0.5 + batch_size = 19 + + wc = WelfordCovariance() + seen = 0 + for start in range(0, len(x), batch_size): + wc.update( + torch.from_numpy(x[start : start + batch_size]), + torch.from_numpy(y[start : start + batch_size]), + ) + seen += len(x[start : start + batch_size]) + px, py = x[:seen], y[:seen] + assert wc.n_samples == seen + assert wc.variance_x.item() == pytest.approx(float(np.var(px)), rel=1e-12) + assert wc.variance_y.item() == pytest.approx(float(np.var(py)), rel=1e-12) + assert wc.covariance.item() == pytest.approx(float(np.cov(px, py, bias=True)[0, 1]), rel=1e-12) + if seen >= 2 and float(np.std(py)) > 0: + assert wc.correlation().item() == pytest.approx(float(np.corrcoef(px, py)[0, 1]), rel=1e-10) + + +def test_welford_covariance_multi_batch_matches_single_batch(): + rng = np.random.default_rng(6) + x = torch.from_numpy(rng.standard_normal(900)) + y = torch.from_numpy(rng.standard_normal(900)) + + single = WelfordCovariance() + single.update(x, y) + + multi = WelfordCovariance() + for start in range(0, 900, 31): + multi.update(x[start : start + 31], y[start : start + 31]) + + assert multi.mean_x.item() == pytest.approx(single.mean_x.item(), abs=1e-12) + assert multi.mean_y.item() == pytest.approx(single.mean_y.item(), abs=1e-12) + assert multi.covariance.item() == pytest.approx(single.covariance.item(), rel=1e-12) + assert multi.correlation().item() == pytest.approx(single.correlation().item(), rel=1e-12) + + +def test_welford_covariance_merge_matches_concatenated_update(): + rng = np.random.default_rng(7) + x1 = torch.from_numpy(rng.standard_normal(300)) + y1 = torch.from_numpy(rng.standard_normal(300)) + x2 = torch.from_numpy(rng.standard_normal(500)) + y2 = torch.from_numpy(rng.standard_normal(500)) + + merged = WelfordCovariance() + merged.update(x1, y1) + right = WelfordCovariance() + right.update(x2, y2) + merged.merge(right) + + baseline = WelfordCovariance() + baseline.update(torch.cat([x1, x2]), torch.cat([y1, y2])) + + assert merged.covariance.item() == pytest.approx(baseline.covariance.item(), rel=1e-12) + assert merged.correlation().item() == pytest.approx(baseline.correlation().item(), rel=1e-12) + + +def test_welford_covariance_negative_correlation_not_clamped(): + # Verifies the documented difference between variance (clamped at 0) + # and covariance (signed). Clamping covariance would silently bias + # negative correlations toward zero. + rng = np.random.default_rng(12) + x = rng.standard_normal(500) + y = -1.0 * x + rng.standard_normal(500) * 0.1 + + wc = WelfordCovariance() + wc.update(torch.from_numpy(x), torch.from_numpy(y)) + + assert wc.covariance.item() < 0 + assert wc.correlation().item() < -0.9 + + +def test_welford_covariance_numerical_stability_large_mean(): + # Pearson regression case from issue #3662: mean=1e6, std=1 makes the + # naive E[X^2] - E[X]^2 formula return garbage in float32. Welford fed + # float64 inputs recovers the true r. + rng = np.random.default_rng(8) + n = 10_000 + x = rng.standard_normal(n).astype(np.float32) + 1e6 + y = (0.99 * x + rng.standard_normal(n).astype(np.float32) * 0.1).astype(np.float32) + + true_r = float(np.corrcoef(x.astype(np.float64), y.astype(np.float64))[0, 1]) + assert true_r > 0.99 + + wc = WelfordCovariance() + wc.update(torch.from_numpy(x).to(torch.float64), torch.from_numpy(y).to(torch.float64)) + assert wc.correlation().item() == pytest.approx(true_r, rel=1e-4) + + +def test_welford_covariance_shape_mismatch_raises(): + wc = WelfordCovariance() + with pytest.raises(ValueError, match="same shape"): + wc.update(torch.zeros(5), torch.zeros(6)) + + +def test_welford_covariance_empty_batch_is_noop(): + wc = WelfordCovariance() + wc.update(torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])) + before = (wc.n_samples, wc.covariance.item()) + wc.update(torch.tensor([]), torch.tensor([])) + after = (wc.n_samples, wc.covariance.item()) + assert before == after + + +def test_welford_covariance_constant_variable_correlation_safe(): + # When one series is constant the denominator of Pearson r is zero; + # the eps clamp keeps us from returning NaN / inf. + wc = WelfordCovariance() + wc.update(torch.tensor([1.0, 2.0, 3.0, 4.0]), torch.tensor([5.0, 5.0, 5.0, 5.0])) + r = wc.correlation().item() + assert r == 0.0 + assert not (r != r) # not NaN + + +def test_welford_covariance_fresh_instance_has_zero_state(): + wc = WelfordCovariance() + assert wc.n_samples == 0 + assert wc.covariance.item() == 0.0 # --------------------------------------------------------------------------- -# Cross-class sanity: variance_x of WelfordCovariance == variance of -# WelfordVariance fed the same x. Catches drift between the two -# implementations. +# Cross-class sanity: variance_x of WelfordCovariance equals variance of +# WelfordVariance fed the same x. Catches drift between the two implementations. # --------------------------------------------------------------------------- From 9bb7da81bf14eb84fae49488618f370e08805918 Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Sun, 24 May 2026 02:22:13 +0300 Subject: [PATCH 6/9] review: simpler all_gather pattern in module docstring @aaishwarymishra noted in #3750 that the three separate all_gather calls in the distributed-reduction example could be collapsed by gathering the dataclass instance directly. idist.all_gather routes non-tensor inputs through _do_all_gather_object, which is wired for every backend including NCCL via a Gloo subgroup, so a single all_gather on self.welford works and reads cleaner. Also fixes two issues in the prior block: result.variance is a @property and was incorrectly shown as result.variance(), and the field name self.welford was inconsistent with the body of the example. Adds a short note about the pickle hop, so a future reader who ports the pattern to a metric with multi-MB state knows to fall back to packing into a flat tensor before all_gather. --- ignite/metrics/_running_stats.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py index d7eab6278684..29cf0a02e4ed 100644 --- a/ignite/metrics/_running_stats.py +++ b/ignite/metrics/_running_stats.py @@ -23,15 +23,21 @@ import ignite.distributed as idist def compute(self): - ws = self.welford if idist.get_world_size() > 1: - n = idist.all_gather(torch.tensor([ws.n_samples])) - m = idist.all_gather(ws.mean.reshape(1)) - s = idist.all_gather(ws.sum_sq_dev_from_mean.reshape(1)) - ws = WelfordVariance() - for i in range(len(n)): - ws.merge(WelfordVariance(int(n[i]), m[i], s[i])) - return ws.variance + collected = idist.all_gather(self.welford) + merged = WelfordVariance() + for item in collected: + merged.merge(item) + return merged.variance + return self.welford.variance + +``idist.all_gather`` of a dataclass instance routes through +``_do_all_gather_object`` (pickle-backed, available for every backend +including NCCL via a Gloo subgroup). Negligible overhead for the +three small tensors carried by these accumulators. Consumers whose +state is significantly larger than a few KB should pack into a flat +tensor before ``all_gather`` and reconstruct on the other side, to +skip the pickle hop. References: Welford, B. P. (1962). Technometrics 4(3), 419-420. From 5a13b17a87a755e9bc0101ff7f9c9031fb89b323 Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Wed, 27 May 2026 17:04:49 +0300 Subject: [PATCH 7/9] review: trim verbose docstring on covariance no-clamp rationale The regression test test_welford_covariance_negative_correlation_not_clamped pins the no-clamp behaviour; the prose is redundant. --- ignite/metrics/_running_stats.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py index 29cf0a02e4ed..4fcc4dc990f4 100644 --- a/ignite/metrics/_running_stats.py +++ b/ignite/metrics/_running_stats.py @@ -229,14 +229,7 @@ def variance_y(self) -> torch.Tensor: @property def covariance(self) -> torch.Tensor: - """Population covariance of ``(x, y)``. - - No ``torch.clamp`` here because covariance is legitimately signed - (negative correlation gives negative covariance). The variance - properties clamp at zero to guard against float rounding only; - applying the same clamp to covariance would silently bias - negatively-correlated pairs toward zero. - """ + """Population covariance of ``(x, y)``.""" if self.n_samples == 0: return torch.tensor(0.0) return self.sum_product_of_devs / self.n_samples From 263863f0e7c304541bc062e0716fd2bcd41336a6 Mon Sep 17 00:00:00 2001 From: blanky Date: Fri, 19 Jun 2026 18:06:56 +0530 Subject: [PATCH 8/9] reduce docstring --- ignite/metrics/_running_stats.py | 59 ++-------------------- tests/ignite/metrics/test_running_stats.py | 12 ----- 2 files changed, 3 insertions(+), 68 deletions(-) diff --git a/ignite/metrics/_running_stats.py b/ignite/metrics/_running_stats.py index 4fcc4dc990f4..5852328bb1f1 100644 --- a/ignite/metrics/_running_stats.py +++ b/ignite/metrics/_running_stats.py @@ -1,50 +1,3 @@ -"""Numerically stable running variance and covariance via Welford's algorithm. - -Shared internals for metrics that accumulate variance or covariance from -streaming batches without the catastrophic cancellation of the naive -``E[X^2] - E[X]^2`` formula. Intended consumers in follow-up PRs of -#3748: :class:`R2Score` denominator and :class:`PearsonCorrelation` -cross-product. - -State is dtype/device agnostic and takes the dtype/device of the first -batch. Cast to ``float64`` caller-side when stability under large means -matters; the helper does not silently promote. - -:meth:`update` and :meth:`merge` share one formula: ``update`` builds -a single-batch accumulator and calls ``merge``. - -Distributed reduction ---------------------- -``sync_all_reduce`` defaults to ``dist.all_reduce(SUM)``, which is not -the right operation for Welford state (the parallel formula is not a -sum of the per-rank means). The pattern is to gather each rank's -accumulator state and merge pairwise:: - - import ignite.distributed as idist - - def compute(self): - if idist.get_world_size() > 1: - collected = idist.all_gather(self.welford) - merged = WelfordVariance() - for item in collected: - merged.merge(item) - return merged.variance - return self.welford.variance - -``idist.all_gather`` of a dataclass instance routes through -``_do_all_gather_object`` (pickle-backed, available for every backend -including NCCL via a Gloo subgroup). Negligible overhead for the -three small tensors carried by these accumulators. Consumers whose -state is significantly larger than a few KB should pack into a flat -tensor before ``all_gather`` and reconstruct on the other side, to -skip the pickle hop. - -References: - Welford, B. P. (1962). Technometrics 4(3), 419-420. - Chan, T. F., Golub, G. H., LeVeque, R. J. (1979). Updating formulae - and a pairwise algorithm for computing sample variances. -""" - from dataclasses import dataclass, field import torch @@ -200,16 +153,10 @@ def merge(self, other: "WelfordCovariance") -> None: # Three parallel-formula combinations. Coefficient ``n_a * n_b / n_ab`` # is inlined per term so arithmetic stays on the operand dtype/device. - self.sum_sq_dev_x = ( - self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * n_a * n_b / n_ab - ) - self.sum_sq_dev_y = ( - self.sum_sq_dev_y + other.sum_sq_dev_y + delta_y * delta_y * n_a * n_b / n_ab - ) + self.sum_sq_dev_x = self.sum_sq_dev_x + other.sum_sq_dev_x + delta_x * delta_x * n_a * n_b / n_ab + self.sum_sq_dev_y = self.sum_sq_dev_y + other.sum_sq_dev_y + delta_y * delta_y * n_a * n_b / n_ab self.sum_product_of_devs = ( - self.sum_product_of_devs - + other.sum_product_of_devs - + delta_x * delta_y * n_a * n_b / n_ab + self.sum_product_of_devs + other.sum_product_of_devs + delta_x * delta_y * n_a * n_b / n_ab ) self.n_samples = n_ab diff --git a/tests/ignite/metrics/test_running_stats.py b/tests/ignite/metrics/test_running_stats.py index 0c6cfbb1b09b..3c852c9e3845 100644 --- a/tests/ignite/metrics/test_running_stats.py +++ b/tests/ignite/metrics/test_running_stats.py @@ -1,15 +1,9 @@ import numpy as np import pytest import torch - from ignite.metrics._running_stats import WelfordCovariance, WelfordVariance -# --------------------------------------------------------------------------- -# WelfordVariance -# --------------------------------------------------------------------------- - - def test_welford_variance_empty_accumulator(): ws = WelfordVariance() assert ws.n_samples == 0 @@ -326,12 +320,6 @@ def test_welford_covariance_fresh_instance_has_zero_state(): assert wc.covariance.item() == 0.0 -# --------------------------------------------------------------------------- -# Cross-class sanity: variance_x of WelfordCovariance equals variance of -# WelfordVariance fed the same x. Catches drift between the two implementations. -# --------------------------------------------------------------------------- - - def test_variance_x_matches_welford_variance(): rng = np.random.default_rng(9) x = torch.from_numpy(rng.standard_normal(1000)) From b8fd9b616cc5bf7d3da83b8112c81480239810bf Mon Sep 17 00:00:00 2001 From: Joe Munene Date: Fri, 19 Jun 2026 22:59:38 +0300 Subject: [PATCH 9/9] review: add end-to-end metric and distributed reduction test Addresses the review on the running-stats helper: - WelfordCorrelation, a minimal Metric over WelfordCovariance, exercises the helper through the real reset/update/compute lifecycle. - Welford state is not additive, so it cannot use sync_all_reduce. compute() combines it across ranks by all_gather of the per-rank state plus pairwise merge, covered by a distributed test that checks the result against the full gathered data. - An end-to-end test drives the mean=1e6 case from #3662 through the metric and confirms the recovered correlation where the naive formula cancels. --- tests/ignite/metrics/test_running_stats.py | 99 ++++++++++++++++++++++ 1 file changed, 99 insertions(+) diff --git a/tests/ignite/metrics/test_running_stats.py b/tests/ignite/metrics/test_running_stats.py index 3c852c9e3845..580f287949d2 100644 --- a/tests/ignite/metrics/test_running_stats.py +++ b/tests/ignite/metrics/test_running_stats.py @@ -1,7 +1,12 @@ import numpy as np import pytest import torch + +import ignite.distributed as idist +from ignite.exceptions import NotComputableError +from ignite.metrics import Metric from ignite.metrics._running_stats import WelfordCovariance, WelfordVariance +from ignite.metrics.metric import reinit__is_reduced def test_welford_variance_empty_accumulator(): @@ -333,3 +338,97 @@ def test_variance_x_matches_welford_variance(): assert wc.variance_x.item() == pytest.approx(wv.variance.item(), rel=1e-12) assert wc.mean_x.item() == pytest.approx(wv.mean.item(), abs=1e-12) + + +# --------------------------------------------------------------------------- +# End to end: WelfordCovariance driven through a real Metric, including the +# cross-rank reduction. Welford state is not additive, so it cannot go through +# sync_all_reduce; compute() gathers each rank's state and folds the parts +# with merge(). +# --------------------------------------------------------------------------- + + +def _merge_across_ranks(local): + device = idist.device() + n = idist.all_gather(torch.tensor([local.n_samples], device=device)) + parts = { + name: idist.all_gather(getattr(local, name).reshape(1).to(device)) + for name in ("mean_x", "mean_y", "sum_sq_dev_x", "sum_sq_dev_y", "sum_product_of_devs") + } + combined = WelfordCovariance() + for i in range(n.numel()): + combined.merge( + WelfordCovariance( + n_samples=int(n[i].item()), + mean_x=parts["mean_x"][i].cpu(), + mean_y=parts["mean_y"][i].cpu(), + sum_sq_dev_x=parts["sum_sq_dev_x"][i].cpu(), + sum_sq_dev_y=parts["sum_sq_dev_y"][i].cpu(), + sum_product_of_devs=parts["sum_product_of_devs"][i].cpu(), + ) + ) + return combined + + +class WelfordCorrelation(Metric): + # Minimal metric over WelfordCovariance, here to exercise the helper through + # reset/update/compute and across ranks. Metrics migrated in #3748 follow + # this shape. + @reinit__is_reduced + def reset(self): + self._cov = WelfordCovariance() + + @reinit__is_reduced + def update(self, output): + y_pred, y = output + self._cov.update(y_pred.double(), y.double()) + + def compute(self): + cov = self._cov if idist.get_world_size() == 1 else _merge_across_ranks(self._cov) + if cov.n_samples == 0: + raise NotComputableError("WelfordCorrelation needs at least one sample before compute.") + return cov.correlation().item() + + +def test_welford_correlation_metric_recovers_3662_case(): + # The helper, driven through a metric's reset/update/compute, recovers the + # correlation on the mean=1e6 data from #3662 that the naive + # E[X^2] - E[X]^2 formula loses to float32 cancellation. Complements the + # helper-level stability test by exercising the full Metric lifecycle. + rng = np.random.default_rng(8) + n = 10_000 + x = rng.standard_normal(n).astype(np.float32) + 1e6 + y = (0.99 * x + rng.standard_normal(n).astype(np.float32) * 0.1).astype(np.float32) + true_r = float(np.corrcoef(x.astype(np.float64), y.astype(np.float64))[0, 1]) + assert true_r > 0.99 + + m = WelfordCorrelation() + for start in range(0, n, 2000): + m.update((torch.from_numpy(x[start : start + 2000]), torch.from_numpy(y[start : start + 2000]))) + assert m.compute() == pytest.approx(true_r, rel=1e-4) + + +def test_welford_correlation_metric_no_samples_raises(): + m = WelfordCorrelation() + with pytest.raises(NotComputableError, match="at least one sample"): + m.compute() + + +@pytest.mark.usefixtures("distributed") +class TestDistributed: + def test_compute_matches_full_data(self): + rank = idist.get_rank() + device = idist.device() + torch.manual_seed(10 + rank) + + x = torch.rand(100, dtype=torch.float64, device=device) + y = 0.6 * x + torch.rand(100, dtype=torch.float64, device=device) * 0.4 + + m = WelfordCorrelation() + m.update((x, y)) + res = m.compute() + + x_all = idist.all_gather(x).cpu().numpy() + y_all = idist.all_gather(y).cpu().numpy() + ref = float(np.corrcoef(x_all, y_all)[0, 1]) + assert res == pytest.approx(ref, rel=1e-5)