fix: report gradient snr in db

This commit is contained in:
2026-08-09 13:40:27 +08:00
parent 47b3ed4e44
commit 596c35fd71
2 changed files with 22 additions and 1 deletions
+17
View File
@@ -0,0 +1,17 @@
import pytest
import torch
from astrai.trainer.metric_util import GradSNRTracker
def test_grad_snr_is_reported_in_decibels():
model = torch.nn.Linear(1, 1, bias=False)
tracker = GradSNRTracker(beta=0.5, eps=1e-8)
model.weight.grad = torch.tensor([[1.0]])
tracker.update(model)
model.weight.grad = torch.tensor([[3.0]])
tracker.update(model)
# E[g]^2 / Var(g) = 4 / 1 = 4, which is 6.0206 dB.
assert tracker.snr == pytest.approx(10.0 * torch.log10(torch.tensor(4.0)).item())