Commit 8b736e5d authored by Johan Gudmundsson's avatar Johan Gudmundsson
Browse files

bugfix rv metrics comparing value with itself

parent 30a12f8b

src/borch/metrics/rv_metrics.py

+18 −1
Original line number Diff line number Diff line
@@ -13,6 +13,7 @@ one can use ``all_metrics`` to get all valid metrics for that random variable. 
"""

from torch.distributions import constraints

from borch.metrics import metrics

from borch.module import Module



METRICS = {

    constraints.real.__class__: (metrics.mean_squared_error,),
@@ -27,7 +28,23 @@ METRICS = { 




def _call_metric(rv, metric):

    return metric(rv.tensor, rv)

    return metric(rv.tensor, rv.distribution.sample())





def module_metrics(mod):

    """

    Get the metrics of all observed RV`s in the module



    Note:

        Does only give for RVs directly attached to the module

        not submodules.

    """

    if not isinstance(mod, Module):

        return {}

    return {

        key: all_metrics(getattr(mod.posterior, key))

        for key in dict(mod.observed.items())

    }





def all_metrics(rv):

tests/borch/metrics/test_rv_metrics.py

+16 −4
Original line number Diff line number Diff line
@@ -10,6 +10,19 @@ from borch.utils.torch_utils import get_device 
DEVICE = get_device()





def test_module_metrics():

    mod = borch.module.Module()

    mod.observe(hello=torch.ones(1))

    mod.hello = distributions.Normal(0, 1)

    mod_metrics = rv_metrics.module_metrics(mod)

    assert isinstance(mod_metrics, dict)

    assert isinstance(mod_metrics["hello"]["mean_squared_error"], torch.Tensor)





def test_module_metrics_with_torch_module_gives_empty_dict():

    assert rv_metrics.module_metrics(torch.nn.Module()) == {}





class Test_metric(unittest.TestCase):

    def test_outputs_dict(self):

        rv = distributions.Normal(
@@ -24,10 +37,9 @@ class Test_metric(unittest.TestCase): 
            torch.zeros(2, 4, device=DEVICE), torch.ones(2, 4, device=DEVICE)

        )

        rv.tensor = torch.randn(2, 4).to(DEVICE)

        self.assertTrue(

            rv_metrics.mean_squared_error.__name__

            in borch.metrics.rv_metrics.all_metrics(rv)

        )

        calc_metrics = borch.metrics.rv_metrics.all_metrics(rv)

        self.assertTrue(rv_metrics.mean_squared_error.__name__ in calc_metrics)

        assert float(calc_metrics[rv_metrics.mean_squared_error.__name__]) > 0



    def test_outputs_contains_accuracy_for_categorical(self):

        rv = distributions.Categorical(logits=torch.randn(4, device=DEVICE))