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

Posted 2023-03-27

【中文标题】在 Swift 中优化 MLMultiArray 上的“嵌套循环”

【英文标题】：Optimizing 'nested for loop' over MLMultiArray in Swift

【发布时间】：2021-08-11 21:37:06

【问题描述】：

我有一个嵌套的 for 循环（来自 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 编写代码会有帮助吗？怎么办？

谢谢

【参考方案1】：

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

【讨论】：

谢谢！我在哪里可以找到这个 argmax 函数？这听起来正是我需要的。我在这里看不到：github.com/hollance/CoreMLHelpers/tree/master/CoreMLHelpers。我看到 Array 的“argmax”，但没有看到 MLMultiArray。另外，对 MLMultiArray 的访问是通过 index，即使我有 argmax 函数，在给定的代码中如何使用它？也许我可以先将其转换为 (14,50,60) 数组？

您需要传入指向 MLMultiArray 内容的指针。在您的情况下，由于数组是 (14, 50, 60) 您还需要传入 50 * 60 的步幅。

感谢您的评论。对不起，我不明白该问题的建议解决方案。如果您发布答案并且它会起作用 - 我很乐意接受它。谢谢:)

也许在您的情况下，将 argmax 操作添加到模型本身是最简单的。你可以在这里看到一个例子，来自我的书：github.com/hollance/coreml-survival-guide/blob/…

在示例中，它是 14 个通道上的 argmax。因此，对于 50x60 热图中的每个像素，它都会告诉您 14 个类别中哪一个类别的概率最高。