Skip to content

Migration of Supervised contrative Learning tutorial to Keras3 #2398

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hertschuh merged 7 commits into from
Sep 1, 2026
Merged

Migration of Supervised contrative Learning tutorial to Keras3 #2398

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
9d8778a
Migration of Supervised contrative Learning tutorial to Keras3
LakshmiKalaKadali Apr 3, 2026
bac31f4
Run the code on Jax backend
LakshmiKalaKadali Apr 16, 2026
742c241
Removed commentline for training statement in baseline classification…
LakshmiKalaKadali Apr 20, 2026
94c2e0b
Add Corresponding files for Supervised Contrastive Learning Tutorial
maitry63 Jul 17, 2026
e21ca5b
add tested files with jax backend
maitry63 Aug 4, 2026
748cbfd
add tested files with jax backend
maitry63 Aug 4, 2026
5d79f58
fix: testcases
maitry63 Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 60 additions & 53 deletions examples/vision/ipynb/supervised-contrastive-learning.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -11,12 +10,11 @@
"\n",
"**Author:** [Khalid Salama](https://www.linkedin.com/in/khalid-salama-24403144/)<br>\n",
"**Date created:** 2020/11/30<br>\n",
"**Last modified:** 2020/11/30<br>\n",
"**Last modified:** 2026/08/19<br>\n",
"**Description:** Using supervised contrastive learning for image classification."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -36,32 +34,27 @@
"representations of images in different classes.\n",
"2. Training a classifier on top of the frozen encoder.\n",
"\n",
"Note that this example requires [TensorFlow Addons](https://www.tensorflow.org/addons), which you can install using the following command:\n",
"\n",
"```python\n",
"pip install tensorflow-addons\n",
"```\n",
"\n",
"## Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
"outputs": [],
"source": [
"import tensorflow as tf\n",
"import tensorflow_addons as tfa\n",
"import numpy as np\n",
"from tensorflow import keras\n",
"from tensorflow.keras import layers"
"import os\n",
"\n",
"os.environ[\"KERAS_BACKEND\"] = \"jax\" # or \"tensorflow\" or \"torch\"\n",
"\n",
"import keras\n",
"from keras import layers\n",
"from keras import ops"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -72,25 +65,24 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
"outputs": [],
"source": [
"num_classes = 10\n",
"num_classes = 50\n",
"input_shape = (32, 32, 3)\n",
"\n",
"# Load the train and test data splits\n",
"(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()\n",
"\n",
"# Display shapes of train and test datasets\n",
"print(f\"x_train shape: {x_train.shape} - y_train shape: {y_train.shape}\")\n",
"print(f\"x_test shape: {x_test.shape} - y_test shape: {y_test.shape}\")"
"print(f\"x_test shape: {x_test.shape} - y_test shape: {y_test.shape}\")\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -101,7 +93,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
Expand All @@ -120,7 +112,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -134,12 +125,13 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
"outputs": [],
"source": [
"\n",
"def create_encoder():\n",
" resnet = keras.applications.ResNet50V2(\n",
" include_top=False, weights=None, input_shape=input_shape, pooling=\"avg\"\n",
Expand All @@ -159,13 +151,12 @@
"batch_size = 265\n",
"hidden_units = 512\n",
"projection_units = 128\n",
"num_epochs = 50\n",
"num_epochs = 10\n",
"dropout_rate = 0.5\n",
"temperature = 0.05"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -179,14 +170,14 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
"outputs": [],
"source": [
"def create_classifier(encoder, trainable=True):\n",
"\n",
"def create_classifier(encoder, trainable=True):\n",
" for layer in encoder.layers:\n",
" layer.trainable = trainable\n",
"\n",
Expand All @@ -203,11 +194,10 @@
" loss=keras.losses.SparseCategoricalCrossentropy(),\n",
" metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
" )\n",
" return model"
" return model\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -222,7 +212,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
Expand All @@ -235,11 +225,10 @@
"history = classifier.fit(x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs)\n",
"\n",
"accuracy = classifier.evaluate(x_test, y_test)[1]\n",
"print(f\"Test accuracy: {round(accuracy * 100, 2)}%\")"
"print(f\"Test accuracy: {round(accuracy * 100, 2)}%\")\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -260,28 +249,47 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
"outputs": [],
"source": [
"\n",
"class SupervisedContrastiveLoss(keras.losses.Loss):\n",
" def __init__(self, temperature=1, name=None):\n",
" super().__init__(name=name)\n",
" def __init__(self, temperature=0.05, **kwargs):\n",
" super().__init__(**kwargs)\n",
" self.temperature = temperature\n",
Comment on lines +260 to 262

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To support proper serialization and deserialization of the custom loss (e.g., when saving and loading the model), it is best practice in Keras 3 to implement the get_config method.

Suggested change
" def __init__(self, temperature=0.05, **kwargs):\n",
" super().__init__(**kwargs)\n",
" self.temperature = temperature\n",
" def __init__(self, temperature=0.05, **kwargs):\n",
" super().__init__(**kwargs)\n",
" self.temperature = temperature\n",
"\n",
" def get_config(self):\n",
" config = super().get_config()\n",
" config.update({\"temperature\": self.temperature})\n",
" return config\n",

"\n",
" def __call__(self, labels, feature_vectors, sample_weight=None):\n",
" # Normalize feature vectors\n",
" feature_vectors_normalized = tf.math.l2_normalize(feature_vectors, axis=1)\n",
" # Compute logits\n",
" logits = tf.divide(\n",
" tf.matmul(\n",
" feature_vectors_normalized, tf.transpose(feature_vectors_normalized)\n",
" ),\n",
" def call(self, labels, feature_vectors):\n",
" feature_vectors = ops.normalize(feature_vectors, axis=1)\n",
"\n",
" logits = ops.divide(\n",
" ops.matmul(feature_vectors, ops.transpose(feature_vectors)),\n",
" self.temperature,\n",
" )\n",
" return tfa.losses.npairs_loss(tf.squeeze(labels), logits)\n",
"\n",
" # Create a mask to find positive pairs (images of same class)\n",
" labels = ops.cast(labels, \"int32\")\n",
" labels = ops.reshape(labels, (-1, 1))\n",
" mask = ops.cast(ops.equal(labels, ops.transpose(labels)), \"float32\")\n",
"\n",
" batch_size = ops.shape(logits)[0]\n",
" logits_mask = 1.0 - ops.eye(batch_size)\n",
Comment on lines +275 to +278

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To ensure full compatibility with mixed precision training (e.g., mixed_float16 or mixed_bfloat16) and to prevent potential runtime dtype mismatch errors in strict backends like PyTorch or JAX, it is highly recommended to use the dynamic dtype of the input logits (i.e., logits.dtype) instead of hardcoding "float32" and default float for ops.eye.

Suggested change
" mask = ops.cast(ops.equal(labels, ops.transpose(labels)), \"float32\")\n",
"\n",
" batch_size = ops.shape(logits)[0]\n",
" logits_mask = 1.0 - ops.eye(batch_size)\n",
" mask = ops.cast(ops.equal(labels, ops.transpose(labels)), logits.dtype)\n",
"\n",
" batch_size = ops.shape(logits)[0]\n",
" logits_mask = 1.0 - ops.eye(batch_size, dtype=logits.dtype)\n",

" mask = mask * logits_mask\n",
"\n",
" logits_max = ops.max(logits, axis=1, keepdims=True)\n",
" logits_exp = ops.exp(logits - logits_max) * logits_mask\n",
"\n",
" log_prob = (logits - logits_max) - ops.log(\n",
" ops.sum(logits_exp, axis=1, keepdims=True) + 1e-8\n",
" )\n",
"\n",
" mean_log_prob_pos = ops.sum(mask * log_prob, axis=1) / (\n",
" ops.sum(mask, axis=1) + 1e-8\n",
" )\n",
"\n",
" return ops.subtract(0.0, ops.mean(mean_log_prob_pos))\n",

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using ops.subtract(0.0, ...) is non-idiomatic and less readable. We can simply use the unary negation operator - on the tensor, which is fully supported across all Keras backends.

Suggested change
" return ops.subtract(0.0, ops.mean(mean_log_prob_pos))\n",
" return -ops.mean(mean_log_prob_pos)\n"

"\n",
"\n",
"def add_projection_head(encoder):\n",
Expand All @@ -291,11 +299,10 @@
" model = keras.Model(\n",
" inputs=inputs, outputs=outputs, name=\"cifar-encoder_with_projection-head\"\n",
" )\n",
" return model"
" return model\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -306,7 +313,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
Expand All @@ -323,12 +330,14 @@
"encoder_with_projection_head.summary()\n",
"\n",
"history = encoder_with_projection_head.fit(\n",
" x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs\n",
" x=x_train,\n",
" y=y_train,\n",
" batch_size=batch_size,\n",
" epochs=num_epochs,\n",
")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -339,7 +348,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 0,
"metadata": {
"colab_type": "code"
},
Expand All @@ -354,7 +363,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -364,7 +372,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"colab_type": "text"
Expand All @@ -381,7 +388,7 @@
"improve its effectiveness. See the [Supervised Contrastive Learning](https://arxiv.org/abs/2004.11362)\n",
"paper for more details.\n",
"\n",
"You can use the trained model hosted on [Hugging Face Hub](https://huggingface.co/keras-io/supervised-contrastive-learning-cifar10) \n",
"You can use the trained model hosted on [Hugging Face Hub](https://huggingface.co/keras-io/supervised-contrastive-learning-cifar10)\n",
"and try the demo on [Hugging Face Spaces](https://huggingface.co/spaces/keras-io/supervised-contrastive-learning)."
]
}
Expand Down
Loading