0

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

Willeke
  • 14,578
  • 4
  • 19
  • 47
jonb
  • 845
  • 1
  • 13
  • 36

1 Answers1

0

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

Matthijs Hollemans
  • 7,706
  • 2
  • 16
  • 23
  • Thanks! Where can I find this argmax function? It sounds exactly what I need. I don't see it here: https://github.com/hollance/CoreMLHelpers/tree/master/CoreMLHelpers. I see "argmax" for Array, but not for MLMultiArray. In addition, the access to the MLMultiArray is by *index*, even if I had argmax function, how do I use it in the given code? Maybe I can convert it to (14,50,60) array first? – jonb May 23 '21 at 17:38
  • You need to pass in the pointer to the MLMultiArray's contents. In your case, since the array is (14, 50, 60) you also need to pass in a stride of 50 * 60. – Matthijs Hollemans May 23 '21 at 20:11
  • Thanks for the comment. I am sorry, I don't understand the proposed solution for the question. If you post an answer and it will work - I will gladly accept it. Thanks :) – jonb May 23 '21 at 20:35
  • Perhaps in your case it would be easiest to add the argmax operation into the model itself. You can see an example here, from my book: https://github.com/hollance/coreml-survival-guide/blob/4dfcbb97c065726a3da240c55d90b2075959801d/Scripts/deeplab.py#L94-L109 – Matthijs Hollemans May 24 '21 at 10:34
  • I am trying to add it in an additional Keras Lambda layer. This is not a single argmax, but argmax for each one of the 14 heatmaps of size (50,60). As I understand in your example is a single argmax. Thanks – jonb May 24 '21 at 12:38
  • 1
    In the example it's an argmax over the 14 channels. So for every pixel in the 50x60 heatmap it tells you which one of the 14 classes has the highest probability. – Matthijs Hollemans May 24 '21 at 20:30