ops.py 文件源码

def full_cosine_similarity(matrix1, matrix2):
    """
    Expect 2 matrices P and Q of dimension (d, n1) and (d, n2) respectively.
    Return a matrix A of dimension (n1, n2) with the result of comparing each
    vector to one another. A[i, j] represents the cosine similarity between
    vectors P[:, i] and Q[:, j].
    """
    n1 = matrix1.size(1)
    n2 = matrix2.size(1)
    d = matrix1.size(0)
    assert d == matrix2.size(0)

    # -> (d, n1, 1)
    t1 = matrix1.view(d, n1, 1)
    # -> (d, n1, n2)
    t1 = t1.repeat(1, 1, n2)

    # -> (d, 1, n2)
    t2 = matrix2.view(d, 1, n2)
    # -> (d, n1, n2)
    t2 = t2.repeat(1, n1, 1).contiguous()

    t1_x_t2 = torch.mul(t1, t2)  # (d, n1, n2)
    dotprod = torch.sum(t1_x_t2, 0).squeeze()  # (n1, n2)

    norm1 = torch.norm(t1, 2, 0)  # (n1, n2)
    norm2 = torch.norm(t2, 2, 0)  # (n1, n2)
    col_norm = torch.mul(norm1, norm2).squeeze()  # (n1, n2)

    return torch.div(dotprod, col_norm)  # (n1, n2)