基于PyTorch批量计算度量矩阵U的GPU加速优化问题

需求说明 #

给定批量特征集F（形状为[batch_size, data_point_num, features_num]），需计算度量矩阵U，其中每个元素满足：

U[i,j] = (f_i - f_j)^T · M · (f_i - f_j)

M为features_num×features_num的固定矩阵。原使用torch.einsum实现，但GPU训练时速度较慢，尝试用torch.bmm重写时得到错误形状（目标形状为[batch_size, data_point_num, data_point_num]），需提供正确的向量化实现方案。

原einsum实现代码 #

import torch

batch_size = 10
data_point_num = 9
features_num = 3
F = torch.randn(batch_size, data_point_num, features_num)
M = torch.randn(features_num, features_num)

F_expanded = F[:, :, None, :]
diff = F_expanded - F[:, None, :, :]
U_einsum = torch.einsum('bijk,kl,bijl->bij', diff, M, diff)

错误实现分析 #

你尝试的代码将diff重塑为[batch_size, data_point_num², features_num]后做bmm，导致输出形状变为[10,81,81]，原因是这种方式把所有(f_i-f_j)向量视为独立样本，计算了它们之间的两两矩阵乘积，而非每个(f_i-f_j)自身与M的二次型结果。

正确向量化实现方案 #

方案1：利用广播与逐元素乘积求和 #

无需调整维度，直接通过广播实现矩阵乘法，最后在特征维度求和得到标量结果，代码简洁且高效：

# 计算特征差，形状[batch_size, data_point_num, data_point_num, features_num]
diff = F.unsqueeze(2) - F.unsqueeze(1)
# 先计算diff与M的矩阵乘法，再和原diff做逐元素乘积，最后在特征维度求和
U = torch.sum((diff @ M) * diff, dim=-1)

方案2：利用bmm批量计算 #

通过维度重组，将每个批次内的向量组批量传入bmm计算，适合对矩阵乘法优化更敏感的场景：

# 计算特征差，形状[batch_size, data_point_num, data_point_num, features_num]
diff = F.unsqueeze(2) - F.unsqueeze(1)
# 调整diff形状为[batch_size * data_point_num, data_point_num, features_num]
diff_reshaped = diff.reshape(-1, data_point_num, features_num)
# 扩展M形状以匹配batch维度，形状[batch_size * data_point_num, features_num, features_num]
M_reshaped = M.unsqueeze(0).repeat(diff_reshaped.shape[0], 1, 1)
# 批量计算矩阵乘法：(f_i-f_j) · M
diff_M = torch.bmm(diff_reshaped, M_reshaped)
# 批量计算点积并提取对角元素（对应每个(i,j)的二次型结果），最后调整回原形状
U = torch.bmm(diff_M, diff_reshaped.transpose(1, 2)).diagonal(dim1=1, dim2=2).reshape(batch_size, data_point_num, data_point_num)

正确性验证 #

可以通过对比两种实现与原einsum结果的误差确认正确性：

# 方案1结果对比
print(torch.allclose(U, U_einsum, atol=1e-6))  # 应输出True

# 方案2结果对比
print(torch.allclose(U, U_einsum, atol=1e-6))  # 应输出True

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