包含用CosFace、SphereFace以及ArcFace的Loss函数训练MINST数据集的代码,顺便对TensorFlow练练手
环境是TensorFlow 2.8.0
# TensorFlow and tf.keras
import tensorflow as tf
from tensorflow import keras
# Helper libraries
import numpy as np
import matplotlib.pyplot as plt
print(tf.__version__)
2.8.0
np.random.seed(1)
tf.random.set_seed(1)
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
class custom_layer_softmax(tf.keras.layers.Layer):
def __init__(self, output_nums):
super(custom_layer_softmax, self).__init__()
self.output_nums = output_nums
#self.D = tf.keras.layers.Dense(10)
def build(self, input_shape):
self.w = self.add_weight(
shape=(input_shape[-1], self.output_nums),
initializer="random_normal",
trainable=True,
)
self.b = self.add_weight(
shape=(self.output_nums,),
initializer=tf.zeros_initializer(),
trainable=True,
)
# 为了适应下面的类,加入一个多余的参数
def call(self, y_pred, y_true):
y_pred = tf.matmul(y_pred, self.w)+self.b
return tf.math.softmax(y_pred)
class custom_layer_arcface(tf.keras.layers.Layer):
def __init__(self, output_nums):
super(custom_layer_arcface, self).__init__()
self.output_nums = output_nums
self.m = tf.constant(0.15)
def build(self, input_shape):
self.w = self.add_weight(
shape=(input_shape[-1], self.output_nums),
initializer="random_normal",
trainable=True,
)
def call(self, y_pred, y_true):
_x = tf.math.l2_normalize(y_pred, 1)
_w = tf.transpose(tf.math.l2_normalize(tf.transpose(self.w), 1))
out = tf.matmul(_x, _w)
y_true = tf.one_hot(y_true, self.output_nums)
orig_cos_theta = tf.reduce_sum((y_true*out), 1, True)
#arc(-1) arc(1)的导数是无穷大,所以要避免。否则会出现NaN,查了好久
orig_cos_theta = tf.clip_by_value(orig_cos_theta, -1+0.00001, 1-0.00001)
theta = tf.acos(orig_cos_theta)
cos_theta = tf.cos(theta + self.m)
out += (cos_theta - orig_cos_theta) * y_true
out = tf.norm(y_pred, 2, 1, True)*out
return tf.math.softmax(out)
class custom_layer_cosface(tf.keras.layers.Layer):
def __init__(self, output_nums):
super(custom_layer_cosface, self).__init__()
self.output_nums = output_nums
self.m = tf.constant(0.1)
def build(self, input_shape):
self.w = self.add_weight(
shape=(input_shape[-1], self.output_nums),
initializer="random_normal",
trainable=True,
)
def call(self, y_pred, y_true):
_x = tf.math.l2_normalize(y_pred, 1)
_w = tf.transpose(tf.math.l2_normalize(tf.transpose(self.w), 1))
out = tf.matmul(_x, _w)
y_true = tf.one_hot(y_true, self.output_nums)
out -= self.m * y_true
out = tf.norm(y_pred, 2, 1, True)*out
return tf.math.softmax(out)
class custom_layer_sphereface(tf.keras.layers