from keras import backend as K from keras.layers import Layer from keras import regularizers import tensorflow as tf class ArcFace(Layer): def __init__(self, n_classes=10, s=30.0, m=0.50, regularizer=None, **kwargs): super(ArcFace, self).__init__(**kwargs) self.n_classes = n_classes self.s = s self.m = m self.regularizer = regularizers.get(regularizer) def build(self, input_shape): super(ArcFace, self).build(input_shape[0]) self.W = self.add_weight(name='W', shape=(input_shape[0][-1], self.n_classes), initializer='glorot_uniform', trainable=True, regularizer=self.regularizer) def call(self, inputs): x, y = inputs c = K.shape(x)[-1] # normalize feature x = tf.nn.l2_normalize(x, axis=1) # normalize weights W = tf.nn.l2_normalize(self.W, axis=0) # dot product logits = x @ W # add margin # clip logits to prevent zero division when backward theta = tf.acos(K.clip(logits, -1.0 + K.epsilon(), 1.0 - K.epsilon())) target_logits = tf.cos(theta + self.m) # sin = tf.sqrt(1 - logits**2) # cos_m = tf.cos(logits) # sin_m = tf.sin(logits) # target_logits = logits * cos_m - sin * sin_m # logits = logits * (1 - y) + target_logits * y # feature re-scale logits *= self.s out = tf.nn.softmax(logits) return out def compute_output_shape(self, input_shape): return (None, self.n_classes) class SphereFace(Layer): def __init__(self, n_classes=10, s=30.0, m=1.35, regularizer=None, **kwargs): super(SphereFace, self).__init__(**kwargs) self.n_classes = n_classes self.s = s self.m = m self.regularizer = regularizers.get(regularizer) def build(self, input_shape): super(SphereFace, self).build(input_shape[0]) self.W = self.add_weight(name='W', shape=(input_shape[0][-1], self.n_classes), initializer='glorot_uniform', trainable=True, regularizer=self.regularizer) def call(self, inputs): x, y = inputs c = K.shape(x)[-1] # normalize feature x = tf.nn.l2_normalize(x, axis=1) # normalize weights W = tf.nn.l2_normalize(self.W, axis=0) # dot product logits = x @ W # add margin # clip logits to prevent zero division when backward theta = tf.acos(K.clip(logits, -1.0 + K.epsilon(), 1.0 - K.epsilon())) target_logits = tf.cos(self.m * theta) # logits = logits * (1 - y) + target_logits * y # feature re-scale logits *= self.s out = tf.nn.softmax(logits) return out def compute_output_shape(self, input_shape): return (None, self.n_classes) class CosFace(Layer): def __init__(self, n_classes=10, s=30.0, m=0.35, regularizer=None, **kwargs): super(CosFace, self).__init__(**kwargs) self.n_classes = n_classes self.s = s self.m = m self.regularizer = regularizers.get(regularizer) def build(self, input_shape): super(CosFace, self).build(input_shape[0]) self.W = self.add_weight(name='W', shape=(input_shape[0][-1], self.n_classes), initializer='glorot_uniform', trainable=True, regularizer=self.regularizer) def call(self, inputs): x, y = inputs c = K.shape(x)[-1] # normalize feature x = tf.nn.l2_normalize(x, axis=1) # normalize weights W = tf.nn.l2_normalize(self.W, axis=0) # dot product logits = x @ W # add margin target_logits = logits - self.m # logits = logits * (1 - y) + target_logits * y # feature re-scale logits *= self.s out = tf.nn.softmax(logits) return out def compute_output_shape(self, input_shape): return (None, self.n_classes)