Repository: pytorch/ignite

Issue:
Loss metric to use required_output_keys
## 🚀 Feature

Currently, if we have custom metrics that require data other then `y_pred` and `y`, [we suggest](https://discuss.pytorch.org/t/how-access-inputs-in-custom-ignite-metric/91221/6) to do the following: 
```python
metrics = {
    "Accuracy": Accuracy(),
    "Loss": Loss(criterion, output_transform=lambda out_dict: (out_dict["y_pred"], out_dict["y"])),
    "CustomMetric": CustomMetric()
}

evaluator = create_supervised_evaluator(
    model, 
    metrics=metrics, 
    output_transform=lambda x, y, y_pred: {"x": x, "y": y, "y_pred": y_pred}
)
```

where `CustomMetric` is defined as 
```python
class CustomMetric(Metric):

    required_output_keys = ("y_pred", "y", "x")
```

The idea is to extend this for `Loss` metric to support `required_output_keys`. The main issue with `Loss` now is with `(prediction, target, kwargs)` optional input, where `kwargs` is a dict for extra args for criterion function.

Discussion:
01-vyom: Following is the demo code to check the validity of the above change:

```python
import torch
import torch.nn as nn
from torch.nn.functional import nll_loss

from ignite.metrics import Accuracy, Loss
from ignite.engine import create_supervised_evaluator

model = nn.Linear(10, 3)

metrics = {
    "Accuracy": Accuracy(),
    "Loss": Loss(nll_loss)
}

# global criterion kwargs
criterion_kwargs = {"reduction": 'mean'}
# criterion_kwargs = {}

evaluator = create_supervised_evaluator(
    model,
    metrics=metrics,
    output_transform=lambda x, y, y_pred: {
        "x": x, "y": y, "y_pred": y_pred, "criterion_kwargs": criterion_kwargs}
)
data = [
    (torch.rand(4, 10), torch.randint(0, 3, size=(4, ))),
    (torch.rand(4, 10), torch.randint(0, 3, size=(4, ))),
    (torch.rand(4, 10), torch.randint(0, 3, size=(4, )))
]
res = evaluator.run(data)
```

As the `required_output_keys` contains `criterion_kwargs`, the user has to pass an empty dictionary for the original case with no `criterion_kwargs` for criterion. 

vfdev-5: @01-vyom thanks for the draft PR ! 
Let's now add some tests here:
- https://github.com/pytorch/ignite/blob/master/tests/ignite/metrics/test_loss.py
like here: 
- https://github.com/pytorch/ignite/blob/692487ce01ec2d212ca1693ebe53a46f255f5584/tests/ignite/metrics/test_metric.py#L98

Instead of `DummyMetric1`, we can create a metric derived from `Loss` and override update method to check passed args similarly to here: 
https://github.com/pytorch/ignite/blob/692487ce01ec2d212ca1693ebe53a46f255f5584/tests/ignite/metrics/test_metric.py#L28-L29

Please let me know if need more explanations.

01-vyom: Ok, I will add a `test_output_mapping` test with a dummy `Loss` and an integration test similar to here:
https://github.com/pytorch/ignite/blob/692487ce01ec2d212ca1693ebe53a46f255f5584/tests/ignite/metrics/test_metric.py#L894

Patch:
diff --git a/ignite/metrics/loss.py b/ignite/metrics/loss.py
index 3cf90a75..fbda23b8 100644
--- a/ignite/metrics/loss.py
+++ b/ignite/metrics/loss.py
@@ -30,9 +30,52 @@ class Loss(Metric):
             metric's device to be the same as your ``update`` arguments ensures the ``update`` method is
             non-blocking. By default, CPU.
 
+    Attributes:
+        required_output_keys: dictionary defines required keys to be found in ``engine.state.output`` if the
+            latter is a dictionary. Default, ``("y_pred", "y", "criterion_kwargs")``. This is useful when the
+            criterion function requires additional arguments, which can be passed using ``criterion_kwargs``.
+            See notes below for an example.
+
+    Note:
+
+        Let's implement a Loss metric that requires ``x``, ``y_pred``, ``y`` and ``criterion_kwargs`` as input
+        for ``criterion`` function. In the example below we show how to setup standard metric like Accuracy
+        and the Loss metric using an ``evaluator`` created with
+        :meth:`~ignite.engine.create_supervised_evaluator` method.
+
+        .. code-block:: python
+
+            import torch
+            import torch.nn as nn
+            from torch.nn.functional import nll_loss
+
+            from ignite.metrics import Accuracy, Loss
+            from ignite.engine import create_supervised_evaluator
+
+            model = ...
+
+            criterion = nll_loss
+
+            metrics = {
+                "Accuracy": Accuracy(),
+                "Loss": Loss(criterion)
+            }
+
+            # global criterion kwargs
+            criterion_kwargs = {...}
+
+            evaluator = create_supervised_evaluator(
+                model,
+                metrics=metrics,
+                output_transform=lambda x, y, y_pred: {
+                    "x": x, "y": y, "y_pred": y_pred, "criterion_kwargs": criterion_kwargs}
+            )
+
+            res = evaluator.run(data)
+
     """
 
-    required_output_keys = None
+    required_output_keys = ("y_pred", "y", "criterion_kwargs")
 
     def __init__(
         self,


Test patch:
diff --git a/tests/ignite/metrics/test_loss.py b/tests/ignite/metrics/test_loss.py
index 23749b96..43adfc19 100644
--- a/tests/ignite/metrics/test_loss.py
+++ b/tests/ignite/metrics/test_loss.py
@@ -1,4 +1,5 @@
 import os
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -7,8 +8,48 @@ from torch import nn
 from torch.nn.functional import nll_loss
 
 import ignite.distributed as idist
+from ignite.engine import State
 from ignite.exceptions import NotComputableError
-from ignite.metrics import Loss
+from ignite.metrics import Loss, Precision
+
+
+class DummyLoss1(Loss):
+    def __init__(self, loss_fn, true_output, output_transform=lambda x: x):
+        super(DummyLoss1, self).__init__(loss_fn, output_transform=output_transform)
+        print(true_output)
+        self.true_output = true_output
+
+    def reset(self):
+        pass
+
+    def compute(self):
+        pass
+
+    def update(self, output):
+
+        assert output == self.true_output
+
+
+def test_output_as_mapping_without_criterion_kwargs():
+    y_pred = torch.Tensor([[2.0], [-2.0]])
+    y = torch.zeros(2)
+    criterion_kwargs = {}
+
+    loss_metric = DummyLoss1(nll_loss, true_output=(y_pred, y, criterion_kwargs))
+    state = State(output=({"y_pred": y_pred, "y": y, "criterion_kwargs": {}}))
+    engine = MagicMock(state=state)
+    loss_metric.iteration_completed(engine)
+
+
+def test_output_as_mapping_with_criterion_kwargs():
+    y_pred = torch.Tensor([[2.0], [-2.0]])
+    y = torch.zeros(2)
+    criterion_kwargs = {"reduction": "sum"}
+
+    loss_metric = DummyLoss1(nll_loss, true_output=(y_pred, y, criterion_kwargs))
+    state = State(output=({"y_pred": y_pred, "y": y, "criterion_kwargs": {"reduction": "sum"}}))
+    engine = MagicMock(state=state)
+    loss_metric.iteration_completed(engine)
 
 
 def y_test_1(requires_grad=False, device=None):
@@ -252,3 +293,48 @@ def test_multinode_distrib_nccl_gpu(distributed_context_multi_node_nccl):
     device = idist.device()
     _test_distrib_compute_on_criterion(device, y_test_1(), y_test_2())
     _test_distrib_accumulator_device(device, y_test_1())
+
+
+def test_override_required_output_keys():
+    # https://github.com/pytorch/ignite/issues/1415
+    from ignite.engine import create_supervised_evaluator
+
+    counter = [0]
+
+    class DummyLoss2(Loss):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+        def update(self, output):
+            y_pred, y, criterion_kwargs = output
+            assert y_pred.shape == (4, 3)
+            assert y.shape == (4,)
+            assert criterion_kwargs == c_kwargs
+            assert y.equal(data[counter[0]][1])
+            counter[0] += 1
+
+        def reset(self):
+            pass
+
+        def compute(self):
+            pass
+
+    model = nn.Linear(10, 3)
+
+    metrics = {"Precision": Precision(), "DummyLoss2": DummyLoss2(nll_loss)}
+
+    # global criterion kwargs
+    c_kwargs = {"reduction": "sum"}
+
+    evaluator = create_supervised_evaluator(
+        model,
+        metrics=metrics,
+        output_transform=lambda x, y, y_pred: {"x": x, "y": y, "y_pred": y_pred, "criterion_kwargs": c_kwargs},
+    )
+
+    data = [
+        (torch.rand(4, 10), torch.randint(0, 3, size=(4,))),
+        (torch.rand(4, 10), torch.randint(0, 3, size=(4,))),
+        (torch.rand(4, 10), torch.randint(0, 3, size=(4,))),
+    ]
+    evaluator.run(data)