Dask Distributed高维数据距离矩阵计算异常排查

你的代码返回长度不符合预期是多个逻辑错误共同导致的，具体问题和修正方式如下：

现存核心bug #

• 笔误：主执行逻辑中调用的函数名为similarity_matrix，但你实际定义的距离计算函数是distance_matrix，运行时会直接抛出名称错误。
• 距离计算函数逻辑错误：
  1. map_blocks传入的u、v是形状为(块大小,20,10)的分块数组，不是单个实例，你当前的循环索引逻辑完全不匹配分块维度。
  2. Python原生列表没有apply方法，追加元素应该用append，现有代码运行到这一步会直接抛出属性错误。
  3. 距离计算的循环层级错误：两个实例的距离是各自20个向量两两L1距离的均值，共400个值取平均，你当前的循环嵌套到了向量内部的标量维度，计算逻辑完全错误。
• Dask调度逻辑错误：
  1. 给map_blocks传入两个相同的X时，Dask默认按块索引对齐计算，只会计算相同索引分块内部的实例距离，跨分块的实例对（比如第0块和第1块的实例）完全不会被计算，不可能得到完整的1000*1000距离矩阵。
  2. 你设置drop_axis=[1,2]强行把输出压成一维，且chunks配置不符合二维距离矩阵的结构，最终拼接结果的长度必然不符合预期。
• 性能问题：三层Python循环就算逻辑修正，运行速度也极慢，完全发挥不了Dask并行的优势。

修正方案 #

不要在分块内手动维护下三角的一维索引，直接用广播实现向量化的块距离计算，通过Dask的blockwise接口生成所有分块对的计算任务，最后统一提取下三角即可，逻辑简单不易出错。

修正后完整代码 #

import numpy as np
import dask.array as da
from dask.distributed import Client

def block_dist(A, B):
    # 输入两个分块：A形状(m,20,10)、B形状(n,20,10)
    # 输出m*n的距离分块，每个位置对应两个实例的平均L1距离
    abs_diff = np.abs(A[:, None, :, None, :] - B[None, :, None, :, :])
    l1_dist = abs_diff.sum(axis=-1)
    return l1_dist.mean(axis=(1, 3))

def distance_matrix(X, chunk_size=100):
    # 按实例维度分块，向量维度不分块
    X_da = da.from_array(X, chunks=(chunk_size, 20, 10))
    # 自动调度所有分块对的距离计算
    dist_matrix = da.blockwise(
        block_dist,
        "ij",
        X_da, "ik",
        X_da, "jk",
        dtype=float
    )
    # 提取下三角矩阵（包含对角线）
    return da.tril(dist_matrix).compute()

if __name__ == '__main__':
    workers = 6
    # 生成测试数据
    X = np.random.rand(1000, 20, 10)
    client = Client(n_workers=workers)
    results = distance_matrix(X)
    client.close()
    # 验证输出形状为(1000,1000)
    print("结果形状:", results.shape)
    print(results)

补充说明 #

• 分块大小chunk_size可以根据机器内存调整，默认100的分块下每个分块的距离矩阵仅100*100=10000个浮点数，内存占用极低，6核CPU下运行时间在秒级。
• 如果你需要长度为1000*1001/2的一维下三角展平结果，可以在计算完成后通过results[np.tril_indices(1000)]提取，比在分块内手动计算索引的容错率高很多。

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