How to compute AUC for one class in multi-class classification in Keras?

Viewed 2060

I'm training a CNN on a multi-class problem and want to use AUC computed only for one of the classes as a metric. The last layer is a softmax layer with num_classes neurons, labels are sparsely encoded, and data are batched. I believe I need to write a custom function and then pass it as a metric. What would be the best way to do it in Keras? Can I somehow use the built-in function(s)?

Is there a way to train a model, save it and evaluate this metric for each class afterwards? (through evaluate()?) Or does evaluate only compute metrics specified when compiling?

2 Answers

I had to implement a multiclass AUC metric for a project a while ago, and I'm sharing that code here. I was using tf.keras instead of standalone Keras, but you can easily adapt it to use standalone Keras if you really need to. Also, it's only been tested using TensorFlow 2 and later.

Here it is:

import tensorflow as tf


class MulticlassAUC(tf.keras.metrics.AUC):
    """AUC for a single class in a muliticlass problem.

    Parameters
    ----------
    pos_label : int
        Label of the positive class (the one whose AUC is being computed).

    from_logits : bool, optional (default: False)
        If True, assume predictions are not standardized to be between 0 and 1.
        In this case, predictions will be squeezed into probabilities using the
        softmax function.

    sparse : bool, optional (default: True)
        If True, ground truth labels should be encoded as integer indices in the
        range [0, n_classes-1]. Otherwise, ground truth labels should be one-hot
        encoded indicator vectors (with a 1 in the true label position and 0
        elsewhere).

    **kwargs : keyword arguments
        Keyword arguments for tf.keras.metrics.AUC.__init__(). For example, the
        curve type (curve='ROC' or curve='PR').
    """

    def __init__(self, pos_label, from_logits=False, sparse=True, **kwargs):
        super().__init__(**kwargs)

        self.pos_label = pos_label
        self.from_logits = from_logits
        self.sparse = sparse

    def update_state(self, y_true, y_pred, **kwargs):
        """Accumulates confusion matrix statistics.

        Parameters
        ----------
        y_true : tf.Tensor
            The ground truth values. Either an integer tensor of shape
            (n_examples,) (if sparse=True) or a one-hot tensor of shape
            (n_examples, n_classes) (if sparse=False).

        y_pred : tf.Tensor
            The predicted values, a tensor of shape (n_examples, n_classes).

        **kwargs : keyword arguments
            Extra keyword arguments for tf.keras.metrics.AUC.update_state
            (e.g., sample_weight).
        """
        if self.sparse:
            y_true = tf.math.equal(y_true, self.pos_label)
            y_true = tf.squeeze(y_true)
        else:
            y_true = y_true[..., self.pos_label]

        if self.from_logits:
            y_pred = tf.nn.softmax(y_pred, axis=-1)
        y_pred = y_pred[..., self.pos_label]

        super().update_state(y_true, y_pred, **kwargs)

Example usage:

(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()

model = tf.keras.Sequential(layers=[
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(units=512, activation='relu'),
    tf.keras.layers.Dense(units=256, activation='relu'),
    tf.keras.layers.Dense(units=128, activation='relu'),
    tf.keras.layers.Dense(units=10, activation='softmax'),
])

model.compile(
    loss='sparse_categorical_crossentropy',
    metrics=[MulticlassAUC(pos_label=0)],  # Track AUC for class 0 only
)

model.fit(x_train, y_train, epochs=10)

One alternative is to build a custom function which takes as input the predictions matrix with probabilities and the labels matrix which has the one hot encoded predictions. Then you can use sklearn's roc_auc_score from sklearn.metrics

def compute_aucs(predictions, labels):
    num_classes = predictions.shape[1]
    aucs = np.zeros(num_classes)
    for i in range(num_classes):
        aucs[i] = roc_auc_score(labels[:, i], predictions[:, i])
    return aucs
Related