add TransformedLoss to allow adding non-linearity to computed loss (#61)

tests/test_ignite.py

@@ -34,8 +34,11 @@ if sys.version_info.major >= 3:
             model = torch.nn.Sequential(prepare, aev_computer, nnp, Flatten())
             container = torchani.ignite.Container({'energies': model})
             optimizer = torch.optim.Adam(container.parameters())
+            loss = torchani.ignite.TransformedLoss(
+                torchani.ignite.MSELoss('energies'),
+                lambda x: torch.exp(x) - 1)
             trainer = create_supervised_trainer(
-                container, optimizer, torchani.ignite.MSELoss('energies'))
+                container, optimizer, loss)
             evaluator = create_supervised_evaluator(container, metrics={
                 'RMSE': torchani.ignite.RMSEMetric('energies')
             })

torchani/ignite/__init__.py

 from .container import Container
-from .loss_metrics import DictLoss, DictMetric, MSELoss, RMSEMetric
+from .loss_metrics import DictLoss, DictMetric, MSELoss, RMSEMetric, \
+    TransformedLoss
 
-__all__ = ['Container', 'DictLoss', 'DictMetric', 'MSELoss', 'RMSEMetric']
+__all__ = ['Container', 'DictLoss', 'DictMetric', 'MSELoss', 'RMSEMetric',
+           'TransformedLoss']

torchani/ignite/loss_metrics.py

@@ -57,5 +57,16 @@ def MSELoss(key, per_atom=True):
         return DictLoss(key, torch.nn.MSELoss())
 
 
+class TransformedLoss(_Loss):
+
+    def __init__(self, origin, transform):
+        super(TransformedLoss, self).__init__()
+        self.origin = origin
+        self.transform = transform
+
+    def forward(self, input, other):
+        return self.transform(self.origin(input, other))
+
+
 def RMSEMetric(key):
     return DictMetric(key, RootMeanSquaredError())