{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\nimport tensorflow_addons as tfa\nimport numpy as np\nfrom tensorflow import keras\nfrom tensorflow.keras import layers","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_classes = 10\ninput_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\nprint(f\"x_train shape: {x_train.shape} - y_train shape: {y_train.shape}\")\nprint(f\"x_test shape: {x_test.shape} - y_test shape: {y_test.shape}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_augmentation = keras.Sequential(\n    [\n        layers.experimental.preprocessing.Normalization(),\n        layers.experimental.preprocessing.RandomFlip(\"horizontal\"),\n        layers.experimental.preprocessing.RandomRotation(0.02),\n        layers.experimental.preprocessing.RandomWidth(0.2),\n        layers.experimental.preprocessing.RandomHeight(0.2),\n    ]\n)\n\n# Setting the state of the normalization layer.\ndata_augmentation.layers[0].adapt(x_train)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_encoder():\n    resnet = keras.applications.ResNet50V2(\n        include_top=False, weights=None, input_shape=input_shape, pooling=\"avg\"\n    )\n#     b4 = keras.applications.EfficientNetB4(\n#             include_top=False,\n#             weights=\"imagenet\",\n#             input_shape=input_shape,\n#             pooling=\"None\",\n#         )\n# #     b4.add(layers.GlobalMaxPooling2D(name=\"gap\"))\n    inputs = keras.Input(shape=input_shape)\n    augmented = data_augmentation(inputs)\n    outputs = resnet(augmented)\n    model = keras.Model(inputs=inputs, outputs=outputs, name=\"cifar10-encoder\")\n    return model\n\n\nencoder = create_encoder()\nencoder.summary()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_classifier(encoder, trainable=True):\n\n    for layer in encoder.layers:\n        layer.trainable = trainable\n\n    inputs = keras.Input(shape=input_shape)\n    features = encoder(inputs)\n    features = layers.Dropout(dropout_rate)(features)\n    features = layers.Dense(hidden_units, activation=\"relu\")(features)\n    features = layers.Dropout(dropout_rate)(features)\n    outputs = layers.Dense(num_classes, activation=\"softmax\")(features)\n\n    model = keras.Model(inputs=inputs, outputs=outputs, name=\"cifar10-classifier\")\n    model.compile(\n        optimizer=keras.optimizers.Adam(learning_rate),\n        loss=keras.losses.SparseCategoricalCrossentropy(),\n        metrics=[keras.metrics.SparseCategoricalAccuracy()],\n    )\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learning_rate = 0.001\nbatch_size = 265\nhidden_units = 512\nprojection_units = 128\nnum_epochs = 50\ndropout_rate = 0.5\ntemperature = 0.05","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SupervisedContrastiveLoss(keras.losses.Loss):\n    def __init__(self, temperature=1, name=None):\n        super(SupervisedContrastiveLoss, self).__init__(name=name)\n        self.temperature = temperature\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            self.temperature,\n        )\n        return tfa.losses.npairs_loss(tf.squeeze(labels), logits)\n\n\ndef add_projection_head(encoder):\n    inputs = keras.Input(shape=input_shape)\n    features = encoder(inputs)\n    outputs = layers.Dense(projection_units, activation=\"relu\")(features)\n    model = keras.Model(inputs=inputs, outputs=outputs, name=\"cifar-encoder_with_projection-head\")\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"encoder = create_encoder()\n\nencoder_with_projection_head = add_projection_head(encoder)\nencoder_with_projection_head.compile(\n    optimizer=keras.optimizers.Adam(learning_rate),\n    loss=SupervisedContrastiveLoss(temperature),\n)\n\nencoder_with_projection_head.summary()\n\nhistory = encoder_with_projection_head.fit( x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"classifier = create_classifier(encoder, trainable=False)\n\nhistory = classifier.fit(x=x_train, y=y_train, batch_size=batch_size, epochs=num_epochs)\n\naccuracy = classifier.evaluate(x_test, y_test)[1]\nprint(f\"Test accuracy: {round(accuracy * 100, 2)}%\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}