add DictMetric (#31)

tests/test_ignite.py

@@ -4,7 +4,9 @@ if sys.version_info.major >= 3:
     import os
     import unittest
     import torch
-    from ignite.engine import create_supervised_trainer
+    from ignite.metrics import RootMeanSquaredError
+    from ignite.engine import create_supervised_trainer, \
+        create_supervised_evaluator
     import torchani
     import torchani.data
 
@@ -37,11 +39,17 @@ if sys.version_info.major >= 3:
             nnp = Flatten(nnp)
             batch_nnp = torchani.models.BatchModel(nnp)
             container = torchani.ignite.Container({'energies': batch_nnp})
-            loss = torchani.ignite.DictLosses({'energies': torch.nn.MSELoss()})
+            loss = torchani.ignite.DictLoss('energies', torch.nn.MSELoss())
             optimizer = torch.optim.SGD(container.parameters(),
                                         lr=0.001, momentum=0.8)
             trainer = create_supervised_trainer(container, optimizer, loss)
             trainer.run(loader, max_epochs=10)
+            metric = torchani.ignite.DictMetric('energies',
+                                                RootMeanSquaredError())
+            evaluator = create_supervised_evaluator(container, metrics={
+                'RMSE': metric
+            })
+            evaluator.run(loader)
 
     if __name__ == '__main__':
         unittest.main()

torchani/ignite/__init__.py

 from .container import Container
-from .dict_loss import DictLosses
+from .loss_metrics import DictLoss, DictMetric
 
-__all__ = ['Container', 'DictLosses']
+__all__ = ['Container', 'DictLoss', 'DictMetric']

torchani/ignite/dict_loss.py

-from torch.nn.modules.loss import _Loss
-
-
-class DictLosses(_Loss):
-
-    def __init__(self, losses):
-        super(DictLosses, self).__init__()
-        self.losses = losses
-
-    def forward(self, input, other):
-        total = 0
-        for i in self.losses:
-            total += self.losses[i](input[i], other[i])
-        return total

torchani/ignite/loss_metrics.py

+from torch.nn.modules.loss import _Loss
+from ignite.metrics.metric import Metric
+
+
+class DictLoss(_Loss):
+
+    def __init__(self, key, loss):
+        super(DictLoss, self).__init__()
+        self.key = key
+        self.loss = loss
+
+    def forward(self, input, other):
+        return self.loss(input[self.key], other[self.key])
+
+
+class DictMetric(Metric):
+
+    def __init__(self, key, metric):
+        self.key = key
+        self.metric = metric
+        super(DictMetric, self).__init__()
+
+    def reset(self):
+        self.metric.reset()
+
+    def update(self, output):
+        y_pred, y = output
+        self.metric.update((y_pred[self.key], y[self.key]))
+
+    def compute(self):
+        self.metric.compute()