CohenKappa compatibility with tf.Dataset

[ammarchalifah]

Hello TF Addons Community!

I need to use CohenKappa as the main metric in my research. However I can't use CohenKappa directly because my model is a regression model, so first I need to clip and round the prediction result. I tried to make customized metric with super function.

import tensorflow_addons as tfa
from tensorflow_addons.metrics import CohenKappa
from tensorflow.keras.metrics import Metric
from tensorflow_addons.utils.types import AcceptableDTypes, FloatTensorLike

from typeguard import typechecked
from typing import Optional

from tensorflow.python.ops import math_ops
from tensorflow.python.keras.utils import losses_utils
from tensorflow.python.keras.utils import metrics_utils

from tensorflow.python.autograph.core import ag_ctx
from tensorflow.python.autograph.impl import api as autograph
from tensorflow.python.framework import ops

class CohenKappaMetric(CohenKappa):
    def __init__(
        self,
        num_classes: FloatTensorLike,
        name: str = "cohen_kappa",
        weightage: Optional[str] = None,
        sparse_labels: bool = False,
        regression: bool = False,
        dtype: AcceptableDTypes = None,
      ):
      """Creates a `CohenKappa` instance."""
      super().__init__(num_classes=num_classes, name=name, weightage=weightage, sparse_labels=sparse_labels,
                        regression=regression,dtype=dtype)

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_pred = tf.clip_by_value(y_pred, 0, 4)
        y_pred = tf.math.round(y_pred)

        y_true = math_ops.cast(y_true, self._dtype)
        y_pred = math_ops.cast(y_pred, self._dtype)
        [y_true, y_pred], sample_weight = \
            metrics_utils.ragged_assert_compatible_and_get_flat_values([y_true, y_pred], sample_weight)
        y_pred, y_true = losses_utils.squeeze_or_expand_dimensions(y_pred, y_true)

        y_pred = ops.convert_to_tensor_v2_with_dispatch(y_pred)
        y_true = math_ops.cast(y_true, y_pred.dtype)

        print('cohen kappa true ',y_true)
        print('cohen kappa pred ',y_pred)

        #ag_fn = autograph.tf_convert(self._fn, ag_ctx.control_status_ctx())

        return super().update_state(y_true, y_pred, sample_weight)

Then, I tried to test this implementation by using this simple test, and it works fine!

m = CohenKappaMetric(num_classes=5, weightage='quadratic', sparse_labels=True)
m.reset_states()
for i in range(5):
  y_true = np.random.randint(0, 4, size = (10))
  y_pred = np.random.uniform(0, 4, size = (10))
  m.update_state(y_true , y_pred)
result = m.result()
result.numpy()

that has the output shown below

cohen kappa true  tf.Tensor([3. 3. 2. 1. 1. 0. 0. 3. 0. 2.], shape=(10,), dtype=float32)
cohen kappa pred  tf.Tensor([3. 0. 0. 3. 2. 0. 4. 3. 2. 2.], shape=(10,), dtype=float32)
cohen kappa true  tf.Tensor([2. 1. 2. 2. 2. 0. 0. 3. 2. 0.], shape=(10,), dtype=float32)
cohen kappa pred  tf.Tensor([1. 3. 2. 3. 4. 1. 1. 1. 1. 1.], shape=(10,), dtype=float32)
cohen kappa true  tf.Tensor([2. 2. 3. 3. 0. 1. 2. 0. 0. 1.], shape=(10,), dtype=float32)
cohen kappa pred  tf.Tensor([2. 3. 2. 4. 2. 4. 3. 0. 3. 4.], shape=(10,), dtype=float32)
cohen kappa true  tf.Tensor([2. 1. 2. 3. 3. 0. 3. 1. 3. 2.], shape=(10,), dtype=float32)
cohen kappa pred  tf.Tensor([3. 0. 2. 1. 1. 2. 1. 2. 1. 1.], shape=(10,), dtype=float32)
cohen kappa true  tf.Tensor([0. 3. 3. 3. 2. 2. 3. 0. 2. 3.], shape=(10,), dtype=float32)
cohen kappa pred  tf.Tensor([2. 1. 4. 3. 2. 0. 1. 3. 1. 1.], shape=(10,), dtype=float32)
-0.032239795

Based on this simple test, I initially thought that my implementation is correct and ready to do the job. Then, I use this custom metric on my model training. For performance optimization, I use the tf Dataset object as the data generator. Then, I trained it.

from tensorflow.keras.metrics import MeanAbsoluteError

class MAEMetric(MeanAbsoluteError):
  def __init__(
      self,
      name = 'mean_absolute_error',
      dtype=None
  ):
    super().__init__(name = name, dtype = dtype)

  def update_state(self, y_true, y_pred, sample_weight=None):
    print('mae true ',y_true)
    print('mae pred ',y_pred)
    super().update_state(y_true,y_pred,sample_weight)

efficientnet = model_backbone(include_top=False, weights='imagenet', input_shape=(SIZE,SIZE,3))

from tensorflow.keras.layers.experimental.preprocessing import Rescaling, RandomFlip, RandomRotation, RandomZoom

dummy_model = Sequential([
    Rescaling(1/.255, input_shape = (224, 224, 3)),
    RandomFlip(seed = 2019),
    RandomRotation((-0.5, 0.5), fill_mode = 'constant', seed = 2019),
    RandomZoom(0.1),
    layers.GlobalAveragePooling2D(),
    layers.Dense(10), 
    layers.Dense(1)])

dummy_model.compile(
    loss='mse',
    optimizer=Adam(lr=0.0001),
    metrics = [MAEMetric(), CohenKappaMetric(num_classes=5, weightage='quadratic')]
)

dummy_model.fit(
    train_ds,
    epochs = 9,
    validation_data = val_ds
)

The training's output looks like this

Epoch 1/9
mae true  Tensor("ExpandDims:0", shape=(None, 1), dtype=int64)
mae pred  Tensor("sequential_7/dense_17/BiasAdd:0", shape=(None, 1), dtype=float32)
cohen kappa true  Tensor("Cast_3:0", shape=(None, 1), dtype=float32)
cohen kappa pred  Tensor("Round:0", shape=(None, 1), dtype=float32)
mae true  Tensor("ExpandDims:0", shape=(None, 1), dtype=int64)
mae pred  Tensor("sequential_7/dense_17/BiasAdd:0", shape=(None, 1), dtype=float32)
cohen kappa true  Tensor("Cast_3:0", shape=(None, 1), dtype=float32)
cohen kappa pred  Tensor("Round:0", shape=(None, 1), dtype=float32)
108/110 [============================>.] - ETA: 0s - loss: 65777.1684 - mean_absolute_error: 240.5696 - cohen_kappa: nanmae true  Tensor("ExpandDims:0", shape=(None, 1), dtype=int64)
mae pred  Tensor("sequential_7/dense_17/BiasAdd:0", shape=(None, 1), dtype=float32)
cohen kappa true  Tensor("Cast_3:0", shape=(None, 1), dtype=float32)
cohen kappa pred  Tensor("Round:0", shape=(None, 1), dtype=float32)
110/110 [==============================] - 3s 19ms/step - loss: 65625.9973 - mean_absolute_error: 240.2580 - cohen_kappa: nan - val_loss: 51789.2344 - val_mean_absolute_error: 211.3784 - val_cohen_kappa: nan
Epoch 2/9
110/110 [==============================] - 2s 18ms/step - loss: 44576.4421 - mean_absolute_error: 195.6172 - cohen_kappa: nan - val_loss: 34642.3633 - val_mean_absolute_error: 170.1765 - val_cohen_kappa: nan
Epoch 3/9
110/110 [==============================] - 2s 17ms/step - loss: 29814.1161 - mean_absolute_error: 157.0602 - cohen_kappa: nan - val_loss: 22724.2129 - val_mean_absolute_error: 136.7035 - val_cohen_kappa: nan
Epoch 4/9
110/110 [==============================] - 2s 17ms/step - loss: 19166.4171 - mean_absolute_error: 123.2973 - cohen_kappa: nan - val_loss: 14697.2520 - val_mean_absolute_error: 109.1846 - val_cohen_kappa: nan
Epoch 5/9
110/110 [==============================] - 2s 17ms/step - loss: 12316.2697 - mean_absolute_error: 98.0135 - cohen_kappa: nan - val_loss: 9658.8457 - val_mean_absolute_error: 87.9779 - val_cohen_kappa: nan
Epoch 6/9
110/110 [==============================] - 2s 17ms/step - loss: 8056.7846 - mean_absolute_error: 78.5180 - cohen_kappa: nan - val_loss: 6636.1865 - val_mean_absolute_error: 72.0047 - val_cohen_kappa: nan
Epoch 7/9
110/110 [==============================] - 2s 18ms/step - loss: 5632.7191 - mean_absolute_error: 64.8466 - cohen_kappa: nan - val_loss: 4909.8643 - val_mean_absolute_error: 60.5663 - val_cohen_kappa: nan
Epoch 8/9
110/110 [==============================] - 2s 18ms/step - loss: 4177.0316 - mean_absolute_error: 55.0338 - cohen_kappa: nan - val_loss: 3967.1941 - val_mean_absolute_error: 52.6669 - val_cohen_kappa: nan
Epoch 9/9
110/110 [==============================] - 2s 18ms/step - loss: 3410.0622 - mean_absolute_error: 48.9486 - cohen_kappa: nan - val_loss: 3448.6873 - val_mean_absolute_error: 47.3190 - val_cohen_kappa: nan
<tensorflow.python.keras.callbacks.History at 0x7f001c38b9d0>

Here's the issue: the custom MAE works just fine, however my custom cohen_kappa results in nan for every epoch! I kinda confident that my implementation is correct based on the previous test, however, I think the problem lies in the conversion between tf objects. I find that MAE's implementation involves complex data transformation and currently too advance for my skill. So, if you guys can't add the metrics API to support tf.Dataset, may you help me to find any workaround for this?

Thanks!

Relevant information

• Are you willing to contribute it (yes/no): YES, however my tech skill is really limited. I'm happy to contribute tho
• Are you willing to maintain it going forward? (yes/no): Again, I'm happy to contribute and maintain it but my tech skill is still too low
• Is there already an implementation in another framework? (if so, where): tensorflow metrics

[ammarchalifah]

I've found the solution, i.e. re-implement it with minor change. Here's the practical solution:

class QuadraticCohenKappa(tf.keras.metrics.Metric):
    def __init__(self, num_classes, name="cohen_kappa", **kwargs):
        super(QuadraticCohenKappa, self).__init__(name=name, **kwargs)
        self.num_classes = num_classes
        self.conf_mtx = self.add_weight(
            "conf_mtx",
            shape=(self.num_classes, self.num_classes),
            initializer=tf.keras.initializers.zeros,
            dtype=tf.float32,
        )

    def _safe_squeeze(self, y):
        y = tf.squeeze(y)

        # Check for scalar result
        if tf.rank(y) == 0:
            y = tf.expand_dims(y, 0)

        return y

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_pred = tf.reshape(tf.math.round(tf.clip_by_value(y_pred, 0, 4)), shape=(-1, 1))
        y_pred = tf.cast(y_pred, "int32")
        y_true = tf.cast(y_true, "int32")

        y_pred = self._safe_squeeze(y_pred)
        y_true = self._safe_squeeze(y_true)

        new_conf_mtx = tf.math.confusion_matrix(
            labels=y_true,
            predictions=y_pred,
            num_classes=self.num_classes,
            weights=sample_weight,
            dtype=tf.float32,
        )
        self.conf_mtx.assign_add(new_conf_mtx)

    def result(self):
        nb_ratings = tf.shape(self.conf_mtx)[0]
        weight_mtx = tf.ones([nb_ratings, nb_ratings], dtype=tf.float32)
        weight_mtx += tf.cast(tf.range(nb_ratings), dtype=tf.float32)
        weight_mtx = tf.cast(weight_mtx, dtype=self.dtype)
        weight_mtx = tf.pow((weight_mtx - tf.transpose(weight_mtx)), 2)
        weight_mtx = tf.cast(weight_mtx, dtype=self.dtype)

        # 3. Get counts
        actual_ratings_hist = tf.reduce_sum(self.conf_mtx, axis=1)
        pred_ratings_hist = tf.reduce_sum(self.conf_mtx, axis=0)

        # 4. Get the outer product
        out_prod = pred_ratings_hist[..., None] * actual_ratings_hist[None, ...]

        # 5. Normalize the confusion matrix and outer product
        conf_mtx = self.conf_mtx / tf.reduce_sum(self.conf_mtx)
        out_prod = out_prod / tf.reduce_sum(out_prod)

        conf_mtx = tf.cast(conf_mtx, dtype=self.dtype)
        out_prod = tf.cast(out_prod, dtype=self.dtype)

        # 6. Calculate Kappa score
        numerator = tf.reduce_sum(conf_mtx * weight_mtx)
        denominator = tf.reduce_sum(out_prod * weight_mtx)
        return tf.cond(
            tf.math.is_nan(denominator),
            true_fn=lambda: 0.0,
            false_fn=lambda: 1 - (numerator / denominator),
        )

    def reset_states(self):
        # The state of the metric will be reset at the start of each epoch.
        for v in self.variables:
            K.set_value(
                v,
                np.zeros((self.num_classes, self.num_classes), v.dtype.as_numpy_dtype),
            )

Feel free to point out if there's something wrong with this implementation. Thanks

[AakashKumarNain]

Thanks @ammarchalifah for pointing it out. Can you submit a PR for adding support for tf.data to the existing codebase?

[ammarchalifah]

Sure, I'll work on it.

[ammarchalifah]

Quick update: after I tried to explore the code and trying to look on what to fix, I found out that if we pass True to sparse_labels argument in CohenKappa, then it can be directly used on tf.data API. The point is, my initial issue was invalid.

[seanpmorgan]

TensorFlow Addons is transitioning to a minimal maintenance and release mode. New features will not be added to this repository. For more information, please see our public messaging on this decision: TensorFlow Addons Wind Down

Please consider sending feature requests / contributions to other repositories in the TF community with a similar charters to TFA: Keras Keras-CV Keras-NLP