wenmin-wu/ds-skills/arcface-metric-learning

import tensorflow as tf
import math
class ArcMarginProduct(tf.keras.layers.Layer):
    def __init__(self, n_classes, s=30, m=0.50, **kwargs):
        super().__init__(**kwargs)
        self.n_classes = n_classes
        self.s = s
        self.cos_m = tf.math.cos(m)
        self.sin_m = tf.math.sin(m)
        self.th = tf.math.cos(math.pi - m)
        self.mm = tf.math.sin(math.pi - m) * m
    def build(self, input_shape):
        self.W = self.add_weight(shape=(input_shape[0][-1], self.n_classes),
                                 initializer='glorot_uniform', trainable=True)
    def call(self, inputs):
        X, y = inputs
        cosine = tf.matmul(tf.math.l2_normalize(X, axis=1),
                           tf.math.l2_normalize(self.W, axis=0))
        sine = tf.math.sqrt(1.0 - tf.math.pow(cosine, 2))
        phi = cosine * self.cos_m - sine * self.sin_m
        one_hot = tf.cast(tf.one_hot(y, depth=self.n_classes), dtype=phi.dtype)
        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
        return output * self.s
# Training
backbone = tf.keras.applications.EfficientNetB3(include_top=False, pooling='avg')
x = backbone.output
margin = ArcMarginProduct(n_classes=num_products, s=30, m=0.5)
output = margin([x, label_input])
model.fit(...)
# Inference: extract embeddings (discard ArcFace head)
embedder = tf.keras.Model(inputs=model.input[0], outputs=model.layers[-4].output)
embeddings = embedder.predict(test_data)