Repository: Project-MONAI/MONAI

Issue:
Issues with compute_confusion_matrix metric
**Describe the bug**
There are 2 minor issues in `compute_confusion_matrix` metric:
1. If reduction mode is `SUMBATCH`, should not do SUM on `not_nans`, otherwise we don't know which class is nan when computing the final result.
2. If `compute_sample` is False, should return the confusion matrix elements of current batch directly, then compute the final metric value for the whole epoch. Current implementation computes metric for every batch, if batch size changed, the metric value will be different.
So in summary, should have 2 modes:
(1) Compute metric for every sample image first, then average or sum the whole epoch to get the final result.
(2) Keep the confusion matrix elements of every image, then average or sum the whole epoch, then compute metric.



Discussion:
Hi @yiheng-wang-nv ,

Could you please help fix this issue as we discussed?
Thanks in advance.


Resolving patch:
diff --git a/monai/handlers/confusion_matrix.py b/monai/handlers/confusion_matrix.py
--- a/monai/handlers/confusion_matrix.py
+++ b/monai/handlers/confusion_matrix.py
@@ -13,7 +13,8 @@
 
 import torch
 
-from monai.metrics import ConfusionMatrixMetric
+from monai.metrics import ConfusionMatrixMetric, compute_confusion_matrix_metric
+from monai.metrics.utils import MetricReduction, do_metric_reduction
 from monai.utils import exact_version, optional_import
 
 NotComputableError, _ = optional_import("ignite.exceptions", "0.4.2", exact_version, "NotComputableError")
@@ -31,8 +32,7 @@ def __init__(
         self,
         include_background: bool = True,
         metric_name: str = "hit_rate",
-        compute_sample: bool = True,
-        output_class: bool = False,
+        compute_sample: bool = False,
         output_transform: Callable = lambda x: x,
         device: Optional[torch.device] = None,
     ) -> None:
@@ -48,10 +48,8 @@ def __init__(
                 ``"informedness"``, ``"markedness"``]
                 Some of the metrics have multiple aliases (as shown in the wikipedia page aforementioned),
                 and you can also input those names instead.
-            compute_sample: if ``True``, each sample's metric will be computed first. Defaults to ``True``.
-            output_class: if ``True``, scores for each class will be returned. The final result is each class's score.
-                If average these scores, you will get the macro average of each class. Otherwise, the micro average score
-                of each class will be returned. Defaults to ``False``.
+            compute_sample: if ``True``, each sample's metric will be computed first.
+                If ``False``, the confusion matrix for all samples will be accumulated first. Defaults to ``False``.
             output_transform: transform the ignite.engine.state.output into [y_pred, y] pair.
             device: device specification in case of distributed computation usage.
 
@@ -63,15 +61,25 @@ def __init__(
             include_background=include_background,
             metric_name=metric_name,
             compute_sample=compute_sample,
-            output_class=output_class,
+            reduction=MetricReduction.MEAN,
         )
         self._sum = 0.0
         self._num_examples = 0
+        self.compute_sample = compute_sample
+        self.metric_name = metric_name
+        self._total_tp = 0
+        self._total_fp = 0
+        self._total_tn = 0
+        self._total_fn = 0
 
     @reinit__is_reduced
     def reset(self) -> None:
         self._sum = 0.0
         self._num_examples = 0
+        self._total_tp = 0
+        self._total_fp = 0
+        self._total_tn = 0
+        self._total_fn = 0
 
     @reinit__is_reduced
     def update(self, output: Sequence[torch.Tensor]) -> None:
@@ -86,20 +94,34 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
         if len(output) != 2:
             raise ValueError(f"output must have length 2, got {len(output)}.")
         y_pred, y = output
-        score, not_nans = self.confusion_matrix(y_pred, y)
-        not_nans = int(not_nans.item())
-
-        # add all items in current batch
-        self._sum += score.item() * not_nans
-        self._num_examples += not_nans
-
-    @sync_all_reduce("_sum", "_num_examples")
-    def compute(self) -> float:
+        if self.compute_sample is True:
+            score, not_nans = self.confusion_matrix(y_pred, y)
+            not_nans = int(not_nans.item())
+
+            # add all items in current batch
+            self._sum += score.item() * not_nans
+            self._num_examples += not_nans
+        else:
+            confusion_matrix = self.confusion_matrix(y_pred, y)
+            confusion_matrix, _ = do_metric_reduction(confusion_matrix, MetricReduction.SUM)
+            self._total_tp += confusion_matrix[0].item()
+            self._total_fp += confusion_matrix[1].item()
+            self._total_tn += confusion_matrix[2].item()
+            self._total_fn += confusion_matrix[3].item()
+
+    @sync_all_reduce("_sum", "_num_examples", "_total_tp", "_total_fp", "_total_tn", "_total_fn")
+    def compute(self):
         """
         Raises:
             NotComputableError: When ``compute`` is called before an ``update`` occurs.
 
         """
-        if self._num_examples == 0:
-            raise NotComputableError("ConfusionMatrix metric must have at least one example before it can be computed.")
-        return self._sum / self._num_examples
+        if self.compute_sample is True:
+            if self._num_examples == 0:
+                raise NotComputableError(
+                    "ConfusionMatrix metric must have at least one example before it can be computed."
+                )
+            return self._sum / self._num_examples
+        else:
+            confusion_matrix = torch.tensor([self._total_tp, self._total_fp, self._total_tn, self._total_fn])
+            return compute_confusion_matrix_metric(self.metric_name, confusion_matrix)
diff --git a/monai/metrics/__init__.py b/monai/metrics/__init__.py
--- a/monai/metrics/__init__.py
+++ b/monai/metrics/__init__.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .confusion_matrix import ConfusionMatrixMetric, get_confusion_matrix
+from .confusion_matrix import ConfusionMatrixMetric, compute_confusion_matrix_metric, get_confusion_matrix
 from .hausdorff_distance import compute_hausdorff_distance
 from .meandice import DiceMetric, compute_meandice
 from .occlusion_sensitivity import compute_occlusion_sensitivity
diff --git a/monai/metrics/confusion_matrix.py b/monai/metrics/confusion_matrix.py
--- a/monai/metrics/confusion_matrix.py
+++ b/monai/metrics/confusion_matrix.py
@@ -10,7 +10,7 @@
 # limitations under the License.
 
 import warnings
-from typing import Union
+from typing import Sequence, Union
 
 import torch
 
@@ -39,25 +39,32 @@ class ConfusionMatrixMetric:
             ``"informedness"``, ``"markedness"``]
             Some of the metrics have multiple aliases (as shown in the wikipedia page aforementioned),
             and you can also input those names instead.
-        compute_sample: if ``True``, each sample's metric will be computed first. Defaults to ``False``.
-        output_class: if ``True``, scores for each class will be returned. The final result is each class's score.
-            If average these scores, you will get the macro average of each class. Otherwise, the micro average score
-            of each class will be returned. Defaults to ``False``.
+            Except for input only one metric, multiple metrics are also supported via input a sequence of metric names, such as
+            ("sensitivity", "precision", "recall"), if ``compute_sample`` is ``True``, multiple ``f`` and ``not_nans`` will be
+            returned with the same order as input names when calling the class.
+        compute_sample: if ``True``, each sample's metric will be computed first. If ``False``, the confusion matrix for each image
+            (the output of function ``get_confusion_matrix``) will be returned. In this way, users should achieve the confusion
+            matrixes for all images during an epoch and then use ``compute_confusion_matrix_metric`` to calculate the metric.
+            Defaults to ``False``.
+        reduction: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``,
+            ``"mean_channel"``, ``"sum_channel"``}
+            Define the mode to reduce computation result of 1 batch data. Reduction will only be employed when
+            ``compute_sample`` is ``True``. Defaults to ``"mean"``.
 
     """
 
     def __init__(
         self,
         include_background: bool = True,
-        metric_name: str = "hit_rate",
+        metric_name: Union[Sequence[str], str] = "hit_rate",
         compute_sample: bool = False,
-        output_class: bool = False,
+        reduction: Union[MetricReduction, str] = MetricReduction.MEAN,
     ) -> None:
         super().__init__()
         self.include_background = include_background
         self.metric_name = metric_name
         self.compute_sample = compute_sample
-        self.output_class = output_class
+        self.reduction = reduction
 
     def __call__(self, y_pred: torch.Tensor, y: torch.Tensor):
         """
@@ -91,22 +98,22 @@ def __call__(self, y_pred: torch.Tensor, y: torch.Tensor):
         )
 
         if self.compute_sample:
-            confusion_matrix = compute_confusion_matrix_metric(self.metric_name, confusion_matrix)
-            if self.output_class:
-                f, not_nans = do_metric_reduction(confusion_matrix, MetricReduction.MEAN_BATCH)
+            if isinstance(self.metric_name, str):
+                confusion_matrix = compute_confusion_matrix_metric(self.metric_name, confusion_matrix)
+                f, not_nans = do_metric_reduction(confusion_matrix, self.reduction)
+                return f, not_nans
             else:
-                f, not_nans = do_metric_reduction(confusion_matrix, MetricReduction.MEAN)
+                if len(self.metric_name) < 1:
+                    raise ValueError("the sequence should at least has on metric name.")
+                results = []
+                for metric_name in self.metric_name:
+                    sub_confusion_matrix = compute_confusion_matrix_metric(metric_name, confusion_matrix)
+                    f, not_nans = do_metric_reduction(sub_confusion_matrix, self.reduction)
+                    results.append(f)
+                    results.append(not_nans)
+                return results
         else:
-            if self.output_class:
-                f, _ = do_metric_reduction(confusion_matrix, MetricReduction.SUM_BATCH)
-            else:
-                f, _ = do_metric_reduction(confusion_matrix, MetricReduction.SUM)
-            f = compute_confusion_matrix_metric(self.metric_name, f)
-            nans = torch.isnan(f)
-            not_nans = (~nans).float()
-            f[nans] = 0
-            not_nans = not_nans.sum(dim=0)
-        return f, not_nans
+            return confusion_matrix
 
 
 def get_confusion_matrix(
@@ -145,22 +152,21 @@ def get_confusion_matrix(
         raise ValueError("y_pred and y should have same shapes.")
 
     # get confusion matrix related metric
-    with torch.no_grad():
-        batch_size, n_class = y_pred.shape[:2]
-        # convert to [BNS], where S is the number of pixels for one sample.
-        # As for classification tasks, S equals to 1.
-        y_pred = y_pred.view(batch_size, n_class, -1)
-        y = y.view(batch_size, n_class, -1)
-        tp = ((y_pred + y) == 2).float()
-        tn = ((y_pred + y) == 0).float()
-
-        tp = tp.sum(dim=[2])
-        tn = tn.sum(dim=[2])
-        p = y.sum(dim=[2])
-        n = y.shape[-1] - p
-
-        fn = p - tp
-        fp = n - tn
+    batch_size, n_class = y_pred.shape[:2]
+    # convert to [BNS], where S is the number of pixels for one sample.
+    # As for classification tasks, S equals to 1.
+    y_pred = y_pred.view(batch_size, n_class, -1)
+    y = y.view(batch_size, n_class, -1)
+    tp = ((y_pred + y) == 2).float()
+    tn = ((y_pred + y) == 0).float()
+
+    tp = tp.sum(dim=[2])
+    tn = tn.sum(dim=[2])
+    p = y.sum(dim=[2])
+    n = y.shape[-1] - p
+
+    fn = p - tp
+    fp = n - tn
 
     return torch.stack([tp, fp, tn, fn], dim=-1)
 
@@ -202,6 +208,7 @@ def compute_confusion_matrix_metric(metric_name: str, confusion_matrix: torch.Te
     # calculate metric
     numerator: torch.Tensor
     denominator: Union[torch.Tensor, float]
+    nan_tensor = torch.tensor(float("nan"), device=confusion_matrix.device)
     if metric == "tpr":
         numerator, denominator = tp, p
     elif metric == "tnr":
@@ -219,8 +226,8 @@ def compute_confusion_matrix_metric(metric_name: str, confusion_matrix: torch.Te
     elif metric == "for":
         numerator, denominator = fn, (fn + tn)
     elif metric == "pt":
-        tpr = torch.where(p > 0, tp / p, torch.tensor(float("nan")))
-        tnr = torch.where(n > 0, tn / n, torch.tensor(float("nan")))
+        tpr = torch.where(p > 0, tp / p, nan_tensor)
+        tnr = torch.where(n > 0, tn / n, nan_tensor)
         numerator = torch.sqrt(tpr * (1.0 - tnr)) + tnr - 1.0
         denominator = tpr + tnr - 1.0
     elif metric == "ts":
@@ -228,8 +235,8 @@ def compute_confusion_matrix_metric(metric_name: str, confusion_matrix: torch.Te
     elif metric == "acc":
         numerator, denominator = (tp + tn), (p + n)
     elif metric == "ba":
-        tpr = torch.where(p > 0, tp / p, torch.tensor(float("nan")))
-        tnr = torch.where(n > 0, tn / n, torch.tensor(float("nan")))
+        tpr = torch.where(p > 0, tp / p, nan_tensor)
+        tnr = torch.where(n > 0, tn / n, nan_tensor)
         numerator, denominator = (tpr + tnr), 2.0
     elif metric == "f1":
         numerator, denominator = tp * 2.0, (tp * 2.0 + fn + fp)
@@ -237,18 +244,18 @@ def compute_confusion_matrix_metric(metric_name: str, confusion_matrix: torch.Te
         numerator = tp * tn - fp * fn
         denominator = torch.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))
     elif metric == "fm":
-        tpr = torch.where(p > 0, tp / p, torch.tensor(float("nan")))
-        ppv = torch.where((tp + fp) > 0, tp / (tp + fp), torch.tensor(float("nan")))
+        tpr = torch.where(p > 0, tp / p, nan_tensor)
+        ppv = torch.where((tp + fp) > 0, tp / (tp + fp), nan_tensor)
         numerator = torch.sqrt(ppv * tpr)
         denominator = 1.0
     elif metric == "bm":
-        tpr = torch.where(p > 0, tp / p, torch.tensor(float("nan")))
-        tnr = torch.where(n > 0, tn / n, torch.tensor(float("nan")))
+        tpr = torch.where(p > 0, tp / p, nan_tensor)
+        tnr = torch.where(n > 0, tn / n, nan_tensor)
         numerator = tpr + tnr - 1.0
         denominator = 1.0
     elif metric == "mk":
-        ppv = torch.where((tp + fp) > 0, tp / (tp + fp), torch.tensor(float("nan")))
-        npv = torch.where((tn + fn) > 0, tn / (tn + fn), torch.tensor(float("nan")))
+        ppv = torch.where((tp + fp) > 0, tp / (tp + fp), nan_tensor)
+        npv = torch.where((tn + fn) > 0, tn / (tn + fn), nan_tensor)
         npv = tn / (tn + fn)
         numerator = ppv + npv - 1.0
         denominator = 1.0
@@ -256,7 +263,7 @@ def compute_confusion_matrix_metric(metric_name: str, confusion_matrix: torch.Te
         raise NotImplementedError("the metric is not implemented.")
 
     if isinstance(denominator, torch.Tensor):
-        result = torch.where(denominator != 0, numerator / denominator, torch.tensor(float("nan")))
+        result = torch.where(denominator != 0, numerator / denominator, nan_tensor)
     else:
         result = numerator / denominator
     return result


Test patch:
diff --git a/tests/min_tests.py b/tests/min_tests.py
--- a/tests/min_tests.py
+++ b/tests/min_tests.py
@@ -37,6 +37,8 @@ def run_testsuit():
         "test_handler_checkpoint_saver",
         "test_handler_classification_saver",
         "test_handler_lr_scheduler",
+        "test_handler_confusion_matrix",
+        "test_handler_confusion_matrix_dist",
         "test_handler_mean_dice",
         "test_handler_rocauc",
         "test_handler_rocauc_dist",
@@ -93,7 +95,6 @@ def run_testsuit():
         "test_zoom_affine",
         "test_zoomd",
         "test_compute_occlusion_sensitivity",
-        "test_handler_confusion_matrix",
     ]
     assert sorted(exclude_cases) == sorted(set(exclude_cases)), f"Duplicated items in {exclude_cases}"
 
diff --git a/tests/test_compute_confusion_matrix.py b/tests/test_compute_confusion_matrix.py
--- a/tests/test_compute_confusion_matrix.py
+++ b/tests/test_compute_confusion_matrix.py
@@ -57,6 +57,8 @@
 data_clf: Dict[Any, Any] = {
     "y_pred": torch.tensor([[1, 0, 0], [0, 0, 1]]),
     "y": torch.tensor([[1, 0, 0], [0, 1, 0]]),
+    "compute_sample": False,
+    "include_background": True,
 }
 
 # 1. test confusion matrix
@@ -72,7 +74,8 @@
 
 # 2. test metric with compute_sample
 TEST_CASES_COMPUTE_SAMPLE = []
-result_with_class = [
+TEST_CASES_COMPUTE_SAMPLE_MULTI_METRICS = []
+result_mean_batch = [
     torch.tensor([0.5000, 0.5000, 0.5000]),
     torch.tensor([0.1667, 0.2500, 0.3333]),
     torch.tensor([0.8333, 0.7500, 0.6667]),
@@ -91,7 +94,7 @@
     torch.tensor([0.3333, 0.2500, 0.1667]),
     torch.tensor([0.3333, 0.3333, 0.0833]),
 ]
-result_without_class = [
+result_mean = [
     torch.tensor([0.5000]),
     torch.tensor([0.2500]),
     torch.tensor([0.7500]),
@@ -129,134 +132,78 @@
     "bm",
     "mk",
 ]
+result: Any = None
 for idx in range(len(metric_names)):
-    for output_class in [True, False]:
+    for reduction in ["mean", "mean_batch"]:
         TEST_CASE: List[Any] = [data.copy()]
         TEST_CASE[0]["compute_sample"] = True
         TEST_CASE[0]["include_background"] = True
         TEST_CASE[0]["metric_name"] = metric_names[idx]
-        TEST_CASE[0]["output_class"] = output_class
-        if not output_class:
-            result = result_without_class[idx]
-        else:
-            result = result_with_class[idx]
+        TEST_CASE[0]["reduction"] = reduction
+        if reduction == "mean_batch":
+            result = result_mean_batch[idx]
+        elif reduction == "mean":
+            result = result_mean[idx]
         TEST_CASE.append(result)
         TEST_CASES_COMPUTE_SAMPLE.append(TEST_CASE)
 
+# one input to compute multiple metrics
+for reduction in ["mean", "mean_batch"]:
+    TEST_CASE_MULTIPLE: List[Any] = [data.copy()]
+    TEST_CASE_MULTIPLE[0]["compute_sample"] = True
+    TEST_CASE_MULTIPLE[0]["include_background"] = True
+    TEST_CASE_MULTIPLE[0]["metric_name"] = metric_names
+    TEST_CASE_MULTIPLE[0]["reduction"] = reduction
+    if reduction == "mean_batch":
+        result = result_mean_batch
+    elif reduction == "mean":
+        result = result_mean
+    TEST_CASE_MULTIPLE.append(result)
+    TEST_CASES_COMPUTE_SAMPLE_MULTI_METRICS.append(TEST_CASE_MULTIPLE)
+
 # 3. test metric with compute_sample, denominator may have zeros
 TEST_CASES_COMPUTE_SAMPLE_NAN = []
 metric_names = ["tpr", "tnr"]
-result_with_class = [
-    torch.tensor([0.0000, 0.2500, 0.0000]),
-    torch.tensor([0.8333, 0.8333, 0.2222]),
+result_sum = [
+    torch.tensor([0.5000]),
+    torch.tensor([4.8333]),
+]
+not_nans_sum = [
+    torch.tensor([6]),
+    torch.tensor([8]),
+]
+result_sum_batch = [
+    torch.tensor([0.0000, 0.5000, 0.0000]),
+    torch.tensor([1.6667, 2.5000, 0.6667]),
 ]
-not_nans_class = [
+not_nans_sum_batch = [
     torch.tensor([3.0, 2.0, 1.0]),
     torch.tensor([2.0, 3.0, 3.0]),
 ]
 for idx in range(2):
-    TEST_CASE = [data_nan.copy()]
-    TEST_CASE[0]["compute_sample"] = True
-    TEST_CASE[0]["include_background"] = True
-    TEST_CASE[0]["output_class"] = True
-    TEST_CASE[0]["metric_name"] = metric_names[idx]
-    TEST_CASE.append(result_with_class[idx])
-    TEST_CASE.append(not_nans_class[idx])
-    TEST_CASES_COMPUTE_SAMPLE_NAN.append(TEST_CASE)
-
-# 4. test metric without compute_sample
-TEST_CASES_NO_COMPUTE_SAMPLE = []
-result_with_class = [
-    torch.tensor([0.5000, 0.5000, 0.5000]),
-    torch.tensor([0.1667, 0.2500, 0.3333]),
-    torch.tensor([0.8333, 0.7500, 0.6667]),
-    torch.tensor([0.5000, 0.6667, 0.3333]),
-    torch.tensor([0.8333, 0.6000, 0.8000]),
-    torch.tensor([0.5000, 0.5000, 0.5000]),
-    torch.tensor([0.5000, 0.3333, 0.6667]),
-    torch.tensor([0.1667, 0.4000, 0.2000]),
-    torch.tensor([0.3660, 0.4142, 0.4495]),
-    torch.tensor([0.3333, 0.4000, 0.2500]),
-    torch.tensor([0.7500, 0.6250, 0.6250]),
-    torch.tensor([0.6667, 0.6250, 0.5833]),
-    torch.tensor([0.5000, 0.5714, 0.4000]),
-    torch.tensor([0.3333, 0.2582, 0.1491]),
-    torch.tensor([0.5000, 0.5774, 0.4082]),
-    torch.tensor([0.3333, 0.2500, 0.1667]),
-    torch.tensor([0.3333, 0.2667, 0.1333]),
-]
-result_without_class = [
-    torch.tensor([0.5000]),
-    torch.tensor([0.2500]),
-    torch.tensor([0.7500]),
-    torch.tensor([0.5000]),
-    torch.tensor([0.7500]),
-    torch.tensor([0.5000]),
-    torch.tensor([0.5000]),
-    torch.tensor([0.2500]),
-    torch.tensor([0.4142]),
-    torch.tensor([0.3333]),
-    torch.tensor([0.6667]),
-    torch.tensor([0.6250]),
-    torch.tensor([0.5000]),
-    torch.tensor([0.2500]),
-    torch.tensor([0.5000]),
-    torch.tensor([0.2500]),
-    torch.tensor([0.2500]),
-]
-metric_names = [
-    "tpr",
-    "fpr",
-    "tnr",
-    "ppv",
-    "npv",
-    "fnr",
-    "fdr",
-    "for",
-    "pt",
-    "ts",
-    "acc",
-    "ba",
-    "f1",
-    "mcc",
-    "fm",
-    "bm",
-    "mk",
-]
-for idx in range(len(metric_names)):
-    for output_class in [True, False]:
-        TEST_CASE = [data.copy()]
-        TEST_CASE[0]["compute_sample"] = False
+    for reduction in ["sum", "sum_batch"]:
+        TEST_CASE = [data_nan.copy()]
+        TEST_CASE[0]["compute_sample"] = True
         TEST_CASE[0]["include_background"] = True
+        TEST_CASE[0]["reduction"] = reduction
         TEST_CASE[0]["metric_name"] = metric_names[idx]
-        TEST_CASE[0]["output_class"] = output_class
-        if not output_class:
-            result = result_without_class[idx]
-        else:
-            result = result_with_class[idx]
-        TEST_CASE.append(result)
-        TEST_CASES_NO_COMPUTE_SAMPLE.append(TEST_CASE)
+        if reduction == "sum":
+            TEST_CASE.append(result_sum[idx])
+            TEST_CASE.append(not_nans_sum[idx])
+        elif reduction == "sum_batch":
+            TEST_CASE.append(result_sum_batch[idx])
+            TEST_CASE.append(not_nans_sum_batch[idx])
+        TEST_CASES_COMPUTE_SAMPLE_NAN.append(TEST_CASE)
 
-# 5. test classification task, denominator may have zeros
-TEST_CASES_CLF_NAN = []
-metric_names = ["tpr", "tnr"]
-result_with_class = [
-    torch.tensor([1.0000, 0.0000, 0.0000]),
-    torch.tensor([1.0000, 1.0000, 0.5000]),
-]
-not_nans_class = [
-    torch.tensor(2),
-    torch.tensor(3),
-]
-for idx in range(2):
-    TEST_CASE = [data_clf.copy()]
-    TEST_CASE[0]["compute_sample"] = False
-    TEST_CASE[0]["include_background"] = True
-    TEST_CASE[0]["output_class"] = True
-    TEST_CASE[0]["metric_name"] = metric_names[idx]
-    TEST_CASE.append(result_with_class[idx])
-    TEST_CASE.append(not_nans_class[idx])
-    TEST_CASES_CLF_NAN.append(TEST_CASE)
+# 4. test classification task
+result_clf = torch.tensor(
+    [
+        [[1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 1.0, 0.0]],
+        [[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0]],
+    ]
+)
+
+TEST_CASES_CLF = [data_clf.copy(), result_clf]
 
 
 class TestConfusionMatrix(unittest.TestCase):
@@ -265,10 +212,10 @@ def test_value(self, input_data, expected_value):
         # include or ignore background
         input_data["include_background"] = True
         result = get_confusion_matrix(**input_data)
-        np.testing.assert_allclose(result, expected_value, atol=1e-4)
+        np.testing.assert_allclose(result, expected_value, atol=1e-4, rtol=1e-4)
         input_data["include_background"] = False
         result = get_confusion_matrix(**input_data)
-        np.testing.assert_allclose(result, expected_value[:, 1:, :], atol=1e-4)
+        np.testing.assert_allclose(result, expected_value[:, 1:, :], atol=1e-4, rtol=1e-4)
 
     @parameterized.expand(TEST_CASES_COMPUTE_SAMPLE)
     def test_compute_sample(self, input_data, expected_value):
@@ -278,17 +225,20 @@ def test_compute_sample(self, input_data, expected_value):
         vals["y"] = params.pop("y")
         metric = ConfusionMatrixMetric(**params)
         result, _ = metric(**vals)
-        np.testing.assert_allclose(result, expected_value, atol=1e-4)
+        np.testing.assert_allclose(result, expected_value, atol=1e-4, rtol=1e-4)
 
-    @parameterized.expand(TEST_CASES_NO_COMPUTE_SAMPLE)
-    def test_no_compute_sample(self, input_data, expected_value):
+    @parameterized.expand(TEST_CASES_COMPUTE_SAMPLE_MULTI_METRICS)
+    def test_compute_sample_multiple_metrics(self, input_data, expected_values):
         params = input_data.copy()
         vals = dict()
         vals["y_pred"] = params.pop("y_pred")
         vals["y"] = params.pop("y")
         metric = ConfusionMatrixMetric(**params)
-        result, _ = metric(**vals)
-        np.testing.assert_allclose(result, expected_value, atol=1e-4)
+        results = metric(**vals)
+        for idx in range(0, len(results), 2):
+            result = results[idx]
+            expected_value = expected_values[int(idx / 2)]
+            np.testing.assert_allclose(result, expected_value, atol=1e-4, rtol=1e-4)
 
     @parameterized.expand(TEST_CASES_COMPUTE_SAMPLE_NAN)
     def test_compute_sample_with_nan(self, input_data, expected_value, expected_not_nans):
@@ -298,19 +248,18 @@ def test_compute_sample_with_nan(self, input_data, expected_value, expected_not_
         vals["y"] = params.pop("y")
         metric = ConfusionMatrixMetric(**params)
         result, not_nans = metric(**vals)
-        np.testing.assert_allclose(result, expected_value, atol=1e-4)
-        np.testing.assert_allclose(not_nans, expected_not_nans, atol=1e-4)
+        np.testing.assert_allclose(result, expected_value, atol=1e-4, rtol=1e-4)
+        np.testing.assert_allclose(not_nans, expected_not_nans, atol=1e-4, rtol=1e-4)
 
-    @parameterized.expand(TEST_CASES_CLF_NAN)
-    def test_clf_with_nan(self, input_data, expected_value, expected_not_nans):
+    @parameterized.expand([TEST_CASES_CLF])
+    def test_clf_with_nan(self, input_data, expected_value):
         params = input_data.copy()
         vals = dict()
         vals["y_pred"] = params.pop("y_pred")
         vals["y"] = params.pop("y")
         metric = ConfusionMatrixMetric(**params)
-        result, not_nans = metric(**vals)
-        np.testing.assert_allclose(result, expected_value, atol=1e-4)
-        np.testing.assert_allclose(not_nans, expected_not_nans, atol=1e-4)
+        result = metric(**vals)
+        np.testing.assert_allclose(result, expected_value, atol=1e-4, rtol=1e-4)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_handler_confusion_matrix.py b/tests/test_handler_confusion_matrix.py
--- a/tests/test_handler_confusion_matrix.py
+++ b/tests/test_handler_confusion_matrix.py
@@ -10,14 +10,46 @@
 # limitations under the License.
 
 import unittest
+from typing import Any, Dict
 
 import torch
 from parameterized import parameterized
 
 from monai.handlers import ConfusionMatrix
 
-TEST_CASE_1 = [{"include_background": True, "metric_name": "f1", "compute_sample": True, "output_class": False}, 0.75]
-TEST_CASE_2 = [{"include_background": False, "metric_name": "ppv", "compute_sample": True, "output_class": False}, 1.0]
+TEST_CASE_1 = [{"include_background": True, "metric_name": "f1", "compute_sample": False}, 0.75]
+TEST_CASE_2 = [{"include_background": False, "metric_name": "ppv", "compute_sample": False}, 1.0]
+
+TEST_CASE_SEG_1 = [{"include_background": True, "metric_name": "tpr", "compute_sample": True}, 0.8333]
+TEST_CASE_SEG_2 = [{"include_background": True, "metric_name": "tpr", "compute_sample": False}, 0.7]
+
+data_1: Dict[Any, Any] = {
+    "y_pred": torch.tensor(
+        [
+            [[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]]],
+            [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]],
+        ]
+    ),
+    "y": torch.tensor(
+        [
+            [[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]]],
+            [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]],
+        ]
+    ),
+}
+
+data_2: Dict[Any, Any] = {
+    "y_pred": torch.tensor(
+        [
+            [[[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [0.0, 0.0]]],
+        ]
+    ),
+    "y": torch.tensor(
+        [
+            [[[1.0, 1.0], [1.0, 1.0]], [[1.0, 1.0], [1.0, 1.0]], [[1.0, 1.0], [1.0, 1.0]]],
+        ]
+    ),
+}
 
 
 class TestHandlerConfusionMatrix(unittest.TestCase):
@@ -38,6 +70,23 @@ def test_compute(self, input_params, expected_avg):
         avg_metric = metric.compute()
         self.assertAlmostEqual(avg_metric, expected_avg, places=4)
 
+    @parameterized.expand([TEST_CASE_SEG_1, TEST_CASE_SEG_2])
+    def test_compute_seg(self, input_params, expected_avg):
+        metric = ConfusionMatrix(**input_params)
+
+        y_pred = data_1["y_pred"]
+        y = data_1["y"]
+        metric.update([y_pred, y])
+
+        y_pred = data_2["y_pred"]
+        y = data_2["y"]
+        metric.update([y_pred, y])
+
+        avg_metric = metric.compute()
+        if input_params["compute_sample"] is False:
+            avg_metric = avg_metric.item()
+        self.assertAlmostEqual(avg_metric, expected_avg, places=4)
+
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
     def test_shape_mismatch(self, input_params, _expected):
         metric = ConfusionMatrix(**input_params)
diff --git a/tests/test_handler_confusion_matrix_dist.py b/tests/test_handler_confusion_matrix_dist.py
new file mode 100644
--- /dev/null
+++ b/tests/test_handler_confusion_matrix_dist.py
@@ -0,0 +1,72 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from monai.handlers import ConfusionMatrix
+
+
+def main():
+    for compute_sample in [True, False]:
+        dist.init_process_group(backend="nccl", init_method="env://")
+
+        torch.cuda.set_device(dist.get_rank())
+        metric = ConfusionMatrix(include_background=True, metric_name="tpr", compute_sample=compute_sample)
+
+        if dist.get_rank() == 0:
+            y_pred = torch.tensor(
+                [
+                    [[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]]],
+                    [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]],
+                ],
+                device=torch.device("cuda:0"),
+            )
+            y = torch.tensor(
+                [
+                    [[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]]],
+                    [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]],
+                ],
+                device=torch.device("cuda:0"),
+            )
+            metric.update([y_pred, y])
+
+        if dist.get_rank() == 1:
+            y_pred = torch.tensor(
+                [[[[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [0.0, 0.0]]]],
+                device=torch.device("cuda:1"),
+            )
+            y = torch.tensor(
+                [[[[1.0, 1.0], [1.0, 1.0]], [[1.0, 1.0], [1.0, 1.0]], [[1.0, 1.0], [1.0, 1.0]]]],
+                device=torch.device("cuda:1"),
+            )
+            metric.update([y_pred, y])
+
+        avg_metric = metric.compute()
+        if compute_sample is False:
+            avg_metric = avg_metric.item()
+            np.testing.assert_allclose(avg_metric, 0.7, rtol=1e-04, atol=1e-04)
+        else:
+            np.testing.assert_allclose(avg_metric, 0.8333, rtol=1e-04, atol=1e-04)
+
+        dist.destroy_process_group()
+
+
+# suppose to execute on 2 rank processes
+# python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_PER_NODE
+#        --nnodes=NUM_NODES --node_rank=INDEX_CURRENT_NODE
+#        --master_addr="192.168.1.1" --master_port=1234
+#        test_handler_confusion_matrix_dist.py
+
+if __name__ == "__main__":
+    main()