Skip to content

rslearn.train.metrics

metrics

Metric output classes for non-scalar metrics.

NonScalarMetricOutput dataclass

Bases: ABC

Base class for non-scalar metric outputs that need special logging.

Subclasses should implement the log_to_wandb method to define how the metric should be logged (only supports logging to Weights & Biases).

Source code in rslearn/train/metrics.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
@dataclass
class NonScalarMetricOutput(ABC):
    """Base class for non-scalar metric outputs that need special logging.

    Subclasses should implement the log_to_wandb method to define how the metric
    should be logged (only supports logging to Weights & Biases).
    """

    @abstractmethod
    def log_to_wandb(self, name: str) -> None:
        """Log this metric to wandb.

        Args:
            name: the metric name
        """
        pass

log_to_wandb abstractmethod 

log_to_wandb(name: str) -> None

Log this metric to wandb.

Parameters:

Name Type Description Default
name str

the metric name

 required
Source code in rslearn/train/metrics.py 
23
24
25
26
27
28
29
30
@abstractmethod
def log_to_wandb(self, name: str) -> None:
    """Log this metric to wandb.

    Args:
        name: the metric name
    """
    pass

ConfusionMatrixOutput dataclass

Bases: NonScalarMetricOutput

Confusion matrix metric output.

Parameters:

Name Type Description Default
confusion_matrix Tensor

confusion matrix of shape (num_classes, num_classes) where cm[i, j] is the count of samples with true label i and predicted label j.

 required
class_names list[str] | None

optional list of class names for axis labels

 None
Source code in rslearn/train/metrics.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@dataclass
class ConfusionMatrixOutput(NonScalarMetricOutput):
    """Confusion matrix metric output.

    Args:
        confusion_matrix: confusion matrix of shape (num_classes, num_classes)
            where cm[i, j] is the count of samples with true label i and predicted
            label j.
        class_names: optional list of class names for axis labels
    """

    confusion_matrix: torch.Tensor
    class_names: list[str] | None = None

    def _expand_confusion_matrix(self) -> tuple[list[int], list[int]]:
        """Expand confusion matrix to (preds, labels) pairs for wandb.

        Returns:
            Tuple of (preds, labels) as lists of integers.
        """
        cm = self.confusion_matrix.detach().cpu()

        # Handle extra dimensions from distributed reduction
        if cm.dim() > 2:
            cm = cm.sum(dim=0)

        total = cm.sum().item()
        if total == 0:
            return [], []

        preds = []
        labels = []
        for true_label in range(cm.shape[0]):
            for pred_label in range(cm.shape[1]):
                count = cm[true_label, pred_label].item()
                if count > 0:
                    preds.extend([pred_label] * int(count))
                    labels.extend([true_label] * int(count))

        return preds, labels

    def log_to_wandb(self, name: str) -> None:
        """Log confusion matrix to wandb.

        Args:
            name: the metric name (e.g., "val_confusion_matrix")
        """
        preds, labels = self._expand_confusion_matrix()

        if len(preds) == 0:
            logger.warning(f"No samples to log for {name}")
            return

        num_classes = self.confusion_matrix.shape[0]
        if self.class_names is None:
            class_names = [str(i) for i in range(num_classes)]
        else:
            class_names = self.class_names

        wandb.log(
            {
                name: wandb.plot.confusion_matrix(
                    preds=preds,
                    y_true=labels,
                    class_names=class_names,
                    title=name,
                ),
            },
        )

log_to_wandb 

log_to_wandb(name: str) -> None

Log confusion matrix to wandb.

Parameters:

Name Type Description Default
name str

the metric name (e.g., "val_confusion_matrix")

 required
Source code in rslearn/train/metrics.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def log_to_wandb(self, name: str) -> None:
    """Log confusion matrix to wandb.

    Args:
        name: the metric name (e.g., "val_confusion_matrix")
    """
    preds, labels = self._expand_confusion_matrix()

    if len(preds) == 0:
        logger.warning(f"No samples to log for {name}")
        return

    num_classes = self.confusion_matrix.shape[0]
    if self.class_names is None:
        class_names = [str(i) for i in range(num_classes)]
    else:
        class_names = self.class_names

    wandb.log(
        {
            name: wandb.plot.confusion_matrix(
                preds=preds,
                y_true=labels,
                class_names=class_names,
                title=name,
            ),
        },
    )

ConfusionMatrixMetric

Bases: Metric

Confusion matrix metric that works on flattened inputs.

Expects preds of shape (N, C) and labels of shape (N,). Should be wrapped by ClassificationMetric or SegmentationMetric which handle the task-specific preprocessing.

Parameters:

Name Type Description Default
num_classes int

number of classes

 required
class_names list[str] | None

optional list of class names for labeling

 None
Source code in rslearn/train/metrics.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
class ConfusionMatrixMetric(Metric):
    """Confusion matrix metric that works on flattened inputs.

    Expects preds of shape (N, C) and labels of shape (N,).
    Should be wrapped by ClassificationMetric or SegmentationMetric
    which handle the task-specific preprocessing.

    Args:
        num_classes: number of classes
        class_names: optional list of class names for labeling
    """

    def __init__(
        self,
        num_classes: int,
        class_names: list[str] | None = None,
    ):
        """Initialize a new ConfusionMatrixMetric.

        Args:
            num_classes: number of classes
            class_names: optional list of class names for labeling
        """
        super().__init__()
        self.num_classes = num_classes
        self.class_names = class_names
        self.add_state(
            "confusion_matrix",
            default=torch.zeros(num_classes, num_classes, dtype=torch.long),
            dist_reduce_fx="sum",
        )

    def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
        """Update metric.

        Args:
            preds: predictions of shape (N, C) - probabilities
            labels: ground truth of shape (N,) - class indices
        """
        if len(preds) == 0:
            return

        pred_classes = preds.argmax(dim=1)  # (N,)

        for true_label in range(self.num_classes):
            for pred_label in range(self.num_classes):
                count = ((labels == true_label) & (pred_classes == pred_label)).sum()
                self.confusion_matrix[true_label, pred_label] += count

    def compute(self) -> ConfusionMatrixOutput:
        """Returns the confusion matrix wrapped in ConfusionMatrixOutput."""
        return ConfusionMatrixOutput(
            confusion_matrix=self.confusion_matrix,
            class_names=self.class_names,
        )

    def reset(self) -> None:
        """Reset metric."""
        super().reset()

update 

update(preds: Tensor, labels: Tensor) -> None

Update metric.

Parameters:

Name Type Description Default
preds Tensor

predictions of shape (N, C) - probabilities

 required
labels Tensor

ground truth of shape (N,) - class indices

 required
Source code in rslearn/train/metrics.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
    """Update metric.

    Args:
        preds: predictions of shape (N, C) - probabilities
        labels: ground truth of shape (N,) - class indices
    """
    if len(preds) == 0:
        return

    pred_classes = preds.argmax(dim=1)  # (N,)

    for true_label in range(self.num_classes):
        for pred_label in range(self.num_classes):
            count = ((labels == true_label) & (pred_classes == pred_label)).sum()
            self.confusion_matrix[true_label, pred_label] += count

compute 

compute() -> ConfusionMatrixOutput

Returns the confusion matrix wrapped in ConfusionMatrixOutput.

Source code in rslearn/train/metrics.py 
153
154
155
156
157
158
def compute(self) -> ConfusionMatrixOutput:
    """Returns the confusion matrix wrapped in ConfusionMatrixOutput."""
    return ConfusionMatrixOutput(
        confusion_matrix=self.confusion_matrix,
        class_names=self.class_names,
    )

reset 

reset() -> None

Reset metric.

Source code in rslearn/train/metrics.py 
160
161
162
def reset(self) -> None:
    """Reset metric."""
    super().reset()