Swift中MLMultiArray上“嵌套循环”的优化

我有一个嵌套的for循环(from here )

heatmaps 是形状为(14,50,60)的MLMultiArray。 此代码迭代14个子数组的形状(50,60),并为每个值找到最大值。

这是代码:

for k in 0..<keypoint_number {
        for i in 0..<heatmap_w {
            for j in 0..<heatmap_h {
                let index = k*(heatmap_w*heatmap_h) + i*(heatmap_h) + j
                let confidence = heatmaps[index].doubleValue
                guard confidence > 0 else { continue }
                if n_kpoints[k] == nil ||
                    (n_kpoints[k] != nil && n_kpoints[k]!.maxConfidence < confidence) {
                    n_kpoints[k] = PredictedPoint(maxPoint: CGPoint(x: CGFloat(j), y: CGFloat(i)), maxConfidence: confidence)
                }
            }
        }
    }

等效python代码:

for p_ind in range(n_keypoints):
    heat = heatmaps[0, p_ind, :, :]
    ind = np.unravel_index(np.argmax(heat), heat.shape)

Swift功能在iPhone 12 mini上大约需要36毫秒,在iPhone 11 Pro Max上大约需要58毫秒。 我想优化它并减少它的运行时间。

我该怎么做?

如果我用Objective-C编写代码会有帮助吗?如何做到这一点?

谢谢