AM-Softmax的keras实现: 《Additive Margin Softmax for Face Verification》

原答案在对输入进行归一化时有错误，另外m应该是一个固定的超参数不需要训练，已改正

论文地址：Additive Margin Softmax for Face Verification

和L-Softmax、A-Softmax比简单很多，在minist上测试通过，两个500节点的全连接层 + 10个节点的分类层，训练两轮验证集的acc达到97+%。

运行环境：

keras：2.1.5

tensorflow：1.4.0

python：3.6

OS：win7

# -*- coding: utf-8 -*-
from keras import backend as K
from keras.engine.topology import Layer
from keras.layers import Dense, Activation,BatchNormalization
from keras.layers import activations, initializers, regularizers, constraints, Lambda
from keras.engine import InputSpec
import tensorflow as tf
import numpy as np

class AMSoftmax(Layer):
def __init__(self, units, s, m,
kernel_initializer='glorot_uniform',
kernel_regularizer=None,
kernel_constraint=None,
**kwargs
):
if 'input_shape' not in kwargs and 'input_dim' in kwargs:
kwargs['input_shape'] = (kwargs.pop('input_dim'),)
super(AMSoftmax, self).__init__(**kwargs)
self.units = units
self.s = s
self.m = m
self.kernel_initializer = initializers.get(kernel_initializer)
self.kernel_regularizer = regularizers.get(kernel_regularizer)
self.kernel_constraint = constraints.get(kernel_constraint)
self.input_spec = InputSpec(min_ndim=2)
self.supports_masking = True

def build(self, input_shape):
assert len(input_shape) >= 2
input_dim = input_shape[-1]

self.kernel = self.add_weight(shape=(input_dim, self.units),
initializer=self.kernel_initializer,
name='kernel',
regularizer=self.kernel_regularizer,
constraint=self.kernel_constraint)
self.bias = None

self.input_spec = InputSpec(min_ndim=2, axes={-1: input_dim})
self.built = True

def call(self, inputs, **kwargs):
inputs = tf.nn.l2_normalize(inputs, dim=-1)
self.kernel = tf.nn.l2_normalize(self.kernel, dim=(0, 1))   # W归一化

dis_cosin = K.dot(inputs, self.kernel)
psi = dis_cosin - self.m

e_costheta = K.exp(self.s * dis_cosin)
e_psi = K.exp(self.s * psi)
sum_x = K.sum(e_costheta, axis=-1, keepdims=True)

temp = e_psi - e_costheta
temp = temp + sum_x

output = e_psi / temp
return output

def amsoftmax_loss(y_true, y_pred):
d1 = K.sum(y_true * y_pred, axis=-1)
d1 = K.log(K.clip(d1, K.epsilon(), None))
loss = -K.mean(d1, axis=-1)
return loss

minist上测试程序：

import numpy as np
np.random.seed(1337)
from AM_Softmax import *
from keras import backend as K
from keras.layers import Dense,Input,Conv2D,MaxPooling2D,Dropout,BatchNormalization
from keras.models import Model
from keras.optimizers import SGD, Adam
import os
from keras.backend.tensorflow_backend import set_session

batch_size = 200
nb_classes = 10
nb_epoch = 2

os.environ["CUDA_VISIBLE_DEVICES"] = '1'
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.8
set_session(tf.Session(config=config))

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

X_train, Y_train = mnist.train.images,mnist.train.labels
X_test, Y_test = mnist.test.images, mnist.test.labels
X_train = X_train.reshape(-1, 28, 28,1).astype('float32')
X_test = X_test.reshape(-1,28, 28,1).astype('float32')

#打印训练数据和测试数据的维度
print(X_train.shape,X_test.shape,Y_train.shape,Y_test.shape)

#修改维度
X_train = X_train.reshape(55000,784)
X_test = X_test.reshape(10000,784)
print(X_train.shape, X_test.shape, Y_train.shape, Y_test.shape)

# 将X_train, X_test的数据格式转为float32存储
X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
# 归一化
X_train /= 255
X_test /= 255
# 打印出训练集和测试集的信息
print(X_train.shape[0], 'train samples')
print(X_test.shape[0], 'test samples')

x_input = Inp
4000
ut(shape=(784,))
y = Dense(500, activation='relu')(x_input)
y = Dropout(0.2)(y)
y = Dense(500, activation='relu')(y)
y = Dropout(0.2)(y)

output = AMSoftmax(10, 10, 0.35)(y)
#output = Dense(10, activation='softmax')(y)
model = Model(inputs=x_input, outputs=output)
model.summary()

adam = Adam()
model.compile(loss=amsoftmax_loss,
optimizer=adam,
metrics=['accuracy'])

history = model.fit(X_train, Y_train,
batch_size=batch_size,
epochs=10,
verbose=1,
validation_data=(X_test, Y_test))

score = model.evaluate(X_test, Y_test, verbose=0)

print('Test score:', score[0])
print('Test accuracy:', score[1])