在 iOS 上释放 ML 力量：Apple Silicon 优化秘诀

Core ML 是 Apple 的机器学习框架，它引入了异步和批处理预测的新 API。这些强大的功能可以在处理机器学习模型时显著提高应用的性能。本文将探讨如何实现这些新 API 以及它们为 Core ML 集成带来的好处。更多详细信息，请查看 WWDC 2023 讲座：通过异步预测改进 Core ML 集成。

理解异步预测

Core ML 中新的异步预测 API 旨在与 Swift 并发无缝协作。它通过实现并发预测并与其他异步代码良好集成，从而提高应用的性能和响应能力。

异步预测的关键优势：

• 线程安全：无需手动同步
• 支持取消：遵循任务取消，更好地管理资源
• 提高吞吐量：允许并发预测

实现异步预测

我们来看看如何实现异步预测

class ColorizingService {
    private let colorizerModel: colorizer_model

    init() throws {
        let config = MLModelConfiguration()
        self.colorizerModel = try colorizer_model(configuration: config)
    }

    func colorize(_ image: CGImage) async throws -> CGImage {
        // Check for cancellation
        try Task.checkCancellation()

        // Prepare input (non-isolated work)
        let input = try colorizationInputFrom(cgImage: image)

        // Perform prediction
        let output = try await colorizerModel.prediction(input: input)

        // Process output
        return try cgImageFrom(output: output)
    }
}

在这个例子中，我们将 `colorize` 方法设置为异步。`prediction` 调用前加上了 `await` 关键字，使其能够与 Swift 的并发系统一起使用。我们还在方法开始时添加了一个取消检查，以更好地进行任务管理。

批处理预测

虽然异步预测对于随时间处理单个输入非常有用，但当您同时有多个输入可用时，批处理预测是理想的选择。

何时使用批处理预测：

• 当您有一组固定的输入要处理时
• 当您想利用 Core ML 内部优化来处理多个输入时
• 当您不需要取消批处理中的单个预测时

下面是如何实现批处理预测的示例

func colorizeBatch(_ images: [CGImage]) throws -> [CGImage] {
    let inputs = try images.map { try colorizationInputFrom(cgImage: $0) }
    let outputs = try colorizerModel.predictions(inputs: inputs)
    return try outputs.map { try cgImageFrom(output: $0) }
}

案例研究：CLIP-Finder

我们来探讨这些概念如何在实际项目中应用：CLIP-Finder。CLIP-Finder 是一款 iOS 应用，允许用户使用自然语言描述或摄像头输入搜索他们的照片库。它使用针对神经网络引擎优化的 Core ML 模型对照片执行语义搜索。

CLIP-Finder 中的异步预测

CLIP-Finder 在其 `CLIPImageModel` 类中实现了异步预测。

class CLIPImageModel {
    var model: MLModel?

    func performInference(_ pixelBuffer: CVPixelBuffer) async throws -> MLMultiArray? {
        guard let model = model else {
            throw NSError(domain: "ClipImageModel", code: 2, userInfo: [NSLocalizedDescriptionKey: "Model is not loaded"])
        }

        let input = InputFeatureProvider(pixelBuffer: pixelBuffer)

        do {
            let outputFeatures = try await model.prediction(from: input)

            if let multiArray = outputFeatures.featureValue(for: "var_1259")?.multiArrayValue {
                return multiArray
            } else {
                throw NSError(domain: "ClipImageModel", code: 3, userInfo: [NSLocalizedDescriptionKey: "Failed to retrieve MLMultiArray from prediction"])
            }
        } catch {
            print("ClipImageModel: Failed to perform inference: \(error)")
            throw error
        }
    }
}

CLIP-Finder 中的 Turbo 模式

CLIP-Finder 利用异步预测的 Turbo 模式实现更快的图像处理

• 通过相机界面中的按钮激活
• 启用异步相机预测以加快图像搜索
• 通过点击相机界面右下角的“Turbo”按钮激活

CLIP-Finder 中的批处理

CLIP-Finder 使用批处理来预处理照片并将其特征向量存储在 CoreData 数据库中

private func processAndCachePhotos(_ assetsToProcess: [PHAsset]) async {
    // ... (setup code)

    let batchSize = 512
    let batches = stride(from: 0, to: assetsToProcess.count, by: batchSize).map {
        Array(assetsToProcess[$0..<min($0 + batchSize, assetsToProcess.count)])
    }

    for batch in batches {
        let results = await withTaskGroup(of: (String, CVPixelBuffer?).self) { group in
            for asset in batch {
                group.addTask {
                    let identifier = asset.localIdentifier
                    guard let image = await self.requestImage(for: asset, targetSize: targetSize, options: options) else {
                        return (identifier, nil)
                    }
                    let pixelBuffer = Preprocessing.preprocessImageWithCoreImage(image, targetSize: targetSize)
                    return (identifier, pixelBuffer)
                }
            }

            var batchResults: [(String, CVPixelBuffer?)] = []
            for await result in group {
                batchResults.append(result)
            }
            return batchResults
        }

        // ... (process batch results)
    }

    // ... (finalize processing)
}

这种批处理方法具有以下优势

• 高效的初始处理：使用神经网络引擎预处理整个照片库，批处理大小为 512 张照片
• 增量更新：后续应用启动时仅处理新图像
• 数据库维护：当照片从设备中删除时更新数据库

性能分析

同步预测与异步预测

为了理解异步预测的优势，我们将其与同步预测进行比较

同步预测

在同步预测中，任务按顺序执行。正如我们在图中看到的，连续预测之间存在明显的间隙。这种方法可能导致系统资源利用效率低下，并可能降低整体性能，尤其是在处理多个预测时。

异步预测

通过异步预测，图像预处理和预测任务是异步分派的。这允许沿时间线优化执行，最大限度地提高了多个 CoreML 实例在神经网络引擎上同时运行的可能性。如图所示，任务重叠并更有效地利用系统资源，从而提高整体性能。

批处理性能

此 Instruments 跟踪提供了批处理性能的洞察

• 图像预处理 GPU：图像预处理的两个不同 GPU 活动阶段，可能涉及图像大小调整和归一化。
• CLIP 图像预测：神经网络引擎密集活动，批处理大小为 512 张照片，每张在 A17 Pro 芯片上处理大约需要 1 毫秒。
• 神经网络引擎预测：额外的神经网络引擎批处理活动，可能用于进一步处理或特征提取。
• 高效资源利用：跟踪显示了 CLIP-Finder 如何高效利用 GPU 和神经网络引擎，实现并行处理。

详细批处理

批处理的这个详细视图揭示了一个重要的优化：两个预测实例并行发生，显著提高了性能。通过利用神经网络引擎的功能，CLIP-Finder 可以同时处理多个图像，大大减少了批处理操作所需的总时间。

性能比较：同步 vs 异步 vs 批处理

为了更好地理解不同处理模式的性能优势，我们使用 CLIP-Finder 对包含 6,524 张图像的照片库进行了测试。结果清楚地表明了批处理优于同步和异步方法的优势。

结果表明，批处理显著优于同步和异步方法，其中异步处理比同步处理有显著改进。

同步处理

在同步模式下，整个图库的处理大约需要 40.4 秒，平均每张照片 6.19 毫秒。时间线显示了顺序处理，操作之间存在明显的间隙。

异步处理

异步处理显著提高了性能，整个图库的处理大约需要 16.28 秒，平均每张照片 2.49 毫秒。时间线显示操作重叠更多，神经网络引擎和 GPU 的利用率提高，与同步处理相比，资源利用率更高。

批处理

批处理表现出最佳性能，仅需 15.29 秒即可完成图库处理，平均每张照片 2.34 毫秒。时间线显示神经网络引擎密集并行利用，从而实现了最有效的处理方法。

这些结果清楚地说明了批处理的性能优势，特别是对于照片库这样的大型数据集。利用神经网络引擎同时处理多个图像的能力可以显著节省时间并更有效地利用资源。

异步和批处理预测的选择

异步和批处理预测之间的选择取决于您的具体用例

• 在以下情况下使用异步预测：
  • 您处于异步上下文
  • 输入随时间逐个可用
  • 您需要支持取消单个预测
• 在以下情况下使用批处理预测：
  • 您有已知数量的工作要完成
  • 所有输入同时可用
  • 您不需要取消批处理中的单个预测

性能考量

虽然异步和批处理预测可以显著提高性能，但请记住：

• 分析您的工作负载以确保其受益于并发性
• 注意内存使用情况，尤其是在并发处理大输入或输出时
• 如果内存使用是问题，请考虑实施流控制以限制并发预测的数量

结论

Core ML 中新的异步预测 API 和批处理功能提供了强大的方法，可以将机器学习集成到您的 Swift 应用程序中，从而提高性能和响应能力。通过了解何时使用每种方法以及如何有效地实现它们，您可以优化应用程序的 Core ML 集成，以获得最佳的用户体验。CLIP-Finder 项目展示了这些技术如何在实际场景中应用，利用异步和批处理来创建响应迅速且高效的照片搜索应用程序。

参考文献

• CLIP-Finder GitHub 存储库
• WWDC 2023 讲座：通过异步预测改进 Core ML 集成