ArcFace_1.py

class ArcFace(tf.keras.layers.Layer):
    
    # Implementation reference from https://github.com/lyakaap/Landmark2019-1st-and-3rd-Place-Solution/blob/master/src/modeling/metric_learning.py
    
    def __init__(self, n_classes, scale, margin, **kwargs):

        super(ArcFace, self).__init__(**kwargs)

        self.n_classes = n_classes
        self.scale = scale
        self.margin = margin
        self.cos_m = tf.math.cos(margin)
        self.sin_m = tf.math.sin(margin)

    def get_config(self):

        config = super().get_config().copy()
        config.update({'n_classes': self.n_classes,'scale': self.scale,'margin': self.margin})
        return config

    def build(self, input_shape):
        super(ArcFace, self).build(input_shape[0])

        self.W = self.add_weight(
            name='W',
            shape=(int(input_shape[0][-1]), self.n_classes),
            initializer='glorot_uniform',
            dtype='float32',
            trainable=True,
            regularizer=None)

    def call(self, inputs):
        X, y = inputs
        y = tf.cast(y, dtype=tf.int32)

        # Normalizing vectors( Unit Vectors ) to make dot product depend only on angle between vectors.
        cosine = tf.matmul(
            tf.math.l2_normalize(X, axis=1),
            tf.math.l2_normalize(self.W, axis=0)
        )

        # Sin(angle)^2 + Cos(angle)^2 = 1
        sine = tf.math.sqrt(1.0 - tf.math.pow(cosine, 2))
        
        # Cos(angle+margin)=Cos(angle)*Cos(margin)-Sin(angle)*Sin(margin)
        phi = cosine * self.cos_m - sine * self.sin_m

        one_hot = tf.cast(tf.one_hot(y, depth=self.n_classes),dtype=cosine.dtype)
    
        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
        output *= self.scale
        return output