Keras-DSSM之in-batch余弦相似度负采样层

定义余弦相似度层，并在batch内进行负采样

NEG, batch_size = 20, 128

class NegativeCosineLayer():
    """ 自定义batch内负采样并做cosine相似度的层 """
    def __call__(self, inputs):
        def _cosine(x):
            query_encoder, doc_encoder = x
            doc_encoder_fd = doc_encoder
            for i in range(NEG):
                ss = tf.gather(doc_encoder, tf.random.shuffle(tf.range(tf.shape(doc_encoder)[0])))
                doc_encoder_fd = tf.concat([doc_encoder_fd, ss], axis=0)
            query_norm = tf.tile(tf.sqrt(tf.reduce_sum(tf.square(query_encoder), axis=1, keepdims=True)),[NEG + 1, 1])
            doc_norm = tf.sqrt(tf.reduce_sum(tf.square(doc_encoder_fd), axis=1, keepdims=True))
            query_encoder_fd = tf.tile(query_encoder, [NEG + 1, 1])
            prod = tf.reduce_sum(tf.multiply(query_encoder_fd, doc_encoder_fd, name="sim-multiply"), axis