找不到有效的 numpy 广播

Question

这可能是一个简单的问题，但我找不到任何实用的解决方案。我的代码有以下涉及 3 个嵌套 for 循环的代码段。目标是为我的算法为prediction 和ground_truth 图像矩阵创建一个专门的强度矩阵，如下所示：

    for i in range (batch):
      for j in range (img_width):
        for k in range (img_height):
            tensor=prediction[i][j][:]-prediction[i][k][:]
            extracted_intensity_pred[i][j][k]=torch.norm(tensor, 2)
            tensor=ground truth[i][j][:]-ground_truth[i][k][:]
            extracted_intensity_ground_truth[i][j][k]=torch.norm(tensor, 2)

这种嵌套的 for 循环结构大大降低了执行速度。是否有任何可以使用的广播实现（基于 numpy 或 pytorch 张量）？

Answer 1

首先让我们清理一些符号； [:] 什么都不做

但首先尺寸是多少，主要是 3d？

for i in range (batch):
      for j in range (img_width):
        for k in range (img_height):
            tensor = prediction[i,j,:] - prediction[i,k,:]
            # looks like a prediction[:,:,None]-prediction[:,None,:]; making 4d?
            extracted_intensity_pred[i,j,k] = torch.norm(tensor, 2)
            # what can torch.norm work with?

也许只是

 tensor = prediction[:,:,None] - prediction(:,None,:]
 extracted_intensity_pred = torch.norm(tensor, ?)