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:
So, `output_transform` signature used in `Loss` initialization should be defined in a generic way in the `Loss` class ? and this signature is related to the criterion ? 

The output dict is compared to `required_output_keys` and here, we need to handle extra args defined in criterion... Not so easy because requirements are not statically known.

EDIT not sure to really understand ...
@sdesrozis yes, you are right about class params vs instance params. We can adopt two approaches:
1) if user would like to pass kwargs for the criterion, there is no way to use dict as state.output
2) in `Loss` we define `required_output_keys=("y_pred", "y", "criterion_kwargs")` and ask the user to define `state.output` as 
```python
state.output = {
  "y_pred": ...,
  "y": ...,
  "criterion_kwargs": {...}
}
```

Let's take 2nd approach.

Hello @vfdev-5 , I am new to PyTorch ignite and studied this issue for some time and understood some part of it related to loss.py but stuck with criterion. Would it be possible to explain the expectation of this issue and some information about the criterion function as well as a state? output. I am interested to work on this issue.
@madhucharan thanks for proposing a help with this issue. Please read https://pytorch.org/ignite/metrics.html#ignite.metrics.Metric.required_output_keys about `required_output_keys`. 
Please, see our [CONTRIBUTING guidelines](https://github.com/pytorch/ignite/blob/master/CONTRIBUTING.md). You can also take a look at Quick-start and Concepts sections of our [docs](https://pytorch.org/ignite/).


Hi @vfdev-5, Thanks for the information. I went through the code base and its functioning related to this issue.

In pytorch we generally use like this,
`criterion = nn.CrossEntropyLoss()`
in training
we use 
`loss = criterion(output, target)`

Is this the criterion mentioned in the first comment? or something else. Kindly clarify this. 

2. And the state.output is actually an output returned by run() function right? I would like to know why user needs to define state.output dict and how and where it is passed as an argument?
@madhucharan well, in most of the cases a loss function in pytorch takes predictions and targets, but there are also cases where loss function can accept other params (e.g. `reduction` for functional cross entropy) and also inputs like sample x etc.
See the link in the first message (https://discuss.pytorch.org/t/how-access-inputs-in-custom-ignite-metric/91221).

The idea of this issue is to simplify Ignite's part of handling those functions.
Try to play with the code snippet from here : https://discuss.pytorch.org/t/how-access-inputs-in-custom-ignite-metric/91221/6

> And the state.output is actually an output returned by run() function right? I would like to know why user needs to define state.output dict and how and where it is passed as an argument?

Passing a dict as the output of the evaluation step function (which is written to `state.output`) is a nice way to structure the output.
In most of the cases, output is a tuple of `(y_pred, y_true)` which are used to compute metrics.

Hope this helps to understand.
 
@vfdev-5 are we going to create an extended class `customMetrics` where users can use custom operations like `reset`, `update`, and `compute`.
@sparkingdark no, the idea is to improve existing `Loss` class
Okay 
@vfdev-5, I was studying this issue for the last two days as I wanted to contribute to PyTorch as a newbie. I looked at all the code and all the comments, and here is what I understood, we already provide a `required_output_key` argument for the `Metric` class and by that, we provide any extra params we may need while calculating custom metric. We need to extend this to the `Loss` metric, where currently the `required_output_key` is `None` ([here](https://pytorch.org/ignite/master/_modules/ignite/metrics/loss.html#Loss)).  So we can do something like this for custom loss :

```python
class CustomLoss(Loss):

    required_output_keys = ("y_pred", "y", "criterion_kwargs")

    def __init__(
        self,
        loss_fn: Callable,
        output_transform: Callable = lambda x: x,
        batch_size: Callable = len,
        device: Union[str, torch.device] = torch.device("cpu"),
    ):
        super(Loss, self).__init__(output_transform, device=device)
        self._loss_fn = loss_fn
        self._batch_size = batch_size

    def reset(self) -> None:

    def update(self, output: Sequence[Union[torch.Tensor, Dict]]) -> None:

    def compute(self) -> float:

```

Then, while calling custom loss we simply do the following:

```python
metrics = {
    "Accuracy": Accuracy(),
    "CustomLoss": CustomLoss(criterion, output_transform=lambda out_dict: (out_dict["y_pred"], out_dict["y"], out_dict["criterion_kwargs"])),
    "CustomMetric": CustomMetric(output_transform=lambda out_dict: (out_dict["y_pred"], out_dict["y"], out_dict["x"])
}

evaluator = create_supervised_evaluator(
    model, 
    metrics=metrics, 
    output_transform=lambda x, y, y_pred, criterion_kwargs, : {"x": x, "y": y, "y_pred": y_pred, "criterion_kwargs":criterion_kwargs}
)
```

Let me know if I have understood this correctly.
@01-vyom thanks for the time on studying this issue. I agree it is not clearly stated what we would like to do here. Sorry about that.

The only thing to do here is to update current implementation of `Loss` by defining `required_output_keys = ("y_pred", "y", "criterion_kwargs")` instead of None and updating the docs saying that we can now interpret output's keys if it is a dictionary like here : https://pytorch.org/ignite/metrics.html#ignite.metrics.Metric.required_output_keys

The main idea is too simplify the code:
- BEFORE
```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}
)
```

- AFTER
```python
metrics = {
    "Accuracy": Accuracy(),
    "Loss": Loss(criterion),
    "CustomMetric": CustomMetric()
}

evaluator = create_supervised_evaluator(
    model, 
    metrics=metrics, 
    output_transform=lambda x, y, y_pred: {"x": x, "y": y, "y_pred": y_pred}
)
```

And if we are in a use-case where user's criterion requires some kwargs: `criterion(y_pred, y, **kwargs)` then our code should work almost like you suggested above:
```python
metrics = {
    "Accuracy": Accuracy(),
    "Loss": Loss(criterion),
    "CustomMetric": CustomMetric()
}

# 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}
)
```
 


Please, let me know if it stille unclear ?

Yes, it is clear. I will change the argument and add a demo code as well as a statement that we accept a dictionary argument to the docs. I will make a PR soon. Thank you for explaining the issue.

Resolving patch:
diff --git a/ignite/metrics/loss.py b/ignite/metrics/loss.py
--- 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
--- 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.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)