Skip to content

Instantly share code, notes, and snippets.

@Chaitanya-Boyalla
Created July 16, 2021 19:06
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment