Repository: continualai/avalanche
Issue: StreamConfusionMatrix Stalling Problem

**Describe the bug**
The script provided below stalls (or it is super slow) to move to the next experience. 

**To Reproduce**

```
import torch
import torch.optim as optim
import torchvision.transforms as transforms
from avalanche.benchmarks.classic import CORe50
from avalanche.evaluation.metrics import ExperienceForgetting, accuracy_metrics, loss_metrics, timing_metrics, \\
    cpu_usage_metrics, StreamConfusionMatrix, disk_usage_metrics
from avalanche.logging import InteractiveLogger, TextLogger, TensorboardLogger
from avalanche.training.plugins import EvaluationPlugin
from avalanche.training.strategies import Naive
from torchvision.models import resnet34

train_transform = transforms.Compose([
        transforms.ToTensor()
])
eval_transform = transforms.Compose([
    transforms.ToTensor()
])
scenario = CORe50(
    scenario="ni",
    train_transform=train_transform,
    eval_transform=eval_transform)
model = resnet34(pretrained=True)

tb_logger = TensorboardLogger(tb_log_dir="experiment")
interactive_logger = InteractiveLogger()

eval_plugin = EvaluationPlugin(
    accuracy_metrics(minibatch=True, epoch=True, epoch_running=True, experience=True, stream=True),
    loss_metrics(minibatch=True, epoch=True, epoch_running=True, experience=True, stream=True),
    StreamConfusionMatrix(num_classes=50, save_image=True), # <---- disable this to make the problem disappear
    loggers=[interactive_logger, tb_logger]
)
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
criterion = torch.nn.CrossEntropyLoss()
batch_size = 16
epochs = 1
device = torch.device("cuda")

cl_strategy = Naive(
        model=model,
        optimizer=optimizer,
        criterion=criterion,
        train_mb_size=batch_size,
        train_epochs=epochs,
        eval_mb_size=batch_size,
        evaluator=eval_plugin,
        device=device
)

print("Starting experiment...")
print(len(scenario.train_stream))

for experience in scenario.train_stream[:2]:
    print("Start of experience ", experience.current_experience)
    print("Current classes: ", experience.classes_in_this_experience)

    cl_strategy.train(experience, num_workers=8)
    print("Training complete.")

    print("Computing accuracy on the test set")
    cl_strategy.eval(scenario.test_stream, num_workers=8)
    print("End of experience ", experience.current_experience)
```

**Expected behavior**
The script should go over the next experience immediately. 

**Additional context**
See discussion here: https://github.com/ContinualAI/avalanche/discussions/439#discussioncomment-585228


Changed files:
avalanche/training/plugins/evaluation.py, avalanche/evaluation/metric_utils.py, avalanche/logging/tensorboard_logger.py, avalanche/evaluation/metrics/confusion_matrix.py, avalanche/logging/wandb_logger.py

diff --git a/avalanche/evaluation/metric_utils.py b/avalanche/evaluation/metric_utils.py
index bbd0879b..207c7f93 100644
--- a/avalanche/evaluation/metric_utils.py
+++ b/avalanche/evaluation/metric_utils.py
@@ -13,6 +13,7 @@ import io
 from typing import Tuple, TYPE_CHECKING
 
 import matplotlib.pyplot as plt
+import numpy as np
 from PIL import Image
 from sklearn.metrics import ConfusionMatrixDisplay
 from torch import Tensor
@@ -29,19 +30,16 @@ TRAIN = "train"
 
 def default_cm_image_creator(confusion_matrix_tensor: Tensor,
                              display_labels=None,
-                             include_values=True,
-                             xticks_rotation='horizontal',
+                             include_values=False,
+                             xticks_rotation=0,
+                             yticks_rotation=0,
                              values_format=None,
                              cmap='viridis',
-                             dpi=100,
                              image_title=''):
     """
-    The default Confusion Matrix image creator. This utility uses Scikit-learn
-    `ConfusionMatrixDisplay` to create the matplotlib figure. The figure
-    is then converted to a PIL `Image`.
-
-    For more info about the accepted graphic parameters,
-    see `Scikit learn <https://scikit-learn.org/stable/modules/generated/sklearn.metrics.plot_confusion_matrix.html>`_ # noqa
+    The default Confusion Matrix image creator.
+    Code adapted from
+    `Scikit learn <https://scikit-learn.org/stable/modules/generated/sklearn.metrics.plot_confusion_matrix.html>`_ # noqa
 
     :param confusion_matrix_tensor: The tensor describing the confusion matrix.
         This can be easily obtained through Scikit-learn `confusion_matrix`
@@ -52,34 +50,70 @@ def default_cm_image_creator(confusion_matrix_tensor: Tensor,
     :param include_values: Includes values in confusion matrix. Defaults to
         `True`.
     :param xticks_rotation: Rotation of xtick labels. Valid values are
-        'vertical', 'horizontal' or a float point value. Defaults to
-        'horizontal'.
+        float point value. Defaults to 0.
+    :param yticks_rotation: Rotation of ytick labels. Valid values are
+        float point value. Defaults to 0.
     :param values_format: Format specification for values in confusion matrix.
         Defaults to `None`, which means that the format specification is
         'd' or '.2g', whichever is shorter.
     :param cmap: Must be a str or a Colormap recognized by matplotlib.
         Defaults to 'viridis'.
-    :param dpi: The dpi to use to save the image.
     :param image_title: The title of the image. Defaults to an empty string.
     :return: The Confusion Matrix as a PIL Image.
     """
 
-    display = ConfusionMatrixDisplay(
-        confusion_matrix=confusion_matrix_tensor.numpy(),
-        display_labels=display_labels)
-    display.plot(include_values=include_values, cmap=cmap,
-                 xticks_rotation=xticks_rotation, values_format=values_format)
+    fig, ax = plt.subplots()
+
+    cm = confusion_matrix_tensor.numpy()
+    n_classes = cm.shape[0]
+    im_ = ax.imshow(cm, interpolation='nearest', cmap=cmap)
+    cmap_min, cmap_max = im_.cmap(0), im_.cmap(256)
+
+    if include_values:
+        text_ = np.empty_like(cm, dtype=object)
+
+        # print text with appropriate color depending on background
+        thresh = (cm.max() + cm.min()) / 2.0
+
+        for i in range(n_classes):
+            for j in range(n_classes):
+                color = cmap_max if cm[i, j] < thresh else cmap_min
+
+                if values_format is None:
+                    text_cm = format(cm[i, j], '.2g')
+                    if cm.dtype.kind != 'f':
+                        text_d = format(cm[i, j], 'd')
+                        if len(text_d) < len(text_cm):
+                            text_cm = text_d
+                else:
+                    text_cm = format(cm[i, j], values_format)
+
+                text_[i, j] = ax.text(
+                    j, i, text_cm,
+                    ha="center", va="center",
+                    color=color)
+
+    if display_labels is None:
+        display_labels = np.arange(n_classes)
+
+    fig.colorbar(im_, ax=ax)
+
+    ax.set(xticks=np.arange(n_classes),
+           yticks=np.arange(n_classes),
+           xticklabels=display_labels,
+           yticklabels=display_labels,
+           ylabel="True label",
+           xlabel="Predicted label")
+
+    if image_title != '':
+        ax.set_title(image_title)
 
-    display.ax_.set_title(image_title)
+    ax.set_ylim((n_classes - 0.5, -0.5))
+    plt.setp(ax.get_xticklabels(), rotation=xticks_rotation)
+    plt.setp(ax.get_yticklabels(), rotation=yticks_rotation)
 
-    fig = display.figure_
     fig.tight_layout()
-    buf = io.BytesIO()
-    fig.savefig(buf, format='jpg', dpi=dpi)
-    plt.close(fig)
-    buf.seek(0)
-    image = Image.open(buf)
-    return image
+    return fig
 
 
 def get_task_label(strategy: 'BaseStrategy') -> int:
diff --git a/avalanche/evaluation/metrics/confusion_matrix.py b/avalanche/evaluation/metrics/confusion_matrix.py
index 639f26f7..ed8b5aec 100644
--- a/avalanche/evaluation/metrics/confusion_matrix.py
+++ b/avalanche/evaluation/metrics/confusion_matrix.py
@@ -43,8 +43,13 @@ class ConfusionMatrix(Metric[Tensor]):
     predictions Tensors. When passing one-hot/logit vectors, this
     metric will try to infer the number of classes from the vector sizes.
     Otherwise, the maximum label value encountered in the truth/prediction
-    Tensors will be used. It is recommended to set the (initial) number of
-    classes in the constructor.
+    Tensors will be used.
+
+    If the user sets the `num_classes`, then the confusion matrix will always be
+    of size `num_classes, num_classes`. Whenever a prediction or label tensor is
+    provided as logits, only the first `num_classes` units will be considered in
+    the confusion matrix computation. If they are provided as numerical labels,
+    each of them has to be smaller than `num_classes`.
 
     The reset method will bring the metric to its initial state. By default
     this metric in its initial state will return an empty Tensor.
@@ -59,10 +64,14 @@ class ConfusionMatrix(Metric[Tensor]):
         The metric can be updated by using the `update` method while the running
         confusion matrix can be retrieved using the `result` method.
 
-        :param num_classes: The initial number of classes. Defaults to None,
+        :param num_classes: The number of classes. Defaults to None,
             which means that the number of classes will be inferred from
             ground truth and prediction Tensors (see class description for more
-            details).
+            details). If not None, the confusion matrix will always be of size
+            `num_classes, num_classes` and only the first `num_classes` values
+            of output logits or target logits will be considered in the update.
+            If the output or targets are provided as numerical labels,
+            there can be no label greater than `num_classes`.
         :param normalize: how to normalize confusion matrix.
             None to not normalize
         """
@@ -88,24 +97,44 @@ class ConfusionMatrix(Metric[Tensor]):
         if len(true_y) != len(predicted_y):
             raise ValueError('Size mismatch for true_y and predicted_y tensors')
 
-        max_label = -1
+        if len(true_y.shape) > 2:
+            raise ValueError('Confusion matrix supports labels with at'
+                             ' most 2 dimensions')
+        if len(predicted_y.shape) > 2:
+            raise ValueError('Confusion matrix supports predictions with at '
+                             'most 2 dimensions')
+
+        max_label = -1 if self._num_classes is None else self._num_classes - 1
+
+        # SELECT VALID PORTION OF TARGET AND PREDICTIONS
         true_y = torch.as_tensor(true_y)
+        if len(true_y.shape) == 2 and self._num_classes is not None:
+            true_y = true_y[:, :max_label]
         predicted_y = torch.as_tensor(predicted_y)
+        if len(predicted_y.shape) == 2 and self._num_classes is not None:
+            predicted_y = predicted_y[:, :max_label]
 
-        # Check if logits or labels
+        # COMPUTE MAX LABEL AND CONVERT TARGET AND PREDICTIONS IF NEEDED
         if len(predicted_y.shape) > 1:
             # Logits -> transform to labels
-            max_label = max(max_label, predicted_y.shape[1]-1)
+            if self._num_classes is None:
+                max_label = max(max_label, predicted_y.shape[1]-1)
             predicted_y = torch.max(predicted_y, 1)[1]
         else:
             # Labels -> check non-negative
             min_label = torch.min(predicted_y).item()
             if min_label < 0:
                 raise ValueError('Label values must be non-negative values')
+            if self._num_classes is None:
+                max_label = max(max_label, torch.max(predicted_y).item())
+            elif torch.max(predicted_y).item() >= self._num_classes:
+                raise ValueError("Encountered predicted label larger than"
+                                 "num_classes")
 
         if len(true_y.shape) > 1:
             # Logits -> transform to labels
-            max_label = max(max_label, true_y.shape[1]-1)
+            if self._num_classes is None:
+                max_label = max(max_label, true_y.shape[1]-1)
             true_y = torch.max(true_y, 1)[1]
         else:
             # Labels -> check non-negative
@@ -113,11 +142,11 @@ class ConfusionMatrix(Metric[Tensor]):
             if min_label < 0:
                 raise ValueError('Label values must be non-negative values')
 
-        # Initialize or enlarge the confusion matrix
-        max_label = max(max_label, max(torch.max(predicted_y).item(),
-                                       torch.max(true_y).item()))
-        if self._num_classes is not None:
-            max_label = max(max_label, self._num_classes-1)
+            if self._num_classes is None:
+                max_label = max(max_label, torch.max(true_y).item())
+            elif torch.max(true_y).item() >= self._num_classes:
+                raise ValueError("Encountered target label larger than"
+                                 "num_classes")
 
         if max_label < 0:
             raise ValueError('The Confusion Matrix metric can only handle '
@@ -151,7 +180,7 @@ class ConfusionMatrix(Metric[Tensor]):
             return torch.zeros(matrix_shape, dtype=torch.long)
         if self.normalize is not None:
             return ConfusionMatrix._normalize_cm(self._cm_tensor,
-                                                 self._normalize)
+                                                 self.normalize)
         return self._cm_tensor
 
     def reset(self) -> None:
@@ -197,6 +226,10 @@ class StreamConfusionMatrix(PluginMetric[Tensor]):
     The Stream Confusion Matrix metric.
     This plugin metric only works on the eval phase.
 
+    Confusion Matrix computation can be slow if you compute it for a large
+    number of classes. We recommend to set `save_image=False` if the runtime
+    is too large.
+
     At the end of the eval phase, this metric logs the confusion matrix
     relative to all the patterns seen during eval.
 
@@ -204,7 +237,7 @@ class StreamConfusionMatrix(PluginMetric[Tensor]):
     confusion matrix.
     """
 
-    def __init__(self, *,
+    def __init__(self,
                  num_classes: Union[int, Mapping[int, int]] = None,
                  normalize: Literal['true', 'pred', 'all'] = None,
                  save_image: bool = True,
@@ -213,13 +246,18 @@ class StreamConfusionMatrix(PluginMetric[Tensor]):
         """
         Creates an instance of the Stream Confusion Matrix metric.
 
-        :param num_classes: When not None, is used to properly define the
-            amount of rows/columns in the confusion matrix. When None, the
-            matrix will have many rows/columns as the maximum value of the
-            predicted and true pattern labels. Can be either an int, in which
-            case the same value will be used across all experiences, or a
-            dictionary defining the amount of classes for each experience (key =
-             experience label, value = amount of classes). Defaults to None.
+        We recommend to set `save_image=False` if the runtime is too large.
+        In fact, a large number of classes may increase the computation time
+        of this metric.
+
+        :param num_classes: The number of classes. Defaults to None,
+            which means that the number of classes will be inferred from
+            ground truth and prediction Tensors (see class description for more
+            details). If not None, the confusion matrix will always be of size
+            `num_classes, num_classes` and only the first `num_classes` values
+            of output logits or target logits will be considered in the update.
+            If the output or targets are provided as numerical labels,
+            there can be no label greater than `num_classes`.
         :param normalize: Normalizes confusion matrix over the true (rows),
             predicted (columns) conditions or all the population. If None,
             confusion matrix will not be normalized. Valid values are: 'true',
@@ -234,16 +272,18 @@ class StreamConfusionMatrix(PluginMetric[Tensor]):
         super().__init__()
 
         self._save_image: bool = save_image
-        self._matrix: ConfusionMatrix = ConfusionMatrix(num_classes, normalize)
-        self._num_classes: Optional[Union[int, Mapping[int, int]]] = num_classes
-        self._normalize: Optional[Literal['true', 'pred', 'all']] = normalize
+        self.num_classes = num_classes
+        self.normalize = normalize
+        self._matrix: ConfusionMatrix = ConfusionMatrix(num_classes=num_classes,
+                                                        normalize=normalize)
 
         if image_creator is None:
             image_creator = default_cm_image_creator
         self._image_creator: Callable[[Tensor], Image] = image_creator
 
     def reset(self) -> None:
-        self._matrix = ConfusionMatrix()
+        self._matrix = ConfusionMatrix(num_classes=self.num_classes,
+                                       normalize=self.normalize)
 
     def result(self) -> Tensor:
         exp_cm = self._matrix.result()
diff --git a/avalanche/logging/tensorboard_logger.py b/avalanche/logging/tensorboard_logger.py
index d4062d5f..c1fed732 100644
--- a/avalanche/logging/tensorboard_logger.py
+++ b/avalanche/logging/tensorboard_logger.py
@@ -15,9 +15,10 @@ Avalanche experiments. """
 from pathlib import Path
 
 from PIL.Image import Image
+from torch import Tensor
 from torch.utils.tensorboard import SummaryWriter
+from matplotlib.pyplot import Figure
 from torchvision.transforms.functional import to_tensor
-
 from avalanche.evaluation.metric_results import AlternativeValues, MetricValue
 from avalanche.logging import StrategyLogger
 
@@ -61,17 +62,28 @@ class TensorboardLogger(StrategyLogger):
         value = metric_value.value
 
         if isinstance(value, AlternativeValues):
-            value = value.best_supported_value(Image, float, int)
+            value = value.best_supported_value(Image, Tensor,
+                                               Figure, float, int)
 
-        if not isinstance(value, (Image, float, int)):
+        if not isinstance(value, (Image, Tensor, Figure, float, int)):
             # Unsupported type
             return
 
-        if isinstance(value, Image):
+        if isinstance(value, Figure):
+            self.writer.add_figure(name, value,
+                                   global_step=metric_value.x_plot)
+
+        elif isinstance(value, Image):
             self.writer.add_image(name, to_tensor(value),
                                   global_step=metric_value.x_plot)
+
+        elif isinstance(value, Tensor):
+            self.writer.add_histogram(name, value,
+                                      global_step=metric_value.x_plot)
+
         elif isinstance(value, (float, int)):
-            self.writer.add_scalar(name, value, global_step=metric_value.x_plot)
+            self.writer.add_scalar(name, value,
+                                   global_step=metric_value.x_plot)
 
 
 __all__ = [
diff --git a/avalanche/logging/wandb_logger.py b/avalanche/logging/wandb_logger.py
index 98bb18a4..f01674c8 100644
--- a/avalanche/logging/wandb_logger.py
+++ b/avalanche/logging/wandb_logger.py
@@ -11,11 +11,14 @@
 
 """ This module handles all the functionalities related to the logging of
 Avalanche experiments using Weights & Biases. """
-
+import wandb
 from PIL.Image import Image
+from torch import Tensor
+from matplotlib.pyplot import Figure
 
 from avalanche.evaluation.metric_results import AlternativeValues, MetricValue
 from avalanche.logging import StrategyLogger
+import numpy as np
 
 
 class WandBLogger(StrategyLogger):
@@ -60,22 +63,21 @@ class WandBLogger(StrategyLogger):
         value = metric_value.value
 
         if isinstance(value, AlternativeValues):
-            value = value.best_supported_value(Image, float, int)
+            value = value.best_supported_value(Image, Tensor,
+                                               Figure, float, int)
 
-        if not isinstance(value, (Image, float, int)):
+        if not isinstance(value, (Image, Tensor, Figure, float, int)):
             # Unsupported type
             return
 
         if isinstance(value, Image):
-            # Confusion Matrix logging
-            if not self.interactive:
-                self.wandb.log({name: [self.wandb.Image(
-                    value, caption="Confusion Matrix")]})
-            else:
-                # WIP
-                return
-
-        elif isinstance(value, (float, int)):
+            self.wandb.log({name: self.wandb.Image(value)})
+
+        elif isinstance(value, Tensor):
+            value = np.histogram(value.view(-1).numpy())
+            self.wandb.log({name: wandb.Histogram(np_histogram=value)})
+
+        elif isinstance(value, (float, int, Figure)):
             self.wandb.log({name: value})
 
 
diff --git a/avalanche/training/plugins/evaluation.py b/avalanche/training/plugins/evaluation.py
index 73b541f2..10b7decb 100644
--- a/avalanche/training/plugins/evaluation.py
+++ b/avalanche/training/plugins/evaluation.py
@@ -1,7 +1,7 @@
 import warnings
+from copy import copy
 from collections import defaultdict
 from typing import Union, Sequence, TYPE_CHECKING
-from copy import deepcopy
 from avalanche.training.plugins.strategy_plugin import StrategyPlugin
 
 if TYPE_CHECKING:
@@ -96,18 +96,18 @@ class EvaluationPlugin(StrategyPlugin):
 
     def get_last_metrics(self):
         """
-        Return dictionary with metric names as keys and last metrics
-        value as values.
+        Return a shallow copy of dictionary with metric names
+        as keys and last metrics value as values.
 
         :return: a dictionary with full metric
             names as keys and last metric value as value.
         """
-        return deepcopy(self.last_metric_results)
+        return copy(self.last_metric_results)
 
     def get_all_metrics(self):
         """
-        Return all collected metrics. This method should be called
-        only when `collect_all` is set to True.
+        Return the dictionary of all collected metrics.
+        This method should be called only when `collect_all` is set to True.
 
         :return: if `collect_all` is True, returns a dictionary
             with full metric names as keys and a tuple of two lists
@@ -118,7 +118,7 @@ class EvaluationPlugin(StrategyPlugin):
             is False return an empty dictionary
         """
         if self.collect_all:
-            return deepcopy(self.all_metric_results)
+            return self.all_metric_results
         else:
             return {}