cos類似度行列の実装【Pytorch, Tensorflow】

SimCLRなどの対照学習(Contrastive Learning)の手法で，特徴量空間における類似度の指標として用いられるものの一つにCos(コサイン)類似度があります．

Tensorflow，Pytorchそれぞれで実装を行ったので，メモ程度に記録しておきます．（参考までに）

Pytorch

# input_sizeは (batchsize*次元数)

def cosine_matrix(a, b):
    dot = torch.matmul(a, torch.t(b))
    norm = torch.matmul(torch.norm(a, dim=1).unsqueeze(-1), torch.norm(b, dim=0).unsqueeze(0))
    return dot / norm

Tensorflow

def cosine_matrix(a, b):
    a_normed, _ = tf.linalg.normalize(a, axis=-1)
    b_normed, _ = tf.linalg.normalize(b, axis=-1)
    matrix = tf.matmul(a_normed, b_normed, transpose_b=True)
    return matrix