title selfsupervised contrastive learning with simsiam sayak paulhttps twitter comrisingsayak date created 20210319 last modified 20210320 description implementation of a selfsupervised learning method for computer vision accelerator gpu selfsupervised learning ssl is an interesting branch of study in the field of representation learning ssl systems try to formulate a supervised signal from a corpus of unlabeled data points an example is we train a deep neural network to predict the next word from a given set of words in literature these tasks are known as pretext tasks or auxiliary tasks if we train such a networkhttps arxiv orgabs1801 06146 on a huge dataset such as the wikipedia text corpushttps www corpusdata orgwikipedia asp it learns very effective representations that transfer well to downstream tasks language models like berthttps arxiv orgabs1810 04805 gpt3https arxiv orgabs2005 14165 elmohttps allennlp orgelmo all benefit from this much like the language models we can train computer vision models using similar approaches to make things work in computer vision we need to formulate the learning tasks such that the underlying model a deep neural network is able to make sense of the semantic information present in vision data one such task is to a model to contrast between two different versions of the same image the hope is that in this way the model will have learn representations where the similar images are grouped as together possible while the dissimilar images are further away in this example we will be implementing one such system called simsiam proposed in exploring simple siamese representation learninghttps arxiv orgabs2011 10566 it is implemented as the following 1 we create two different versions of the same dataset with a stochastic data augmentation pipeline note that the random initialization seed needs to be the same during create these versions 2 we take a resnet without any classification head backbone and we add a shallow fullyconnected network projection head on top of it collectively this is known as the encoder 3 we pass the output of the encoder through a predictor which is again a shallow fullyconnected network having an autoencoderhttps en wikipedia orgwikiautoencoder like structure 4 we then train our encoder to maximize the cosine similarity between the two different versions of our dataset setup define hyperparameters load the cifar10 dataset defining our data augmentation pipeline as studied in simclrhttps arxiv orgabs2002 05709 having the right data augmentation pipeline is critical for ssl systems to work effectively in computer vision two particular augmentation transforms that seem to matter the most are 1 random resized crops and 2 color distortions most of the other ssl systems for computer vision such as byolhttps arxiv orgabs2006 07733 mocov2https arxiv orgabs2003 04297 swavhttps arxiv orgabs2006 09882 etc include these in their training pipelines with random crops we also apply horizontal flipping affine transformations can disturb the natural range of rgb images hence this is needed as discussed in the simclr paper the series of augmentation transformations except for random crops need to be applied randomly to impose translational invariance it should be noted that an augmentation pipeline is generally dependent on various properties of the dataset we are dealing with for example if images in the dataset are heavily objectcentric then taking random crops with a very high probability may hurt the training performance let s now apply our augmentation pipeline to our dataset and visualize a few outputs convert the data into tensorflow dataset objects here we create two different versions of our dataset without any groundtruth labels we then zip both of these datasets visualize a few augmented images ensure that the different versions of the dataset actually contain identical images notice that the images in samplesimagesone and sampleimagestwo are essentially the same but are augmented differently defining the encoder and the predictor we use an implementation of resnet20 that is specifically configured for the cifar10 dataset the code is taken from the kerasidiomaticprogrammerhttps github comgooglecloudplatformkerasidiomaticprogrammerblobmasterzooresnetresnetcifar10v2 py repository the hyperparameters of these architectures have been referred from section 3 and appendix a of the original paperhttps arxiv orgabs2011 10566 shell wget q https shorturl atqs369 o resnetcifar10v2 py input and backbone projection head note the autoencoderlike structure defining the pretraining loop one of the main reasons behind training networks with these kinds of approaches is to utilize the learned representations for downstream tasks like classification this is why this particular training phase is also referred to as pretraining we start by defining the loss function the s of simsiam emphasize the impact of the stopgradient operator in the paper as it has an important role in the overall optimization negative cosine similarity minimizing this is equivalent to maximizing the similarity we then define our training loop by overriding the trainstep function of the keras model class unpack the data forward pass through the encoder and predictor note that here we are enforcing the network to match the representations of two differently augmented batches of data compute gradients and update the parameters monitor loss pretraining our networks in the interest of this example we will train the model for only 5 epochs in reality this should at least be 100 epochs create a cosine decay learning scheduler create an early stopping callback compile model and start training visualize the training progress of the model if your solution gets very close to 1 minimum value of our loss very quickly with a different dataset and a different backbone architecture that is likely because of representation collapse it is a phenomenon where the encoder yields similar output for all the images in that case additional hyperparameter tuning is required especially in the following areas strength of the color distortions and their probabilities learning rate and its schedule architecture of both the backbone and their projection head evaluating our ssl method the most popularly used method to evaluate a ssl method in computer vision or any other pretraining method as such is to learn a linear classifier on the frozen features of the trained backbone model in this case it is resnet20 and evaluate the classifier on unseen images other methods include finetuninghttps keras ioguidestransferlearning on the source dataset or even a target dataset with 5 or 10 labels present practically we can use the backbone model for any downstream task such as semantic segmentation object detection and so on where the backbone models are usually pretrained with pure supervised learning we first create labeled dataset objects then we shuffle batch and prefetch this dataset for performance we also apply random resized crops as an augmentation but only to the training set extract the backbone resnet20 we then create our linear classifier and train it compile model and start training notes more data and longer pretraining schedule benefit ssl in general ssl is particularly very helpful when you do not have access to very limited labeled training data but you can manage to build a large corpus of unlabeled data recently using an ssl method called swavhttps arxiv orgabs2006 09882 a group of researchers at facebook trained a regnethttps arxiv orgabs2006 09882 on 2 billion images they were able to achieve downstream performance very close to those achieved by pure supervised pretraining for some downstream tasks their method even outperformed the supervised counterparts you can check out their paperhttps arxiv orgpdf2103 01988 pdf to know the details if you are interested to understand why contrastive ssl helps networks learn meaningful representations you can check out the following resources selfsupervised learning the dark matter of intelligencehttps ai facebook comblogselfsupervisedlearningthedarkmatterofintelligence understanding selfsupervised learning using controlled datasets with known structurehttps sslneuips20 github iofilescamerareadys2037764camerareadyunderstandingselfsupervisedlearning pdf self supervised learning ssl is an interesting branch of study in the field of representation learning ssl systems try to formulate a supervised signal from a corpus of unlabeled data points an example is we train a deep neural network to predict the next word from a given set of words in literature these tasks are known as pretext tasks or auxiliary tasks if we train such a network https arxiv org abs 1801 06146 on a huge dataset such as the wikipedia text corpus https www corpusdata org wikipedia asp it learns very effective representations that transfer well to downstream tasks language models like bert https arxiv org abs 1810 04805 gpt 3 https arxiv org abs 2005 14165 elmo https allennlp org elmo all benefit from this much like the language models we can train computer vision models using similar approaches to make things work in computer vision we need to formulate the learning tasks such that the underlying model a deep neural network is able to make sense of the semantic information present in vision data one such task is to a model to _contrast_ between two different versions of the same image the hope is that in this way the model will have learn representations where the similar images are grouped as together possible while the dissimilar images are further away in this example we will be implementing one such system called simsiam proposed in exploring simple siamese representation learning https arxiv org abs 2011 10566 it is implemented as the following 1 we create two different versions of the same dataset with a stochastic data augmentation pipeline note that the random initialization seed needs to be the same during create these versions 2 we take a resnet without any classification head backbone and we add a shallow fully connected network projection head on top of it collectively this is known as the encoder 3 we pass the output of the encoder through a predictor which is again a shallow fully connected network having an autoencoder https en wikipedia org wiki autoencoder like structure 4 we then train our encoder to maximize the cosine similarity between the two different versions of our dataset setup define hyperparameters load the cifar 10 dataset defining our data augmentation pipeline as studied in simclr https arxiv org abs 2002 05709 having the right data augmentation pipeline is critical for ssl systems to work effectively in computer vision two particular augmentation transforms that seem to matter the most are 1 random resized crops and 2 color distortions most of the other ssl systems for computer vision such as byol https arxiv org abs 2006 07733 mocov2 https arxiv org abs 2003 04297 swav https arxiv org abs 2006 09882 etc include these in their training pipelines with random crops we also apply horizontal flipping affine transformations can disturb the natural range of rgb images hence this is needed as discussed in the simclr paper the series of augmentation transformations except for random crops need to be applied randomly to impose translational invariance it should be noted that an augmentation pipeline is generally dependent on various properties of the dataset we are dealing with for example if images in the dataset are heavily object centric then taking random crops with a very high probability may hurt the training performance let s now apply our augmentation pipeline to our dataset and visualize a few outputs convert the data into tensorflow dataset objects here we create two different versions of our dataset without any ground truth labels we then zip both of these datasets visualize a few augmented images ensure that the different versions of the dataset actually contain identical images notice that the images in samples_images_one and sample_images_two are essentially the same but are augmented differently defining the encoder and the predictor we use an implementation of resnet20 that is specifically configured for the cifar10 dataset the code is taken from the keras idiomatic programmer https github com googlecloudplatform keras idiomatic programmer blob master zoo resnet resnet_cifar10_v2 py repository the hyperparameters of these architectures have been referred from section 3 and appendix a of the original paper https arxiv org abs 2011 10566 shell wget q https shorturl at qs369 o resnet_cifar10_v2 py input and backbone projection head note the autoencoder like structure defining the pre training loop one of the main reasons behind training networks with these kinds of approaches is to utilize the learned representations for downstream tasks like classification this is why this particular training phase is also referred to as _pre training_ we start by defining the loss function the s of simsiam emphasize the impact of the stop_gradient operator in the paper as it has an important role in the overall optimization negative cosine similarity minimizing this is equivalent to maximizing the similarity we then define our training loop by overriding the train_step function of the keras model class unpack the data forward pass through the encoder and predictor note that here we are enforcing the network to match the representations of two differently augmented batches of data compute gradients and update the parameters monitor loss pre training our networks in the interest of this example we will train the model for only 5 epochs in reality this should at least be 100 epochs create a cosine decay learning scheduler create an early stopping callback compile model and start training visualize the training progress of the model if your solution gets very close to 1 minimum value of our loss very quickly with a different dataset and a different backbone architecture that is likely because of representation collapse it is a phenomenon where the encoder yields similar output for all the images in that case additional hyperparameter tuning is required especially in the following areas strength of the color distortions and their probabilities learning rate and its schedule architecture of both the backbone and their projection head evaluating our ssl method the most popularly used method to evaluate a ssl method in computer vision or any other pre training method as such is to learn a linear classifier on the frozen features of the trained backbone model in this case it is resnet20 and evaluate the classifier on unseen images other methods include fine tuning https keras io guides transfer_learning on the source dataset or even a target dataset with 5 or 10 labels present practically we can use the backbone model for any downstream task such as semantic segmentation object detection and so on where the backbone models are usually pre trained with pure supervised learning we first create labeled dataset objects then we shuffle batch and prefetch this dataset for performance we also apply random resized crops as an augmentation but only to the training set extract the backbone resnet20 we then create our linear classifier and train it compile model and start training notes more data and longer pre training schedule benefit ssl in general ssl is particularly very helpful when you do not have access to very limited labeled training data but you can manage to build a large corpus of unlabeled data recently using an ssl method called swav https arxiv org abs 2006 09882 a group of researchers at facebook trained a regnet https arxiv org abs 2006 09882 on 2 billion images they were able to achieve downstream performance very close to those achieved by pure supervised pre training for some downstream tasks their method even outperformed the supervised counterparts you can check out their paper https arxiv org pdf 2103 01988 pdf to know the details if you are interested to understand why contrastive ssl helps networks learn meaningful representations you can check out the following resources self supervised learning the dark matter of intelligence https ai facebook com blog self supervised learning the dark matter of intelligence understanding self supervised learning using controlled datasets with known structure https sslneuips20 github io files camerareadys 203 77 64 cameraready understanding_self_supervised_learning pdf
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'

from keras import layers
from keras import regularizers
import keras
import tensorflow as tf

import matplotlib.pyplot as plt
import numpy as np



AUTO = tf.data.AUTOTUNE
BATCH_SIZE = 128
EPOCHS = 5
CROP_TO = 32
SEED = 26

PROJECT_DIM = 2048
LATENT_DIM = 512
WEIGHT_DECAY = 0.0005



(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
print(f"Total training examples: {len(x_train)}")
print(f"Total test examples: {len(x_test)}")




def flip_random_crop(image):
    
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_crop(image, (CROP_TO, CROP_TO, 3))
    return image


def color_jitter(x, strength=[0.4, 0.4, 0.4, 0.1]):
    x = tf.image.random_brightness(x, max_delta=0.8 * strength[0])
    x = tf.image.random_contrast(
        x, lower=1 - 0.8 * strength[1], upper=1 + 0.8 * strength[1]
    )
    x = tf.image.random_saturation(
        x, lower=1 - 0.8 * strength[2], upper=1 + 0.8 * strength[2]
    )
    x = tf.image.random_hue(x, max_delta=0.2 * strength[3])
    
    
    x = tf.clip_by_value(x, 0, 255)
    return x


def color_drop(x):
    x = tf.image.rgb_to_grayscale(x)
    x = tf.tile(x, [1, 1, 3])
    return x


def random_apply(func, x, p):
    if tf.random.uniform([], minval=0, maxval=1) < p:
        return func(x)
    else:
        return x


def custom_augment(image):
    
    
    
    image = flip_random_crop(image)
    image = random_apply(color_jitter, image, p=0.8)
    image = random_apply(color_drop, image, p=0.2)
    return image






ssl_ds_one = tf.data.Dataset.from_tensor_slices(x_train)
ssl_ds_one = (
    ssl_ds_one.shuffle(1024, seed=SEED)
    .map(custom_augment, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)

ssl_ds_two = tf.data.Dataset.from_tensor_slices(x_train)
ssl_ds_two = (
    ssl_ds_two.shuffle(1024, seed=SEED)
    .map(custom_augment, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)


ssl_ds = tf.data.Dataset.zip((ssl_ds_one, ssl_ds_two))


sample_images_one = next(iter(ssl_ds_one))
plt.figure(figsize=(10, 10))
for n in range(25):
    ax = plt.subplot(5, 5, n + 1)
    plt.imshow(sample_images_one[n].numpy().astype("int"))
    plt.axis("off")
plt.show()



sample_images_two = next(iter(ssl_ds_two))
plt.figure(figsize=(10, 10))
for n in range(25):
    ax = plt.subplot(5, 5, n + 1)
    plt.imshow(sample_images_two[n].numpy().astype("int"))
    plt.axis("off")
plt.show()







import resnet_cifar10_v2

N = 2
DEPTH = N * 9 + 2
NUM_BLOCKS = ((DEPTH - 2) // 9) - 1


def get_encoder():
    
    inputs = layers.Input((CROP_TO, CROP_TO, 3))
    x = layers.Rescaling(scale=1.0 / 127.5, offset=-1)(inputs)
    x = resnet_cifar10_v2.stem(x)
    x = resnet_cifar10_v2.learner(x, NUM_BLOCKS)
    x = layers.GlobalAveragePooling2D(name="backbone_pool")(x)

    
    x = layers.Dense(
        PROJECT_DIM,
        use_bias=False,
        kernel_regularizer=regularizers.l2(WEIGHT_DECAY),
    )(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Dense(
        PROJECT_DIM,
        use_bias=False,
        kernel_regularizer=regularizers.l2(WEIGHT_DECAY),
    )(x)
    outputs = layers.BatchNormalization()(x)
    return keras.Model(inputs, outputs, name="encoder")


def get_predictor():
    model = keras.Sequential(
        [
            
            layers.Input((PROJECT_DIM,)),
            layers.Dense(
                LATENT_DIM,
                use_bias=False,
                kernel_regularizer=regularizers.l2(WEIGHT_DECAY),
            ),
            layers.ReLU(),
            layers.BatchNormalization(),
            layers.Dense(PROJECT_DIM),
        ],
        name="predictor",
    )
    return model





def compute_loss(p, z):
    
    
    
    z = tf.stop_gradient(z)
    p = tf.math.l2_normalize(p, axis=1)
    z = tf.math.l2_normalize(z, axis=1)
    
    
    return -tf.reduce_mean(tf.reduce_sum((p * z), axis=1))





class SimSiam(keras.Model):
    def __init__(self, encoder, predictor):
        super().__init__()
        self.encoder = encoder
        self.predictor = predictor
        self.loss_tracker = keras.metrics.Mean(name="loss")

    @property
    def metrics(self):
        return [self.loss_tracker]

    def train_step(self, data):
        
        ds_one, ds_two = data

        
        with tf.GradientTape() as tape:
            z1, z2 = self.encoder(ds_one), self.encoder(ds_two)
            p1, p2 = self.predictor(z1), self.predictor(z2)
            
            
            
            loss = compute_loss(p1, z2) / 2 + compute_loss(p2, z1) / 2

        
        learnable_params = (
            self.encoder.trainable_variables
            + self.predictor.trainable_variables
        )
        gradients = tape.gradient(loss, learnable_params)
        self.optimizer.apply_gradients(zip(gradients, learnable_params))

        
        self.loss_tracker.update_state(loss)
        return {"loss": self.loss_tracker.result()}





num_training_samples = len(x_train)
steps = EPOCHS * (num_training_samples // BATCH_SIZE)
lr_decayed_fn = keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=0.03, decay_steps=steps
)


early_stopping = keras.callbacks.EarlyStopping(
    monitor="loss", patience=5, restore_best_weights=True
)


simsiam = SimSiam(get_encoder(), get_predictor())
simsiam.compile(optimizer=keras.optimizers.SGD(lr_decayed_fn, momentum=0.6))
history = simsiam.fit(ssl_ds, epochs=EPOCHS, callbacks=[early_stopping])


plt.plot(history.history["loss"])
plt.grid()
plt.title("Negative Cosine Similairty")
plt.show()






train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))




train_ds = (
    train_ds.shuffle(1024)
    .map(lambda x, y: (flip_random_crop(x), y), num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)
test_ds = test_ds.batch(BATCH_SIZE).prefetch(AUTO)


backbone = keras.Model(
    simsiam.encoder.input, simsiam.encoder.get_layer("backbone_pool").output
)


backbone.trainable = False
inputs = layers.Input((CROP_TO, CROP_TO, 3))
x = backbone(inputs, training=False)
outputs = layers.Dense(10, activation="softmax")(x)
linear_model = keras.Model(inputs, outputs, name="linear_model")


linear_model.compile(
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
    optimizer=keras.optimizers.SGD(lr_decayed_fn, momentum=0.9),
)
history = linear_model.fit(
    train_ds, validation_data=test_ds, epochs=EPOCHS, callbacks=[early_stopping]
)
_, test_acc = linear_model.evaluate(test_ds)
print("Test accuracy: {:.2f}%".format(test_acc * 100))