通过手动计算相似度并使用矩阵乘法+转置：

import torch
from scipy import spatial
import numpy as np

a = torch.randn(2, 2)
b = torch.randn(3, 2) # different row number, for the fun

# Given that cos_sim(u, v) = dot(u, v) / (norm(u) * norm(v))
#                          = dot(u / norm(u), v / norm(v))
# We fist normalize the rows, before computing their dot products via transposition:
a_norm = a / a.norm(dim=1)[:, None]
b_norm = b / b.norm(dim=1)[:, None]
res = torch.mm(a_norm, b_norm.transpose(0,1))
print(res)
#  0.9978 -0.9986 -0.9985
# -0.8629  0.9172  0.9172

# -------
# Let's verify with numpy/scipy if our computations are correct:
a_n = a.numpy()
b_n = b.numpy()
res_n = np.zeros((2, 3))
for i in range(2):
    for j in range(3):
        # cos_sim(u, v) = 1 - cos_dist(u, v)
        res_n[i, j] = 1 - spatial.distance.cosine(a_n[i], b_n[j])
print(res_n)
# [[ 0.9978022  -0.99855876 -0.99854881]
#  [-0.86285472  0.91716063  0.9172349 ]]

根据benjaminplanche的答案，增加eps以获得数值稳定性：

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt

与Zhang Yu的答案相同，但使用clamp而不是max，并且没有创建新的Tensor。我用timeit做了一个小测试，表明clamp更快，尽管我不精通使用该工具。

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.clamp(a_n, min=eps)
    b_norm = b / torch.clamp(b_n, min=eps)
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt

在计算矩阵的行向量和列向量的相似度时，不需要使用循环。下面是一个例子。

import torch as t
a = t.randn(2,4)
print(a)

# step 1. 计算行向量的长度
len_a = t.sqrt(t.sum(a**2,dim=-1))
print(len_a)

b = len_a.unsqueeze(1).expand(-1,2)
c = len_a.expand(2,-1)
# print(b)
# print(c)

# step2. 计算乘积
x = a @ a.T
print(x)

# step3. 计算最后的结果
res = x/(b*c)
print(res)

您可以展开2个输入批次，执行成对余弦相似性运算，然后转置它：
使用torch.repeat_interleave和torch.repeat的非克隆等同物。

def distance_matrix(x, y, distance_function):
    return distance_function(
        x.view(x.size(0), 1, x.size(1)).expand(x.size(0), y.size(0), x.size(1)).contiguous().view(-1, x.size(1)),
        y.unsqueeze(0).expand(y.size(0), x.size(0), x.size(1)).flatten(end_dim=1),
    ).view(x.size(0), y.size(0))

from torch.nn import functional as F

distance_matrix(x, y, F.cosine_similarity)