如何按指定窗口计算两个PyTorch嵌入向量的相似度？

我现有两个维度为32（可扩展至64等）的PyTorch嵌入向量A1和B1，已通过unsqueeze处理为列向量，代码如下：

import torch

A = torch.rand(32)
A1 = torch.unsqueeze(A, dim = 1)
A1

B = torch.rand(32)
B1 = torch.unsqueeze(B, dim = 1)
B1

需要按指定窗口大小（如4、8、16）对两个向量分段，分别计算每段的相似度，得到长度与分段数匹配的结果向量。例如窗口大小为8时，32维向量会被分为4段，最终输出含4个相似度值的向量：

ExampleOutput = [.4, .5, .6, .7] # 窗口大小为8时的输出示例

已知可通过PyTorch的直方图相关方法实现，求具体实现方案。

核心思路 #

利用torch.histogram生成分段索引，将原向量按窗口大小划分到不同区间，再对每个区间内的子向量计算相似度。这里以余弦相似度为例，你也可以替换为点积、欧氏距离归一化等其他相似度计算方式。

完整代码实现 #

import torch
import torch.nn.functional as F

# 初始化列向量
A = torch.rand(32)
A1 = torch.unsqueeze(A, dim=1)  # 形状 (32, 1)
B = torch.rand(32)
B1 = torch.unsqueeze(B, dim=1)  # 形状 (32, 1)

# 配置窗口大小
window_size = 8
# 校验窗口大小是否能整除向量长度，避免分段不均
assert A1.shape[0] % window_size == 0, "窗口大小必须能整除向量的维度长度"
num_chunks = A1.shape[0] // window_size

# 1. 用torch.histogram生成分段索引
# 创建分段的边界，例如32维向量+窗口8的边界为[0,8,16,24,32]
bins = torch.arange(0, A1.shape[0] + 1, window_size)
# 生成每个元素对应的区间索引
indices = torch.arange(A1.shape[0])
_, bin_indices = torch.histogram(indices, bins=bins)
# 转换为0-based索引
bin_indices = bin_indices - 1

# 2. 按分段索引计算每段的相似度
similarity_results = []
for idx in range(num_chunks):
    # 提取当前分段的子向量
    a_segment = A1[bin_indices == idx].squeeze(dim=1)
    b_segment = B1[bin_indices == idx].squeeze(dim=1)
    # 计算余弦相似度（可替换为其他相似度计算逻辑）
    segment_sim = F.cosine_similarity(a_segment, b_segment, dim=0)
    similarity_results.append(segment_sim.item())

# 转换为PyTorch张量或保持列表形式
result = torch.tensor(similarity_results)
print(result)
# 输出示例：tensor([0.4215, 0.5367, 0.6023, 0.7109])

说明 #

1. 分段逻辑：通过torch.histogram对向量元素的位置索引进行区间划分，确保每个分段的长度严格等于窗口大小。
2. 相似度替换：如果需要其他相似度，可替换F.cosine_similarity部分，例如用点积的话可以写为torch.dot(a_segment, b_segment) / (torch.norm(a_segment) * torch.norm(b_segment))，效果和余弦相似度一致；若用原始点积则直接写torch.dot(a_segment, b_segment)即可。
3. 扩展性：只要窗口大小能整除向量维度（如64维向量用窗口16，会分成4段），代码无需修改即可直接适配。

内容的提问来源于stack exchange，提问作者Droid-Bird