Optimizing 'nested for loop' over MLMultiArray in Swift

Question

I have a nested for loop (from here)

heatmaps is an MLMultiArray at shape of (14, 50, 60). This code iterates 14 sub-arrays of shape (50,60), and finding the maximum for each value.

This is the code:

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)
                }
            }
        }
    }

Equivalent python code:

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

The Swift function takes around ~36ms with iPhone 12 mini, and ~58ms with iPhone 11 Pro Max. I want to optimize it and reduce it's runtime.

How can I do it?

Will it help if I write the code in Objective-C? How can it be done?

Thanks

Answer 1

Accelerate framework has an optimized argmax function. You can also find this in CoreMLHelpers on GitHub.