{"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"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":19231,"databundleVersionId":1413778,"sourceType":"competition"}],"dockerImageVersionId":30008,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Keras Implementation of ArcFace, CosFace, AdaCos\n\nIn this notebook you may find the implementation of layers that you may need for this competition.\n\n## Why Softmax and Triplet Loss is bad?\nThere are two main lines research to train CNN for face recognition, one that train a multi-class classifier using softmax classifier and the other that learn the embeddings such as the triplet loss. However both have their drawbacks:\n - For the softmax loss, the more you add the different identities for recognition, the more number of parameters will increase.\n - And for the triplet loss there is a combinatorial explosion in the number of face triplets for large scale dataset leading to large number of iterations.\n\nBelow is the visualization of Softmax trained on MNIST. As you can see, the points are quite scattered even for such a simple dataset. We will compare this to other losses in next sections.\n![](https://github.com/4uiiurz1/keras-arcface/raw/master/figures/mnist_vgg8_3d.png)","metadata":{}},{"cell_type":"code","source":"import math\nimport tensorflow as tf\nimport tensorflow.keras.backend as K\nfrom tensorflow.keras.layers import Layer\nfrom tensorflow.keras.initializers import Constant\nfrom tensorflow.python.keras.utils import tf_utils\n\n\ndef _resolve_training(layer, training):\n    if training is None:\n        training = K.learning_phase()\n    if isinstance(training, int):\n        training = bool(training)\n    if not layer.trainable:\n        # When the layer is not trainable, override the value\n        training = False\n    return tf_utils.constant_value(training)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ArcFace\n\nPaper source: https://arxiv.org/pdf/1801.07698.pdf\n\nAn additive angular margin loss is proposed in arcface to further improve the descriminative power of the face recognition model and stabilize the training process. The arc-cosine function is used to calculate angle between the current feature and target weight. ArcFace directly optimizes the geodesic distance margin by virtue of exact correspondence between angle and arc in the normalized hypersphere.\n\n$$L_A = -\\frac{1}{N} \\sum_{i=1}^N \\log \\frac{e^{s(\\cos (\\theta_{y_i} + m))}}{e^{s(\\cos (\\theta_{y_i} + m))} + \\sum_{j=1, j\\ne y_i}^N e^{s \\cos \\theta_j}}$$\n\nBelow is the visualization for a network trained on MNIST with ArcFace.\n![](https://github.com/4uiiurz1/keras-arcface/raw/master/figures/mnist_vgg8_arcface_3d.png)","metadata":{}},{"cell_type":"code","source":"class ArcFace(Layer):\n    \"\"\"\n    Implementation of ArcFace layer. Reference: https://arxiv.org/abs/1801.07698\n    \n    Arguments:\n      num_classes: number of classes to classify\n      s: scale factor\n      m: margin\n      regularizer: weights regularizer\n    \"\"\"\n    def __init__(self,\n                 num_classes,\n                 s=30.0,\n                 m=0.5,\n                 regularizer=None,\n                 name='arcface',\n                 **kwargs):\n        \n        super().__init__(name=name, **kwargs)\n        self._n_classes = num_classes\n        self._s = float(s)\n        self._m = float(m)\n        self._regularizer = regularizer\n\n    def build(self, input_shape):\n        embedding_shape, label_shape = input_shape\n        self._w = self.add_weight(shape=(embedding_shape[-1], self._n_classes),\n                                  initializer='glorot_uniform',\n                                  trainable=True,\n                                  regularizer=self._regularizer,\n                                  name='cosine_weights')\n\n    def call(self, inputs, training=None):\n        \"\"\"\n        During training, requires 2 inputs: embedding (after backbone+pool+dense),\n        and ground truth labels. The labels should be sparse (and use\n        sparse_categorical_crossentropy as loss).\n        \"\"\"\n        embedding, label = inputs\n\n        # Squeezing is necessary for Keras. It expands the dimension to (n, 1)\n        label = tf.reshape(label, [-1], name='label_shape_correction')\n\n        # Normalize features and weights and compute dot product\n        x = tf.nn.l2_normalize(embedding, axis=1, name='normalize_prelogits')\n        w = tf.nn.l2_normalize(self._w, axis=0, name='normalize_weights')\n        cosine_sim = tf.matmul(x, w, name='cosine_similarity')\n\n        training = resolve_training_flag(self, training)\n        if not training:\n            # We don't have labels if we're not in training mode\n            return self._s * cosine_sim\n        else:\n            one_hot_labels = tf.one_hot(label,\n                                        depth=self._n_classes,\n                                        name='one_hot_labels')\n            theta = tf.math.acos(K.clip(\n                    cosine_sim, -1.0 + K.epsilon(), 1.0 - K.epsilon()))\n            selected_labels = tf.where(tf.greater(theta, math.pi - self._m),\n                                       tf.zeros_like(one_hot_labels),\n                                       one_hot_labels,\n                                       name='selected_labels')\n            final_theta = tf.where(tf.cast(selected_labels, dtype=tf.bool),\n                                   theta + self._m,\n                                   theta,\n                                   name='final_theta')\n            output = tf.math.cos(final_theta, name='cosine_sim_with_margin')\n            return self._s * output","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CosFace: Large Margin Cosine Loss\n\nPaper source: https://arxiv.org/abs/1801.09414\n\nLarge Margin Cosine Loss (LMCL) which is referred as CosFace, reformulates the traditional softmax loss as a cosine loss by L2 normalizing both features and weight vectors to remove radial variations, based on this a cosine margin term is introduced to further maximize the decision margin in the angular space. As a result of which we get minimum intra-class margin and maximum inter-class margin for accurate face verification.\n\nThe formula for CosFace is pretty similar to ArcFace:\n\n$$L_A = -\\frac{1}{N} \\sum_{i=1}^N \\log \\frac{e^{s(\\cos (\\theta_{y_i}) - m)}}{e^{s(\\cos (\\theta_{y_i}) - m)} + \\sum_{j=1, j\\ne y_i}^N e^{s \\cos \\theta_j}}$$\n\nBelow is the visualization for a network trained on MNIST with CosFace.\n\n![](https://github.com/4uiiurz1/keras-arcface/raw/master/figures/mnist_vgg8_cosface_3d.png)","metadata":{}},{"cell_type":"code","source":"class CosFace(Layer):\n    \"\"\"\n    Implementation of CosFace layer. Reference: https://arxiv.org/abs/1801.09414\n    \n    Arguments:\n      num_classes: number of classes to classify\n      s: scale factor\n      m: margin\n      regularizer: weights regularizer\n    \"\"\"\n    def __init__(self,\n                 num_classes,\n                 s=30.0,\n                 m=0.35,\n                 regularizer=None,\n                 name='cosface',\n                 **kwargs):\n\n        super().__init__(name=name, **kwargs)\n        self._n_classes = num_classes\n        self._s = float(s)\n        self._m = float(m)\n        self._regularizer = regularizer\n\n    def build(self, input_shape):\n        embedding_shape, label_shape = input_shape\n        self._w = self.add_weight(shape=(embedding_shape[-1], self._n_classes),\n                                  initializer='glorot_uniform',\n                                  trainable=True,\n                                  regularizer=self._regularizer)\n\n    def call(self, inputs, training=None):\n        \"\"\"\n        During training, requires 2 inputs: embedding (after backbone+pool+dense),\n        and ground truth labels. The labels should be sparse (and use\n        sparse_categorical_crossentropy as loss).\n        \"\"\"\n        embedding, label = inputs\n\n        # Squeezing is necessary for Keras. It expands the dimension to (n, 1)\n        label = tf.reshape(label, [-1], name='label_shape_correction')\n\n        # Normalize features and weights and compute dot product\n        x = tf.nn.l2_normalize(embedding, axis=1, name='normalize_prelogits')\n        w = tf.nn.l2_normalize(self._w, axis=0, name='normalize_weights')\n        cosine_sim = tf.matmul(x, w, name='cosine_similarity')\n\n        training = _resolve_training(self, training)\n        if not training:\n            # We don't have labels if we're not in training mode\n            return self._s * cosine_sim\n        else:\n            one_hot_labels = tf.one_hot(label,\n                                        depth=self._n_classes,\n                                        name='one_hot_labels')\n            final_theta = tf.where(tf.cast(one_hot_labels, dtype=tf.bool),\n                                   cosine_sim - self._m,\n                                   cosine_sim,\n                                   name='cosine_sim_with_margin')\n            return self._s * output","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# AdaCos: Adaptively Scaling Cosine Logits\n\nPaper source: https://arxiv.org/abs/1905.00292","metadata":{}},{"cell_type":"code","source":"class AdaCos(Layer):\n    \"\"\"\n    Implementation of AdaCos layer. Reference: https://arxiv.org/abs/1905.00292\n    \n    Arguments:\n      num_classes: number of classes to classify\n      is_dynamic: if False, use Fixed AdaCos. Else, use Dynamic Adacos.\n      regularizer: weights regularizer\n    \"\"\"\n    def __init__(self,\n                 num_classes,\n                 is_dynamic=True,\n                 regularizer=None,\n                 **kwargs):\n\n        super().__init__(**kwargs)\n        self._n_classes = num_classes\n        self._init_s = math.sqrt(2) * math.log(num_classes - 1)\n        self._is_dynamic = is_dynamic\n        self._regularizer = regularizer\n\n    def build(self, input_shape):\n        embedding_shape, label_shape = input_shape\n        self._w = self.add_weight(shape=(embedding_shape[-1], self._n_classes),\n                                  initializer='glorot_uniform',\n                                  trainable=True,\n                                  regularizer=self._regularizer)\n        if self._is_dynamic:\n            self._s = self.add_weight(shape=(),\n                                      initializer=Constant(self._init_s),\n                                      trainable=False,\n                                      aggregation=tf.VariableAggregation.MEAN)\n\n    def call(self, inputs, training=None):\n        embedding, label = inputs\n\n        # Squeezing is necessary for Keras. It expands the dimension to (n, 1)\n        label = tf.reshape(label, [-1])\n\n        # Normalize features and weights and compute dot product\n        x = tf.nn.l2_normalize(embedding, axis=1)\n        w = tf.nn.l2_normalize(self._w, axis=0)\n        logits = tf.matmul(x, w)\n\n        # Fixed AdaCos\n        is_dynamic = tf_utils.constant_value(self._is_dynamic)\n        if not is_dynamic:\n            # _s is not created since we are not in dynamic mode\n            output = tf.multiply(self._init_s, logits)\n            return output\n\n        training = _resolve_training(self, training)\n        if not training:\n            # We don't have labels to update _s if we're not in training mode\n            return self._s * logits\n        else:\n            theta = tf.math.acos(\n                    K.clip(logits, -1.0 + K.epsilon(), 1.0 - K.epsilon()))\n            one_hot = tf.one_hot(label, depth=self._n_classes)\n            b_avg = tf.where(one_hot < 1.0,\n                             tf.exp(self._s * logits),\n                             tf.zeros_like(logits))\n            b_avg = tf.reduce_mean(tf.reduce_sum(b_avg, axis=1))\n            theta_class = tf.gather_nd(\n                    theta,\n                    tf.stack([\n                        tf.range(tf.shape(label)[0]),\n                        tf.cast(label, tf.int32)\n                    ], axis=1))\n            mid_index = tf.shape(theta_class)[0] // 2 + 1\n            theta_med = tf.nn.top_k(theta_class, mid_index).values[-1]\n\n            # Since _s is not trainable, this assignment is safe. Also,\n            # tf.function ensures that this will run in the right order.\n            self._s.assign(\n                    tf.math.log(b_avg) /\n                    tf.math.cos(tf.minimum(math.pi/4, theta_med)))\n\n            # Return scaled logits\n            return self._s * logits","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}