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

SimCLRなどの対照学習(Contrastive Learning)の手法で，特徴量空間における類似度の指標として用いられるものの一つにCos(コサイン)類似度があります． Tensorflow，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 def cosine_matrix(a, b): a_normed, _ = tf.linalg.normalize(a