Repository: pytorch/ignite

Issue:
Documentation for ConfusionMatrix is thin on details, and error messages not informative
## 🐛 Bug description

Confusion matrix requires that `y_pred` and `y` are to have different dimensions `(batch,num_classes)` and `(batch,)`, respectively. Getting this wrong results in the following error message, which makes no sense (batch size is 5 here, and both `y_pred` and `y` have shape `(5,1)`:

`y_pred must have shape (batch_size, num_categories, ...) and y must have shape of (batch_size, ...), but given torch.Size([5, 1]) vs torch.Size([5, 1])`

I believe that the `torch.Size([5, 1]) vs torch.Size([5, 1])` message should probably read `torch.Size([5, 1]) vs torch.Size([5,])` to illustrate the problem.

Separately, it would be good to mention in the `ConfusionMatrix` docs that one-hot encoded vectors are required, and that this can be implemented using `from ignite.utils.to_onehot` (I stumbled upon it after realizing I had to one-hot encode, since most people use `BCEWithLogitsLoss()` anyway, and hence their model predictions are not one-hot encoded).

## Environment

 - PyTorch Version (e.g., 1.4): 1.4.0
 - Ignite Version (e.g., 0.3.0): 0.4.2
 - OS (e.g., Linux): MacOS
 - How you installed Ignite (`conda`, `pip`, source): `pip`
 - Python version: 3.8
 - Any other relevant information:


Discussion:
nsinenian: This looks great to me, thanks everyone!

Patch:
diff --git a/ignite/metrics/confusion_matrix.py b/ignite/metrics/confusion_matrix.py
index aaab2bd5..d4fba5d4 100644
--- a/ignite/metrics/confusion_matrix.py
+++ b/ignite/metrics/confusion_matrix.py
@@ -14,13 +14,14 @@ class ConfusionMatrix(Metric):
     """Calculates confusion matrix for multi-class data.
 
     - ``update`` must receive output of the form ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y}``.
-    - `y_pred` must contain logits and has the following shape (batch_size, num_categories, ...)
+    - `y_pred` must contain logits and has the following shape (batch_size, num_classes, ...).
+      If you are doing binary classification, see Note for an example on how to get this.
     - `y` should have the following shape (batch_size, ...) and contains ground-truth class indices
       with or without the background class. During the computation, argmax of `y_pred` is taken to determine
       predicted classes.
 
     Args:
-        num_classes (int): number of classes. See notes for more details.
+        num_classes (int): Number of classes, should be > 1. See notes for more details.
         average (str, optional): confusion matrix values averaging schema: None, "samples", "recall", "precision".
             Default is None. If `average="samples"` then confusion matrix values are normalized by the number of seen
             samples. If `average="recall"` then confusion matrix values are normalized such that diagonal values
@@ -39,6 +40,27 @@ class ConfusionMatrix(Metric):
         contribute to the confusion matrix and others are neglected. For example, if `num_classes=20` and target index
         equal 255 is encountered, then it is filtered out.
 
+        If you are doing binary classification with a single output unit, you may have to transform your network output,
+        so that you have one value for each class. E.g. you can transform your network output into a one-hot vector
+        with:
+
+        .. code-block:: python
+
+            def binary_one_hot_output_transform(output):
+                y_pred, y = output
+                y_pred = torch.sigmoid(y_pred).round().long()
+                y_pred = ignite.utils.to_onehot(y_pred, 2)
+                y = y.long()
+                return y_pred, y
+
+            metrics = {
+                "confusion_matrix": ConfusionMatrix(2, output_transform=binary_one_hot_output_transform),
+            }
+
+            evaluator = create_supervised_evaluator(
+                model, metrics=metrics, output_transform=lambda x, y, y_pred: (y_pred, y)
+            )
+
     """
 
     def __init__(
@@ -51,6 +73,9 @@ class ConfusionMatrix(Metric):
         if average is not None and average not in ("samples", "recall", "precision"):
             raise ValueError("Argument average can None or one of 'samples', 'recall', 'precision'")
 
+        if num_classes <= 1:
+            raise ValueError("Argument num_classes needs to be > 1")
+
         self.num_classes = num_classes
         self._num_examples = 0
         self.average = average
@@ -65,17 +90,18 @@ class ConfusionMatrix(Metric):
         y_pred, y = output[0].detach(), output[1].detach()
 
         if y_pred.ndimension() < 2:
-            raise ValueError(f"y_pred must have shape (batch_size, num_categories, ...), but given {y_pred.shape}")
-
-        if y_pred.shape[1] != self.num_classes:
             raise ValueError(
-                f"y_pred does not have correct number of categories: {y_pred.shape[1]} vs {self.num_classes}"
+                f"y_pred must have shape (batch_size, num_classes (currently set to {self.num_classes}), ...), "
+                f"but given {y_pred.shape}"
             )
 
+        if y_pred.shape[1] != self.num_classes:
+            raise ValueError(f"y_pred does not have correct number of classes: {y_pred.shape[1]} vs {self.num_classes}")
+
         if not (y.ndimension() + 1 == y_pred.ndimension()):
             raise ValueError(
-                "y_pred must have shape (batch_size, num_categories, ...) and y must have "
-                "shape of (batch_size, ...), "
+                f"y_pred must have shape (batch_size, num_classes (currently set to {self.num_classes}), ...) "
+                "and y must have shape of (batch_size, ...), "
                 f"but given {y.shape} vs {y_pred.shape}."
             )
 


Test patch:
diff --git a/tests/ignite/metrics/test_confusion_matrix.py b/tests/ignite/metrics/test_confusion_matrix.py
index 9c51060a..405a06b0 100644
--- a/tests/ignite/metrics/test_confusion_matrix.py
+++ b/tests/ignite/metrics/test_confusion_matrix.py
@@ -19,17 +19,27 @@ def test_no_update():
         cm.compute()
 
 
+def test_num_classes_wrong_input():
+    with pytest.raises(ValueError, match="Argument num_classes needs to be > 1"):
+        ConfusionMatrix(num_classes=1)
+
+
 def test_multiclass_wrong_inputs():
     cm = ConfusionMatrix(10)
 
-    with pytest.raises(ValueError, match=r"y_pred must have shape \(batch_size, num_categories, ...\)"):
+    with pytest.raises(
+        ValueError, match=r"y_pred must have shape \(batch_size, num_classes " r"\(currently set to 10\), ...\)"
+    ):
         cm.update((torch.rand(10), torch.randint(0, 2, size=(10,)).long()))
 
-    with pytest.raises(ValueError, match=r"y_pred does not have correct number of categories:"):
+    with pytest.raises(ValueError, match=r"y_pred does not have correct number of classes:"):
         cm.update((torch.rand(10, 5, 4), torch.randint(0, 2, size=(10,)).long()))
 
     with pytest.raises(
-        ValueError, match=r"y_pred must have shape \(batch_size, num_categories, ...\) " r"and y must have "
+        ValueError,
+        match=r"y_pred must have shape \(batch_size, num_classes "
+        r"\(currently set to 10\), ...\) "
+        r"and y must have ",
     ):
         cm.update((torch.rand(4, 10, 12, 12), torch.randint(0, 10, size=(10,)).long()))