Skip to content

rslearn.train.metrics

metrics

Metric output classes for non-scalar metrics.

NonScalarMetricOutput dataclass

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

Subclasses should implement the platform-specific methods for the loggers they support. Unsupported loggers are skipped with a warning.

Source code in rslearn/train/metrics.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
@dataclass
class NonScalarMetricOutput:
    """Base class for non-scalar metric outputs that need special logging.

    Subclasses should implement the platform-specific methods for the loggers they
    support. Unsupported loggers are skipped with a warning.
    """

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

        Args:
            name: the metric name
        """
        logger.warning("W&B logging is not implemented for metric %s", name)

    def log_to_mlflow(self, name: str, client: Any, run_id: str) -> None:
        """Log this metric to MLflow.

        Args:
            name: the metric name.
            client: the MLflow client associated with the Lightning logger.
            run_id: the MLflow run ID.
        """
        logger.warning("MLflow logging is not implemented for metric %s", name)

log_to_wandb 

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 
24
25
26
27
28
29
30
def log_to_wandb(self, name: str) -> None:
    """Log this metric to wandb.

    Args:
        name: the metric name
    """
    logger.warning("W&B logging is not implemented for metric %s", name)

log_to_mlflow 

log_to_mlflow(name: str, client: Any, run_id: str) -> None

Log this metric to MLflow.

Parameters:

Name Type Description Default
name str

the metric name.

 required
client Any

the MLflow client associated with the Lightning logger.

 required
run_id str

the MLflow run ID.

 required
Source code in rslearn/train/metrics.py 
32
33
34
35
36
37
38
39
40
def log_to_mlflow(self, name: str, client: Any, run_id: str) -> None:
    """Log this metric to MLflow.

    Args:
        name: the metric name.
        client: the MLflow client associated with the Lightning logger.
        run_id: the MLflow run ID.
    """
    logger.warning("MLflow logging is not implemented for metric %s", name)

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 
 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
102
103
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
@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,
                ),
            },
        )

    def log_to_mlflow(self, name: str, client: Any, run_id: str) -> None:
        """Log the confusion matrix as MLflow table and figure artifacts.

        Args:
            name: the metric name (e.g., ``val_confusion_matrix``).
            client: the MLflow client associated with the Lightning logger.
            run_id: the MLflow run ID.
        """
        cm = self.confusion_matrix.detach().cpu()
        if cm.dim() > 2:
            cm = cm.sum(dim=0)

        num_classes = cm.shape[0]
        class_names = self.class_names or [str(i) for i in range(num_classes)]
        data: dict[str, list[Any]] = {
            "true_class": [],
            "predicted_class": [],
            "count": [],
        }
        for true_label in range(cm.shape[0]):
            for pred_label in range(cm.shape[1]):
                data["true_class"].append(class_names[true_label])
                data["predicted_class"].append(class_names[pred_label])
                data["count"].append(int(cm[true_label, pred_label].item()))

        client.log_table(
            run_id=run_id,
            data=data,
            artifact_file=f"{name}.json",
        )

        # TorchMetrics provides the confusion-matrix visualization; MLflow stores the
        # resulting Matplotlib figure as an artifact.
        plotter = MulticlassConfusionMatrix(num_classes=num_classes)
        figure, axes = plotter.plot(val=cm, labels=class_names)
        # Import lazily because Matplotlib is only required by the optional MLflow path.
        from matplotlib import pyplot as plt

        try:
            axes.set_title(name)
            client.log_figure(
                run_id=run_id,
                figure=figure,
                artifact_file=f"{name}.png",
            )
        finally:
            plt.close(figure)

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 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
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_mlflow 

log_to_mlflow(name: str, client: Any, run_id: str) -> None

Log the confusion matrix as MLflow table and figure artifacts.

Parameters:

Name Type Description Default
name str

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

 required
client Any

the MLflow client associated with the Lightning logger.

 required
run_id str

the MLflow run ID.

 required
Source code in rslearn/train/metrics.py 
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
def log_to_mlflow(self, name: str, client: Any, run_id: str) -> None:
    """Log the confusion matrix as MLflow table and figure artifacts.

    Args:
        name: the metric name (e.g., ``val_confusion_matrix``).
        client: the MLflow client associated with the Lightning logger.
        run_id: the MLflow run ID.
    """
    cm = self.confusion_matrix.detach().cpu()
    if cm.dim() > 2:
        cm = cm.sum(dim=0)

    num_classes = cm.shape[0]
    class_names = self.class_names or [str(i) for i in range(num_classes)]
    data: dict[str, list[Any]] = {
        "true_class": [],
        "predicted_class": [],
        "count": [],
    }
    for true_label in range(cm.shape[0]):
        for pred_label in range(cm.shape[1]):
            data["true_class"].append(class_names[true_label])
            data["predicted_class"].append(class_names[pred_label])
            data["count"].append(int(cm[true_label, pred_label].item()))

    client.log_table(
        run_id=run_id,
        data=data,
        artifact_file=f"{name}.json",
    )

    # TorchMetrics provides the confusion-matrix visualization; MLflow stores the
    # resulting Matplotlib figure as an artifact.
    plotter = MulticlassConfusionMatrix(num_classes=num_classes)
    figure, axes = plotter.plot(val=cm, labels=class_names)
    # Import lazily because Matplotlib is only required by the optional MLflow path.
    from matplotlib import pyplot as plt

    try:
        axes.set_title(name)
        client.log_figure(
            run_id=run_id,
            figure=figure,
            artifact_file=f"{name}.png",
        )
    finally:
        plt.close(figure)

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 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
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 
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
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 
211
212
213
214
215
216
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 
218
219
220
def reset(self) -> None:
    """Reset metric."""
    super().reset()