title supervised contrastive learning khalid salamahttps www linkedin cominkhalidsalama24403144 date created 20201130 last modified 20201130 description using supervised contrastive learning for image classification accelerator gpu introduction supervised contrastive learninghttps arxiv orgabs2004 11362 prannay khosla et al is a training methodology that outperforms supervised training with crossentropy on classification tasks essentially training an image classification model with supervised contrastive learning is performed in two phases 1 training an encoder to learn to produce vector representations of input images such 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 setup prepare the data load the train and test data splits display shapes of train and test datasets using image data augmentation setting the state of the normalization layer build the encoder model the encoder model takes the image as input and turns it into a 2048dimensional feature vector build the classification model the classification model adds a fullyconnected layer on top of the encoder plus a softmax layer with the target classes define npairs loss function computes the npairs loss between ytrue and ypred npairs loss expects paired data where a pair is composed of samples from the same labels and each pairs in the minibatch have different labels the loss takes each row of the pairwise similarity matrix ypred as logits and the remapped multiclass labels ytrue as labels see http www neclabs comuploadsimagesdepartmentimagesmediaanalyticspapersnips16npairmetriclearning pdf args ytrue ground truth values of shape batchsize of multiclass labels ypred predicted values of shape batchsize batchsize of similarity matrix between embedding matrices returns npairsloss float scalar experiment 1 train the baseline classification model in this experiment a baseline classifier is trained as usual i e the encoder and the classifier parts are trained together as a single model to minimize the crossentropy loss experiment 2 use supervised contrastive learning in this experiment the model is trained in two phases in the first phase the encoder is pretrained to optimize the supervised contrastive loss described in prannay khosla et al https arxiv orgabs2004 11362 in the second phase the classifier is trained using the trained encoder with its weights freezed only the weights of fullyconnected layers with the softmax are optimized 1 supervised contrastive learning loss function normalize feature vectors compute logits 2 pretrain the encoder 3 train the classifier with the frozen encoder we get to an improved test accuracy conclusion as shown in the experiments using the supervised contrastive learning technique outperformed the conventional technique in terms of the test accuracy note that the same training budget i e number of epochs was given to each technique supervised contrastive learning pays off when the encoder involves a complex architecture like resnet and multiclass problems with many labels in addition large batch sizes and multilayer projection heads improve its effectiveness see the supervised contrastive learninghttps arxiv orgabs2004 11362 paper for more details introduction supervised contrastive learning https arxiv org abs 2004 11362 prannay khosla et al is a training methodology that outperforms supervised training with crossentropy on classification tasks essentially training an image classification model with supervised contrastive learning is performed in two phases 1 training an encoder to learn to produce vector representations of input images such 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 setup prepare the data load the train and test data splits display shapes of train and test datasets using image data augmentation setting the state of the normalization layer build the encoder model the encoder model takes the image as input and turns it into a 2048 dimensional feature vector build the classification model the classification model adds a fully connected layer on top of the encoder plus a softmax layer with the target classes define npairs loss function computes the npairs loss between y_true and y_pred npairs loss expects paired data where a pair is composed of samples from the same labels and each pairs in the minibatch have different labels the loss takes each row of the pair wise similarity matrix y_pred as logits and the remapped multi class labels y_true as labels see http www nec labs com uploads images department images mediaanalytics papers nips16_npairmetriclearning pdf args y_true ground truth values of shape batch_size of multi class labels y_pred predicted values of shape batch_size batch_size of similarity matrix between embedding matrices returns npairs_loss float scalar experiment 1 train the baseline classification model in this experiment a baseline classifier is trained as usual i e the encoder and the classifier parts are trained together as a single model to minimize the crossentropy loss experiment 2 use supervised contrastive learning in this experiment the model is trained in two phases in the first phase the encoder is pretrained to optimize the supervised contrastive loss described in prannay khosla et al https arxiv org abs 2004 11362 in the second phase the classifier is trained using the trained encoder with its weights freezed only the weights of fully connected layers with the softmax are optimized 1 supervised contrastive learning loss function normalize feature vectors compute logits 2 pretrain the encoder 3 train the classifier with the frozen encoder we get to an improved test accuracy conclusion as shown in the experiments using the supervised contrastive learning technique outperformed the conventional technique in terms of the test accuracy note that the same training budget i e number of epochs was given to each technique supervised contrastive learning pays off when the encoder involves a complex architecture like resnet and multi class problems with many labels in addition large batch sizes and multi layer projection heads improve its effectiveness see the supervised contrastive learning https arxiv org abs 2004 11362 paper for more details
import keras
from keras import ops
from keras import layers
from keras.applications.resnet_v2 import ResNet50V2



num_classes = 10
input_shape = (32, 32, 3)


(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()


print(f"x_train shape: {x_train.shape} - y_train shape: {y_train.shape}")
print(f"x_test shape: {x_test.shape} - y_test shape: {y_test.shape}")




data_augmentation = keras.Sequential(
    [
        layers.Normalization(),
        layers.RandomFlip("horizontal"),
        layers.RandomRotation(0.02),
    ]
)


data_augmentation.layers[0].adapt(x_train)




def create_encoder():
    resnet = ResNet50V2(
        include_top=False, weights=None, input_shape=input_shape, pooling="avg"
    )

    inputs = keras.Input(shape=input_shape)
    augmented = data_augmentation(inputs)
    outputs = resnet(augmented)
    model = keras.Model(inputs=inputs, outputs=outputs, name="cifar10-encoder")
    return model


encoder = create_encoder()
encoder.summary()

learning_rate = 0.001
batch_size = 265
hidden_units = 512
projection_units = 128
num_epochs = 50
dropout_rate = 0.5
temperature = 0.05




def create_classifier(encoder, trainable=True):
    for layer in encoder.layers:
        layer.trainable = trainable

    inputs = keras.Input(shape=input_shape)
    features = encoder(inputs)
    features = layers.Dropout(dropout_rate)(features)
    features = layers.Dense(hidden_units, activation="relu")(features)
    features = layers.Dropout(dropout_rate)(features)
    outputs = layers.Dense(num_classes, activation="softmax")(features)

    model = keras.Model(
        inputs=inputs, outputs=outputs, name="cifar10-classifier"
    )
    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate),
        loss=keras.losses.SparseCategoricalCrossentropy(),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )
    return model





def npairs_loss(y_true, y_pred):
    
    y_pred = ops.cast(y_pred, "float32")
    y_true = ops.cast(y_true, y_pred.dtype)
    y_true = ops.cast(ops.equal(y_true, ops.transpose(y_true)), y_pred.dtype)
    y_true /= ops.sum(y_true, 1, keepdims=True)
    loss = ops.categorical_crossentropy(y_true, y_pred, from_logits=True)

    return ops.mean(loss)




encoder = create_encoder()
classifier = create_classifier(encoder)
classifier.summary()

history = classifier.fit(
    x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs
)

accuracy = classifier.evaluate(x_test, y_test)[1]
print(f"Test accuracy: {round(accuracy * 100, 2)}%")





class SupervisedContrastiveLoss(keras.losses.Loss):
    def __init__(self, temperature=1, name=None):
        super().__init__(name=name)
        self.temperature = temperature

    def __call__(self, labels, feature_vectors, sample_weight=None):
        
        feature_vectors_normalized = keras.utils.normalize(
            feature_vectors, axis=1, order=2
        )
        
        logits = ops.divide(
            ops.matmul(
                feature_vectors_normalized,
                ops.transpose(feature_vectors_normalized),
            ),
            self.temperature,
        )
        return npairs_loss(ops.squeeze(labels), logits)


def add_projection_head(encoder):
    inputs = keras.Input(shape=input_shape)
    features = encoder(inputs)
    outputs = layers.Dense(projection_units, activation="relu")(features)
    model = keras.Model(
        inputs=inputs,
        outputs=outputs,
        name="cifar-encoder_with_projection-head",
    )
    return model




encoder = create_encoder()

encoder_with_projection_head = add_projection_head(encoder)
encoder_with_projection_head.compile(
    optimizer=keras.optimizers.Adam(learning_rate),
    loss=SupervisedContrastiveLoss(temperature),
)

encoder_with_projection_head.summary()

history = encoder_with_projection_head.fit(
    x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs
)



classifier = create_classifier(encoder, trainable=False)

history = classifier.fit(
    x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs
)

accuracy = classifier.evaluate(x_test, y_test)[1]
print(f"Test accuracy: {round(accuracy * 100, 2)}%")