title semisupervised image classification using contrastive pretraining with simclr andrs breshttps www linkedin cominandrasberes789190210 date created 20210424 last modified 20210424 description contrastive pretraining with simclr for semisupervised image classification on the stl10 dataset accelerator gpu introduction semisupervised learning semisupervised learning is a machine learning paradigm that deals with partially labeled datasets when applying deep learning in the real world one usually has to gather a large dataset to make it work well however while the cost of labeling scales linearly with the dataset size labeling each example takes a constant time model performance only scales sublinearlyhttps arxiv orgabs2001 08361 with it this means that labeling more and more samples becomes less and less costefficient while gathering unlabeled data is generally cheap as it is usually readily available in large quantities semisupervised learning offers to solve this problem by only requiring a partially labeled dataset and by being labelefficient by utilizing the unlabeled examples for learning as well in this example we will pretrain an encoder with contrastive learning on the stl10https ai stanford eduacoatesstl10 semisupervised dataset using no labels at all and then finetune it using only its labeled subset contrastive learning on the highest level the main idea behind contrastive learning is to learn representations that are invariant to image augmentations in a selfsupervised manner one problem with this objective is that it has a trivial degenerate solution the case where the representations are constant and do not depend at all on the input images contrastive learning avoids this trap by modifying the objective in the following way it pulls representations of augmented versionsviews of the same image closer to each other contracting positives while simultaneously pushing different images away from each other contrasting negatives in representation space one such contrastive approach is simclrhttps arxiv orgabs2002 05709 which essentially identifies the core components needed to optimize this objective and can achieve high performance by scaling this simple approach another approach is simsiamhttps arxiv orgabs2011 10566 keras examplehttps keras ioexamplesvisionsimsiam whose main difference from simclr is that the former does not use any negatives in its loss therefore it does not explicitly prevent the trivial solution and instead avoids it implicitly by architecture design asymmetric encoding paths using a predictor network and batch normalization batchnorm are applied in the final layers for further reading about simclr check out the official google ai blog posthttps ai googleblog com202004advancingselfsupervisedandsemi html and for an overview of selfsupervised learning across both vision and language check out this blog posthttps ai facebook comblogselfsupervisedlearningthedarkmatterofintelligence setup make sure we are able to handle large datasets hyperparameter setup dataset hyperparameters algorithm hyperparameters stronger augmentations for contrastive weaker ones for supervised training dataset during training we will simultaneously load a large batch of unlabeled images along with a smaller batch of labeled images labeled and unlabeled samples are loaded synchronously with batch sizes selected accordingly turning off shuffle to lower resource usage labeled and unlabeled datasets are zipped together load stl10 dataset image augmentations the two most important image augmentations for contrastive learning are the following cropping forces the model to encode different parts of the same image similarly we implement it with the randomtranslationhttps keras ioapilayerspreprocessinglayersimagepreprocessingrandomtranslation and randomzoomhttps keras ioapilayerspreprocessinglayersimagepreprocessingrandomzoom layers color jitter prevents a trivial color histogrambased solution to the task by distorting color histograms a principled way to implement that is by affine transformations in color space in this example we use random horizontal flips as well stronger augmentations are applied for contrastive learning along with weaker ones for supervised classification to avoid overfitting on the few labeled examples we implement random color jitter as a custom preprocessing layer using preprocessing layers for data augmentation has the following two advantages the data augmentation will run on gpu in batches so the training will not be bottlenecked by the data pipeline in environments with constrained cpu resources such as a colab notebook or a personal machine deployment is easier as the data preprocessing pipeline is encapsulated in the model and does not have to be reimplemented when deploying it distorts the color distibutions of images same for all colors different for all colors image augmentation module sample a batch from a dataset apply augmentations encoder architecture define the encoder architecture supervised baseline model a baseline supervised model is trained using random initialization baseline supervised training with random initialization selfsupervised model for contrastive pretraining we pretrain an encoder on unlabeled images with a contrastive loss a nonlinear projection head is attached to the top of the encoder as it improves the quality of representations of the encoder we use the infoncentxentnpairs loss which can be interpreted in the following way 1 we treat each image in the batch as if it had its own class 2 then we have two examples a pair of augmented views for each class 3 each view s representation is compared to every possible pair s one for both augmented versions 4 we use the temperaturescaled cosine similarity of compared representations as logits 5 finally we use categorical crossentropy as the classification loss the following two metrics are used for monitoring the pretraining performance contrastive accuracy simclr table 5https arxiv orgabs2002 05709 selfsupervised metric the ratio of cases in which the representation of an image is more similar to its differently augmented version s one than to the representation of any other image in the current batch selfsupervised metrics can be used for hyperparameter tuning even in the case when there are no labeled examples linear probing accuracyhttps arxiv orgabs1603 08511 linear probing is a popular metric to evaluate selfsupervised classifiers it is computed as the accuracy of a logistic regression classifier trained on top of the encoder s features in our case this is done by training a single dense layer on top of the frozen encoder note that contrary to traditional approach where the classifier is trained after the pretraining phase in this example we train it during pretraining this might slightly decrease its accuracy but that way we can monitor its value during training which helps with experimentation and debugging another widely used supervised metric is the knn accuracyhttps arxiv orgabs1805 01978 which is the accuracy of a knn classifier trained on top of the encoder s features which is not implemented in this example define the contrastive model with modelsubclassing nonlinear mlp as projection head single dense layer for linear probing self contrastiveloss will be defined as a method infonce loss information noisecontrastive estimation ntxent loss normalized temperaturescaled cross entropy cosine similarity the dot product of the l2normalized feature vectors the similarity between the representations of two augmented views of the same image should be higher than their similarity with other views the temperaturescaled similarities are used as logits for crossentropy a symmetrized version of the loss is used here both labeled and unlabeled images are used without labels each image is augmented twice differently the representations are passed through a projection mlp labels are only used in evalutation for an onthefly logistic regression the encoder is used in inference mode here to avoid regularization and updating the batch normalization paramers if they are used for testing the components are used with a trainingfalse flag only the probe metrics are logged at test time contrastive pretraining supervised finetuning of the pretrained encoder we then finetune the encoder on the labeled examples by attaching a single randomly initalized fully connected classification layer on its top supervised finetuning of the pretrained encoder comparison against the baseline the classification accuracies of the baseline and the pretraining finetuning process by comparing the training curves we can see that when using contrastive pretraining a higher validation accuracy can be reached paired with a lower validation loss which means that the pretrained network was able to generalize better when seeing only a small amount of labeled examples improving further architecture the experiment in the original paper demonstrated that increasing the width and depth of the models improves performance at a higher rate than for supervised learning also using a resnet50https keras ioapiapplicationsresnetresnet50function encoder is quite standard in the literature however keep in mind that more powerful models will not only increase training time but will also require more memory and will limit the maximal batch size you can use it has beenhttps arxiv orgabs1905 09272 reportedhttps arxiv orgabs1911 05722 that the usage of batchnorm layers could sometimes degrade performance as it introduces an intrabatch dependency between samples which is why i did not have used them in this example in my experiments however using batchnorm especially in the projection head improves performance hyperparameters the hyperparameters used in this example have been tuned manually for this task and architecture therefore without changing them only marginal gains can be expected from further hyperparameter tuning however for a different task or model architecture these would need tuning so here are my notes on the most important ones batch size since the objective can be interpreted as a classification over a batch of images loosely speaking the batch size is actually a more important hyperparameter than usual the higher the better temperature the temperature defines the softness of the softmax distribution that is used in the crossentropy loss and is an important hyperparameter lower values generally lead to a higher contrastive accuracy a recent trick in alignhttps arxiv orgabs2102 05918 is to learn the temperature s value as well which can be done by defining it as a tf variable and applying gradients on it even though this provides a good baseline value in my experiments the learned temperature was somewhat lower than optimal as it is optimized with respect to the contrastive loss which is not a perfect proxy for representation quality image augmentation strength during pretraining stronger augmentations increase the difficulty of the task however after a point too strong augmentations will degrade performance during finetuning stronger augmentations reduce overfitting while in my experience too strong augmentations decrease the performance gains from pretraining the whole data augmentation pipeline can be seen as an important hyperparameter of the algorithm implementations of other custom image augmentation layers in keras can be found in this repositoryhttps github comberesandrasimageaugmentationlayerskeras learning rate schedule a constant schedule is used here but it is quite common in the literature to use a cosine decay schedulehttps www tensorflow orgapidocspythontfkerasexperimentalcosinedecay which can further improve performance optimizer adam is used in this example as it provides good performance with default parameters sgd with momentum requires more tuning however it could slightly increase performance related works other instancelevel imagelevel contrastive learning methods mocohttps arxiv orgabs1911 05722 v2https arxiv orgabs2003 04297 v3https arxiv orgabs2104 02057 uses a momentumencoder as well whose weights are an exponential moving average of the target encoder swavhttps arxiv orgabs2006 09882 uses clustering instead of pairwise comparison barlowtwinshttps arxiv orgabs2103 03230 uses a cross correlationbased objective instead of pairwise comparison keras implementations of moco and barlowtwins can be found in this repositoryhttps github comberesandrascontrastiveclassificationkeras which includes a colab notebook there is also a new line of works which optimize a similar objective but without the use of any negatives byolhttps arxiv orgabs2006 07733 momentumencoder no negatives simsiamhttps arxiv orgabs2011 10566 keras examplehttps keras ioexamplesvisionsimsiam no momentumencoder no negatives in my experience these methods are more brittle they can collapse to a constant representation i could not get them to work using this encoder architecture even though they are generally more dependent on the modelhttps generallyintelligent aiunderstandingselfsupervisedcontrastivelearning html architecturehttps arxiv orgabs2010 10241 they can improve performance at smaller batch sizes you can use the trained model hosted on hugging face hubhttps huggingface cokerasiosemisupervisedclassificationsimclr and try the demo on hugging face spaceshttps huggingface cospaceskerasiosemisupervisedclassification introduction semi supervised learning semi supervised learning is a machine learning paradigm that deals with partially labeled datasets when applying deep learning in the real world one usually has to gather a large dataset to make it work well however while the cost of labeling scales linearly with the dataset size labeling each example takes a constant time model performance only scales sublinearly https arxiv org abs 2001 08361 with it this means that labeling more and more samples becomes less and less cost efficient while gathering unlabeled data is generally cheap as it is usually readily available in large quantities semi supervised learning offers to solve this problem by only requiring a partially labeled dataset and by being label efficient by utilizing the unlabeled examples for learning as well in this example we will pretrain an encoder with contrastive learning on the stl 10 https ai stanford edu acoates stl10 semi supervised dataset using no labels at all and then fine tune it using only its labeled subset contrastive learning on the highest level the main idea behind contrastive learning is to learn representations that are invariant to image augmentations in a self supervised manner one problem with this objective is that it has a trivial degenerate solution the case where the representations are constant and do not depend at all on the input images contrastive learning avoids this trap by modifying the objective in the following way it pulls representations of augmented versions views of the same image closer to each other contracting positives while simultaneously pushing different images away from each other contrasting negatives in representation space one such contrastive approach is simclr https arxiv org abs 2002 05709 which essentially identifies the core components needed to optimize this objective and can achieve high performance by scaling this simple approach another approach is simsiam https arxiv org abs 2011 10566 keras example https keras io examples vision simsiam whose main difference from simclr is that the former does not use any negatives in its loss therefore it does not explicitly prevent the trivial solution and instead avoids it implicitly by architecture design asymmetric encoding paths using a predictor network and batch normalization batchnorm are applied in the final layers for further reading about simclr check out the official google ai blog post https ai googleblog com 2020 04 advancing self supervised and semi html and for an overview of self supervised learning across both vision and language check out this blog post https ai facebook com blog self supervised learning the dark matter of intelligence setup make sure we are able to handle large datasets hyperparameter setup dataset hyperparameters algorithm hyperparameters corresponds to 200 steps per epoch stronger augmentations for contrastive weaker ones for supervised training dataset during training we will simultaneously load a large batch of unlabeled images along with a smaller batch of labeled images labeled and unlabeled samples are loaded synchronously with batch sizes selected accordingly turning off shuffle to lower resource usage labeled and unlabeled datasets are zipped together load stl10 dataset image augmentations the two most important image augmentations for contrastive learning are the following cropping forces the model to encode different parts of the same image similarly we implement it with the randomtranslation https keras io api layers preprocessing_layers image_preprocessing random_translation and randomzoom https keras io api layers preprocessing_layers image_preprocessing random_zoom layers color jitter prevents a trivial color histogram based solution to the task by distorting color histograms a principled way to implement that is by affine transformations in color space in this example we use random horizontal flips as well stronger augmentations are applied for contrastive learning along with weaker ones for supervised classification to avoid overfitting on the few labeled examples we implement random color jitter as a custom preprocessing layer using preprocessing layers for data augmentation has the following two advantages the data augmentation will run on gpu in batches so the training will not be bottlenecked by the data pipeline in environments with constrained cpu resources such as a colab notebook or a personal machine deployment is easier as the data preprocessing pipeline is encapsulated in the model and does not have to be reimplemented when deploying it distorts the color distibutions of images same for all colors different for all colors image augmentation module sample a batch from a dataset apply augmentations encoder architecture define the encoder architecture supervised baseline model a baseline supervised model is trained using random initialization baseline supervised training with random initialization self supervised model for contrastive pretraining we pretrain an encoder on unlabeled images with a contrastive loss a nonlinear projection head is attached to the top of the encoder as it improves the quality of representations of the encoder we use the infonce nt xent n pairs loss which can be interpreted in the following way 1 we treat each image in the batch as if it had its own class 2 then we have two examples a pair of augmented views for each class 3 each view s representation is compared to every possible pair s one for both augmented versions 4 we use the temperature scaled cosine similarity of compared representations as logits 5 finally we use categorical cross entropy as the classification loss the following two metrics are used for monitoring the pretraining performance contrastive accuracy simclr table 5 https arxiv org abs 2002 05709 self supervised metric the ratio of cases in which the representation of an image is more similar to its differently augmented version s one than to the representation of any other image in the current batch self supervised metrics can be used for hyperparameter tuning even in the case when there are no labeled examples linear probing accuracy https arxiv org abs 1603 08511 linear probing is a popular metric to evaluate self supervised classifiers it is computed as the accuracy of a logistic regression classifier trained on top of the encoder s features in our case this is done by training a single dense layer on top of the frozen encoder note that contrary to traditional approach where the classifier is trained after the pretraining phase in this example we train it during pretraining this might slightly decrease its accuracy but that way we can monitor its value during training which helps with experimentation and debugging another widely used supervised metric is the knn accuracy https arxiv org abs 1805 01978 which is the accuracy of a knn classifier trained on top of the encoder s features which is not implemented in this example define the contrastive model with model subclassing non linear mlp as projection head single dense layer for linear probing self contrastive_loss will be defined as a method infonce loss information noise contrastive estimation nt xent loss normalized temperature scaled cross entropy cosine similarity the dot product of the l2 normalized feature vectors the similarity between the representations of two augmented views of the same image should be higher than their similarity with other views the temperature scaled similarities are used as logits for cross entropy a symmetrized version of the loss is used here both labeled and unlabeled images are used without labels each image is augmented twice differently the representations are passed through a projection mlp labels are only used in evalutation for an on the fly logistic regression the encoder is used in inference mode here to avoid regularization and updating the batch normalization paramers if they are used for testing the components are used with a training false flag only the probe metrics are logged at test time contrastive pretraining supervised finetuning of the pretrained encoder we then finetune the encoder on the labeled examples by attaching a single randomly initalized fully connected classification layer on its top supervised finetuning of the pretrained encoder comparison against the baseline the classification accuracies of the baseline and the pretraining finetuning process by comparing the training curves we can see that when using contrastive pretraining a higher validation accuracy can be reached paired with a lower validation loss which means that the pretrained network was able to generalize better when seeing only a small amount of labeled examples improving further architecture the experiment in the original paper demonstrated that increasing the width and depth of the models improves performance at a higher rate than for supervised learning also using a resnet 50 https keras io api applications resnet resnet50 function encoder is quite standard in the literature however keep in mind that more powerful models will not only increase training time but will also require more memory and will limit the maximal batch size you can use it has been https arxiv org abs 1905 09272 reported https arxiv org abs 1911 05722 that the usage of batchnorm layers could sometimes degrade performance as it introduces an intra batch dependency between samples which is why i did not have used them in this example in my experiments however using batchnorm especially in the projection head improves performance hyperparameters the hyperparameters used in this example have been tuned manually for this task and architecture therefore without changing them only marginal gains can be expected from further hyperparameter tuning however for a different task or model architecture these would need tuning so here are my notes on the most important ones batch size since the objective can be interpreted as a classification over a batch of images loosely speaking the batch size is actually a more important hyperparameter than usual the higher the better temperature the temperature defines the softness of the softmax distribution that is used in the cross entropy loss and is an important hyperparameter lower values generally lead to a higher contrastive accuracy a recent trick in align https arxiv org abs 2102 05918 is to learn the temperature s value as well which can be done by defining it as a tf variable and applying gradients on it even though this provides a good baseline value in my experiments the learned temperature was somewhat lower than optimal as it is optimized with respect to the contrastive loss which is not a perfect proxy for representation quality image augmentation strength during pretraining stronger augmentations increase the difficulty of the task however after a point too strong augmentations will degrade performance during finetuning stronger augmentations reduce overfitting while in my experience too strong augmentations decrease the performance gains from pretraining the whole data augmentation pipeline can be seen as an important hyperparameter of the algorithm implementations of other custom image augmentation layers in keras can be found in this repository https github com beresandras image augmentation layers keras learning rate schedule a constant schedule is used here but it is quite common in the literature to use a cosine decay schedule https www tensorflow org api_docs python tf keras experimental cosinedecay which can further improve performance optimizer adam is used in this example as it provides good performance with default parameters sgd with momentum requires more tuning however it could slightly increase performance related works other instance level image level contrastive learning methods moco https arxiv org abs 1911 05722 v2 https arxiv org abs 2003 04297 v3 https arxiv org abs 2104 02057 uses a momentum encoder as well whose weights are an exponential moving average of the target encoder swav https arxiv org abs 2006 09882 uses clustering instead of pairwise comparison barlowtwins https arxiv org abs 2103 03230 uses a cross correlation based objective instead of pairwise comparison keras implementations of moco and barlowtwins can be found in this repository https github com beresandras contrastive classification keras which includes a colab notebook there is also a new line of works which optimize a similar objective but without the use of any negatives byol https arxiv org abs 2006 07733 momentum encoder no negatives simsiam https arxiv org abs 2011 10566 keras example https keras io examples vision simsiam no momentum encoder no negatives in my experience these methods are more brittle they can collapse to a constant representation i could not get them to work using this encoder architecture even though they are generally more dependent on the model https generallyintelligent ai understanding self supervised contrastive learning html architecture https arxiv org abs 2010 10241 they can improve performance at smaller batch sizes you can use the trained model hosted on hugging face hub https huggingface co keras io semi supervised classification simclr and try the demo on hugging face spaces https huggingface co spaces keras io semi supervised classification
import resource

low, high = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (high, high))

import math
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_datasets as tfds

import keras
from keras import layers



unlabeled_dataset_size = 100000
labeled_dataset_size = 5000
image_size = 96
image_channels = 3


num_epochs = 1
batch_size = 525  
width = 128
temperature = 0.1

contrastive_augmentation = {"min_area": 0.25, "brightness": 0.6, "jitter": 0.2}
classification_augmentation = {
    "min_area": 0.75,
    "brightness": 0.3,
    "jitter": 0.1,
}




def prepare_dataset():
    
    
    steps_per_epoch = (
        unlabeled_dataset_size + labeled_dataset_size
    ) // batch_size
    unlabeled_batch_size = unlabeled_dataset_size // steps_per_epoch
    labeled_batch_size = labeled_dataset_size // steps_per_epoch
    print(
        f"batch size is {unlabeled_batch_size} (unlabeled) + {labeled_batch_size} (labeled)"
    )

    
    unlabeled_train_dataset = (
        tfds.load(
            "stl10", split="unlabelled", as_supervised=True, shuffle_files=False
        )
        .shuffle(buffer_size=10 * unlabeled_batch_size)
        .batch(unlabeled_batch_size)
    )
    labeled_train_dataset = (
        tfds.load(
            "stl10", split="train", as_supervised=True, shuffle_files=False
        )
        .shuffle(buffer_size=10 * labeled_batch_size)
        .batch(labeled_batch_size)
    )
    test_dataset = (
        tfds.load("stl10", split="test", as_supervised=True)
        .batch(batch_size)
        .prefetch(buffer_size=tf.data.AUTOTUNE)
    )

    
    train_dataset = tf.data.Dataset.zip(
        (unlabeled_train_dataset, labeled_train_dataset)
    ).prefetch(buffer_size=tf.data.AUTOTUNE)

    return train_dataset, labeled_train_dataset, test_dataset



train_dataset, labeled_train_dataset, test_dataset = prepare_dataset()





class RandomColorAffine(layers.Layer):
    def __init__(self, brightness=0, jitter=0, **kwargs):
        super().__init__(**kwargs)

        self.brightness = brightness
        self.jitter = jitter

    def get_config(self):
        config = super().get_config()
        config.update({"brightness": self.brightness, "jitter": self.jitter})
        return config

    def call(self, images, training=True):
        if training:
            batch_size = tf.shape(images)[0]

            
            brightness_scales = 1 + tf.random.uniform(
                (batch_size, 1, 1, 1),
                minval=-self.brightness,
                maxval=self.brightness,
            )
            
            jitter_matrices = tf.random.uniform(
                (batch_size, 1, 3, 3), minval=-self.jitter, maxval=self.jitter
            )

            color_transforms = (
                tf.eye(3, batch_shape=[batch_size, 1]) * brightness_scales
                + jitter_matrices
            )
            images = tf.clip_by_value(tf.matmul(images, color_transforms), 0, 1)
        return images



def get_augmenter(min_area, brightness, jitter):
    zoom_factor = 1.0 - math.sqrt(min_area)
    return keras.Sequential(
        [
            keras.Input(shape=(image_size, image_size, image_channels)),
            layers.Rescaling(1 / 255, dtype="uint8"),
            layers.RandomFlip("horizontal"),
            layers.RandomTranslation(zoom_factor / 2, zoom_factor / 2),
            layers.RandomZoom((-zoom_factor, 0.0), (-zoom_factor, 0.0)),
            RandomColorAffine(brightness, jitter),
        ]
    )


def visualize_augmentations(num_images):
    
    images = next(iter(train_dataset))[0][0][:num_images]

    
    augmented_images = zip(
        images,
        get_augmenter(**classification_augmentation)(images),
        get_augmenter(**contrastive_augmentation)(images),
        get_augmenter(**contrastive_augmentation)(images),
    )
    row_titles = [
        "Original:",
        "Weakly augmented:",
        "Strongly augmented:",
        "Strongly augmented:",
    ]
    plt.figure(figsize=(num_images * 2.2, 4 * 2.2), dpi=100)
    for column, image_row in enumerate(augmented_images):
        for row, image in enumerate(image_row):
            plt.subplot(4, num_images, row * num_images + column + 1)
            plt.imshow(image)
            if column == 0:
                plt.title(row_titles[row], loc="left")
            plt.axis("off")
    plt.tight_layout()


visualize_augmentations(num_images=8)





def get_encoder():
    return keras.Sequential(
        [
            keras.Input(shape=(image_size, image_size, image_channels)),
            layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
            layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
            layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
            layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
            layers.Flatten(),
            layers.Dense(width, activation="relu"),
        ],
        name="encoder",
    )





baseline_model = keras.Sequential(
    [
        keras.Input(shape=(image_size, image_size, image_channels)),
        get_augmenter(**classification_augmentation),
        get_encoder(),
        layers.Dense(10),
    ],
    name="baseline_model",
)
baseline_model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")],
)

baseline_history = baseline_model.fit(
    labeled_train_dataset, epochs=num_epochs, validation_data=test_dataset
)

print(
    "Maximal validation accuracy: {:.2f}%".format(
        max(baseline_history.history["val_acc"]) * 100
    )
)





class ContrastiveModel(keras.Model):
    def __init__(self):
        super().__init__()

        self.temperature = temperature
        self.contrastive_augmenter = get_augmenter(**contrastive_augmentation)
        self.classification_augmenter = get_augmenter(
            **classification_augmentation
        )
        self.encoder = get_encoder()
        
        self.projection_head = keras.Sequential(
            [
                keras.Input(shape=(width,)),
                layers.Dense(width, activation="relu"),
                layers.Dense(width),
            ],
            name="projection_head",
        )
        
        self.linear_probe = keras.Sequential(
            [layers.Input(shape=(width,)), layers.Dense(10)],
            name="linear_probe",
        )

        self.encoder.summary()
        self.projection_head.summary()
        self.linear_probe.summary()

    def compile(self, contrastive_optimizer, probe_optimizer, **kwargs):
        super().compile(**kwargs)

        self.contrastive_optimizer = contrastive_optimizer
        self.probe_optimizer = probe_optimizer

        
        self.probe_loss = keras.losses.SparseCategoricalCrossentropy(
            from_logits=True
        )

        self.contrastive_loss_tracker = keras.metrics.Mean(name="c_loss")
        self.contrastive_accuracy = keras.metrics.SparseCategoricalAccuracy(
            name="c_acc"
        )
        self.probe_loss_tracker = keras.metrics.Mean(name="p_loss")
        self.probe_accuracy = keras.metrics.SparseCategoricalAccuracy(
            name="p_acc"
        )

    @property
    def metrics(self):
        return [
            self.contrastive_loss_tracker,
            self.contrastive_accuracy,
            self.probe_loss_tracker,
            self.probe_accuracy,
        ]

    def contrastive_loss(self, projections_1, projections_2):
        
        

        
        projections_1 = tf.math.l2_normalize(projections_1, axis=1)
        projections_2 = tf.math.l2_normalize(projections_2, axis=1)
        similarities = (
            tf.matmul(projections_1, projections_2, transpose_b=True)
            / self.temperature
        )

        
        
        batch_size = tf.shape(projections_1)[0]
        contrastive_labels = tf.range(batch_size)
        self.contrastive_accuracy.update_state(contrastive_labels, similarities)
        self.contrastive_accuracy.update_state(
            contrastive_labels, tf.transpose(similarities)
        )

        
        
        loss_1_2 = keras.losses.sparse_categorical_crossentropy(
            contrastive_labels, similarities, from_logits=True
        )
        loss_2_1 = keras.losses.sparse_categorical_crossentropy(
            contrastive_labels, tf.transpose(similarities), from_logits=True
        )
        return (loss_1_2 + loss_2_1) / 2

    def train_step(self, data):
        (unlabeled_images, _), (labeled_images, labels) = data

        
        images = tf.concat((unlabeled_images, labeled_images), axis=0)
        
        augmented_images_1 = self.contrastive_augmenter(images, training=True)
        augmented_images_2 = self.contrastive_augmenter(images, training=True)
        with tf.GradientTape() as tape:
            features_1 = self.encoder(augmented_images_1, training=True)
            features_2 = self.encoder(augmented_images_2, training=True)
            
            projections_1 = self.projection_head(features_1, training=True)
            projections_2 = self.projection_head(features_2, training=True)
            contrastive_loss = self.contrastive_loss(
                projections_1, projections_2
            )
        gradients = tape.gradient(
            contrastive_loss,
            self.encoder.trainable_weights
            + self.projection_head.trainable_weights,
        )
        self.contrastive_optimizer.apply_gradients(
            zip(
                gradients,
                self.encoder.trainable_weights
                + self.projection_head.trainable_weights,
            )
        )
        self.contrastive_loss_tracker.update_state(contrastive_loss)

        
        preprocessed_images = self.classification_augmenter(
            labeled_images, training=True
        )
        with tf.GradientTape() as tape:
            
            
            features = self.encoder(preprocessed_images, training=False)
            class_logits = self.linear_probe(features, training=True)
            probe_loss = self.probe_loss(labels, class_logits)
        gradients = tape.gradient(
            probe_loss, self.linear_probe.trainable_weights
        )
        self.probe_optimizer.apply_gradients(
            zip(gradients, self.linear_probe.trainable_weights)
        )
        self.probe_loss_tracker.update_state(probe_loss)
        self.probe_accuracy.update_state(labels, class_logits)

        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        labeled_images, labels = data

        
        preprocessed_images = self.classification_augmenter(
            labeled_images, training=False
        )
        features = self.encoder(preprocessed_images, training=False)
        class_logits = self.linear_probe(features, training=False)
        probe_loss = self.probe_loss(labels, class_logits)
        self.probe_loss_tracker.update_state(probe_loss)
        self.probe_accuracy.update_state(labels, class_logits)

        
        return {m.name: m.result() for m in self.metrics[2:]}



pretraining_model = ContrastiveModel()
pretraining_model.compile(
    contrastive_optimizer=keras.optimizers.Adam(),
    probe_optimizer=keras.optimizers.Adam(),
)

pretraining_history = pretraining_model.fit(
    train_dataset, epochs=num_epochs, validation_data=test_dataset
)
print(
    "Maximal validation accuracy: {:.2f}%".format(
        max(pretraining_history.history["val_p_acc"]) * 100
    )
)




finetuning_model = keras.Sequential(
    [
        layers.Input(shape=(image_size, image_size, image_channels)),
        get_augmenter(**classification_augmentation),
        pretraining_model.encoder,
        layers.Dense(10),
    ],
    name="finetuning_model",
)
finetuning_model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")],
)

finetuning_history = finetuning_model.fit(
    labeled_train_dataset, epochs=num_epochs, validation_data=test_dataset
)
print(
    "Maximal validation accuracy: {:.2f}%".format(
        max(finetuning_history.history["val_acc"]) * 100
    )
)





def plot_training_curves(
    pretraining_history, finetuning_history, baseline_history
):
    for metric_key, metric_name in zip(["acc", "loss"], ["accuracy", "loss"]):
        plt.figure(figsize=(8, 5), dpi=100)
        plt.plot(
            baseline_history.history[f"val_{metric_key}"],
            label="supervised baseline",
        )
        plt.plot(
            pretraining_history.history[f"val_p_{metric_key}"],
            label="self-supervised pretraining",
        )
        plt.plot(
            finetuning_history.history[f"val_{metric_key}"],
            label="supervised finetuning",
        )
        plt.legend()
        plt.title(f"Classification {metric_name} during training")
        plt.xlabel("epochs")
        plt.ylabel(f"validation {metric_name}")


plot_training_curves(pretraining_history, finetuning_history, baseline_history)