swift - 在 Swift 中优化 MLMultiArray 上的“嵌套 for 循环”

Question

我有一个嵌套的 for 循环（从这里开始）

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 编写代码会有帮助吗？怎么做到呢？

谢谢

Answer 1

Accelerate 框架具有优化的 argmax 函数。你也可以在 GitHub 上的 CoreMLHelpers 中找到它。