PyTorch无需for循环基于核函数生成类Gram矩阵的实现方法咨询

1. lambda+map实现方法 #

你之前的map实现只能得到对角线结果，是因为map默认按位置逐元素配对输入，仅能覆盖i=j的组合。要生成全量N×N矩阵，需要先把输入张量展开为所有(i,j)两两配对的序列，再喂给map处理：

N = x.shape[0]
# 生成所有配对的a、b序列，长度均为N²，满足a[i*N +j] = x[i]，b[i*N +j] =x[j]
a = x.repeat_interleave(N, dim=0)
b = x.repeat(N, 1)
# 批量计算后reshape为N×N矩阵
G = torch.tensor(list(map(lambda a_, b_: K(a_, b_), a, b))).reshape(N, N)

该方案不需要修改核函数K的定义，也没有显式Python for循环。 #

2. 更高效率的替代方案 #

2.1 使用torch.vmap向量化（推荐） #

map方法本质还是在Python层面遍历序列，性能有限。用PyTorch原生的vmap（向量映射）工具可以在计算图层面实现向量化，完全规避Python层面的循环开销，性能提升非常明显：

from torch import vmap
# 嵌套vmap：先对x的每个元素a，再对x的每个元素b计算K(a,b)
G = vmap(lambda a: vmap(lambda b: K(a, b))(x))(x)

该方案代码更简洁，运行速度远快于map实现，同样不需要修改核函数定义。

2.2 特定核函数的直接矩阵优化 #

如果你使用的核函数可以拆解为矩阵运算（比如你示例的二阶多项式核），可以直接用矩阵运算实现，性能是最高的：

# 和你给出的二阶多项式核完全等价的实现，一行即可完成
G = (1 + x @ x.T) ** 2

注意事项 #

如果N的数值很大，N×N的矩阵会占用大量内存/显存，可以按需采用分块计算的方式降低峰值内存占用。 #

内容的提问来源于stack exchange，提问作者J Kim