如何利用CUDA并行计算高维矩阵的曼哈顿距离？

问题背景 #

尝试用scipy.spatial.distance.cdist计算两个形状为[4,12,64,32]矩阵的曼哈顿距离，但cdist不支持3维以上的输入。手动实现的嵌套循环+CPU版cdist代码运行速度极慢，需要用CUDA并行加速。

原慢代码：

import torch
from scipy.spatial.distance import cdist
import numpy as np

T=4
A=torch.randn([T,12,64,32],dtype=torch.float)
B=torch.randn([T,12,64,32],dtype=torch.float)

def manhattan_dist(self, q, k):
    _, B, H, N, C = q.shape
    out = [[0 for i in range(H)] for i in range(B)]

    for b in range(0, B):
        for h in range(0, H):
            dist = cdist(q.squeeze().detach().cpu().numpy()[b][h], k.squeeze().detach().cpu().numpy()[b][h],
                         metric='cityblock')
            out[b][h] = dist

    out = torch.tensor(np.array(out)).unsqueeze(0).float()
    return out

慢代码的问题分析 #

• 串行循环：嵌套for循环逐个处理B和H维度，完全没有利用GPU的并行计算能力
• 数据传输开销：频繁将GPU张量转换为CPU numpy数组，再转回GPU张量，数据拷贝耗时巨大
• CPU计算限制：scipy.cdist是纯CPU实现，无法利用CUDA加速

CUDA并行实现方案 #

用PyTorch原生张量操作，通过广播机制实现全GPU并行计算，避免循环和数据转换：

import torch

class YourModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
    
    def manhattan_dist(self, q, k):
        # 确保输入张量形状一致，且在同一设备（CUDA/CPU）上
        assert q.shape == k.shape, "输入q和k的形状必须完全匹配"
        
        # 扩展维度实现广播：将q扩展为 [..., N, 1, C]，k扩展为 [..., 1, N, C]
        # 这样差运算会自动覆盖所有点对组合
        q_expanded = q.unsqueeze(-2)  # 在倒数第二维加一个维度
        k_expanded = k.unsqueeze(-3)  # 在倒数第三维加一个维度
        
        # 计算所有点对的绝对值差，再在特征维度（最后一维）求和得到曼哈顿距离
        dist = torch.abs(q_expanded - k_expanded).sum(dim=-1)
        
        # 若需匹配原代码输出形状（如[1, B, H, N, N]），可按需调整维度：
        # dist = dist.unsqueeze(0)
        
        return dist.float()

# 测试示例（全程在CUDA上运行）
T = 4
# 直接创建CUDA张量
A = torch.randn([T, 12, 64, 32], dtype=torch.float).cuda()
B = torch.randn([T, 12, 64, 32], dtype=torch.float).cuda()

# 模型移到CUDA
model = YourModel().cuda()
manhattan_dist = model.manhattan_dist(A, B)

# 输出形状为 [4, 12, 64, 64]，对应每个T/B维度下64个点的两两曼哈顿距离
print(manhattan_dist.shape)

方案优势 #

• 全GPU并行：所有操作基于PyTorch张量实现，自动利用CUDA核心并行计算，彻底摆脱串行循环
• 无数据拷贝：全程在GPU上完成计算，避免GPU与CPU之间的数据传输开销
• 代码简洁高效：通过广播机制自动处理所有维度的点对计算，无需手动遍历

内容的提问来源于stack exchange，提问作者Henry .H