diff --git a/examples/vision/supervised-contrastive-learning.py b/examples/vision/supervised-contrastive-learning.py index 4803e67147..6b45d56899 100644 --- a/examples/vision/supervised-contrastive-learning.py +++ b/examples/vision/supervised-contrastive-learning.py @@ -20,22 +20,20 @@ that representations of images in the same class will be more similar compared to representations of images in different classes. 2. Training a classifier on top of the frozen encoder. +""" -Note that this example requires [TensorFlow Addons](https://www.tensorflow.org/addons), -which you can install using the following command: - -```python -pip install tensorflow-addons -``` - +""" ## Setup """ +import os + +os.environ["KERAS_BACKEND"] = "tensorflow" +import keras +from keras import ops +from keras import layers import tensorflow as tf -import tensorflow_addons as tfa import numpy as np -from tensorflow import keras -from tensorflow.keras import layers """ ## Prepare the data @@ -159,6 +157,19 @@ def create_classifier(encoder, trainable=True): """ +def npairs_loss(y_true, y_pred): + y_pred = ops.convert_to_tensor(y_pred) + y_true = ops.cast(y_true, dtype="uint32") + + # Expand to [batch_size, 1] + y_true = ops.expand_dims(y_true, -1) + y_true = ops.cast(ops.equal(y_true, ops.transpose(y_true)), dtype="uint32") + y_true /= ops.sum(y_true, 1, keepdims=True) + + loss = tf.nn.softmax_cross_entropy_with_logits(logits=y_pred, labels=y_true) + return ops.mean(loss) + + class SupervisedContrastiveLoss(keras.losses.Loss): def __init__(self, temperature=1, name=None): super().__init__(name=name) @@ -174,7 +185,7 @@ def __call__(self, labels, feature_vectors, sample_weight=None): ), self.temperature, ) - return tfa.losses.npairs_loss(tf.squeeze(labels), logits) + return npairs_loss(ops.squeeze(labels), logits) def add_projection_head(encoder):