[ArcFace] 推論結果がクラスラベルに依存しすぎてしまう

arcface_model = create_arcface_with_resnet50v2(input_shape, s=s_value, m=m_value) # モデル生成

predictions = arcface_model.predict([img, cls]) # softmaxの値を推論(画像(img)とそのクラスラベル(cls)を入力)

[1.0000000e+00 0.0000000e+00 0.0000000e+00
 2.3753611e-32 0.0000000e+00 0.0000000e+00 
 0.0000000e+00 0.0000000e+00 2.8887461e-28 
 0.0000000e+00 0.0000000e+00 1.3241292e-28 ]

[0.0000000e+00 1.0000000e+00 0.0000000e+00
 1.1876997e-18 0.0000000e+00 4.6460828e-18 
 0.0000000e+00 0.0000000e+00 3.9981019e-35 
 0.0000000e+00 0.0000000e+00 0.0000000e+00 ]

class Arcfacelayer(Layer):
    # s:softmaxの温度パラメータ, m:margin
    def __init__(self, output_dim, s, m, easy_margin=False):
        self.output_dim = output_dim
        self.s = s
        self.m = m
        self.easy_margin = easy_margin
        super(Arcfacelayer, self).__init__()
    
    # モデル保存に必要
    def get_config(self):
        config = {
            "output_dim" : self.output_dim,
            "s" : self.s,
            "m" : self.m,
            "easy_margin" : self.easy_margin
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))
    

    # 重みの作成
    def build(self, input_shape):
        # Create a trainable weight variable for this layer.
        self.kernel = self.add_weight(name='kernel',
                                        shape=(input_shape[0][1], self.output_dim),
                                        initializer='uniform',
                                        trainable=True)
        super(Arcfacelayer, self).build(input_shape)


    # mainの処理 
    def call(self, x):
        y = x[1]
        x_normalize = tf.math.l2_normalize(x[0]) # x = x'/ ||x'||2
        k_normalize = tf.math.l2_normalize(self.kernel) # Wj = Wj' / ||Wj'||2

        cos_m = K.cos(self.m)
        sin_m = K.sin(self.m)
        th = K.cos(np.pi - self.m)
        mm = K.sin(np.pi - self.m) * self.m

        cosine = K.dot(x_normalize, k_normalize) # W.Txの内積
        sine = K.sqrt(1.0 - K.square(cosine))

        phi = cosine * cos_m - sine * sin_m #cos(θ+m)の加法定理

        if self.easy_margin:
            phi = tf.where(cosine > 0, phi, cosine) 

        else:
            phi = tf.where(cosine > th, phi, cosine - mm) 

        # 正解クラス:cos(θ+m) 他のクラス:cosθ 
        output = (y * phi) + ((1.0 - y) * cosine) 
        output *= self.s

        return output

    def compute_output_shape(self, input_shape):
        return (input_shape[0][0], self.output_dim) #入力[x,y]のためx[0]はinput_shape[0][0]

# ResNet50v2 + ArcFace定義
# 学習に使用
def create_arcface_with_resnet50v2(input_shape, s, m):
    # ResNet50V2の入力層の前に独自の入力層を追加
    input_tensor = input_shape
    
    input_model = Sequential()
    input_model.add(InputLayer(input_shape=input_tensor))
    input_model.add(Conv2D(3, (7, 7), padding='same'))
    input_model.add(BatchNormalization())
    input_model.add(Activation('relu'))
    
    resnet50v2 = ResNet50V2(include_top=False, weights=None, input_tensor=input_model.output)

    resnet50v2.load_weights('save_model(weights_imagenet)/weights_imagenet.hdf5', by_name=True)

    flat = Flatten()(resnet50v2.layers[-1].output)
    dense = Dense(512, activation="relu", name="hidden")(flat)
    
    x = BatchNormalization()(dense)
    
    yinput = Input(shape=(num_classes,)) #ArcFaceで使用
    
    s_cos = Arcfacelayer(num_classes, s, m)([x,yinput]) #outputをクラス数と同じ数に

    prediction = Dense(num_classes, activation="softmax")(s_cos)

    model = Model(inputs=[resnet50v2.input,yinput], outputs=prediction)
    
    return model