Tune your Core ML models
Machine Learning & AI 进阶 24m

调优你的 Core ML 模型

Tune your Core ML models

2021年6月11日

在 Apple 官方观看视频

一句话判断

Core ML 模型的性能调优不只是换小模型——compute unit 选择、模型量化(Quantization)和批量预测优化这三个维度同时调才能榨干最后一滴性能。

这场 Session 讲了什么

很多开发者在集成 Core ML 模型后就「用起来就行了」,忽略了性能优化的巨大空间。这场 Session 系统性地讲解了 Core ML 模型调优的三个方向。

Compute Unit 选择是第一步。Core ML 支持 CPU、GPU 和 Neural Engine 三种计算单元。默认情况下,Core ML 会自动选择,但你可以通过 MLModelConfiguration 手动指定。不同模型在不同计算单元上的性能差异可能高达 10 倍。

模型量化是压缩模型大小的有效手段。Core ML 支持从 Float32 量化到 Float16 和 Int8。Float16 量化几乎没有精度损失但体积减半,Int8 量化体积缩小到 1/4 但精度会明显下降。Session 展示了如何用 coremltools 在 Python 端完成量化。

批量预测优化适用于需要一次处理多条输入的场景。用 MLBatchProvider 替代单条预测可以显著提高吞吐量,特别是在 GPU 和 ANE 上。

值得深挖的点

Compute Unit 的选择不是一成不变的

同一个模型的不同层可能在不同计算单元上性能最优。Core ML 的 All(默认)模式会在加载时分析模型结构,为每层选择最佳计算单元。但在某些情况下,这个自动选择并不理想——比如一个模型在 ANE 上运行比 GPU 慢,因为 ANE 不支持某些操作,需要回退到 CPU。

Session 建议在 app 启动时做一次 benchmark:用同一组测试数据分别在 CPU only、GPU only 和 ANE only 模式下运行模型,记录耗时和内存占用,然后选择最优配置并缓存结果。不同设备上最优配置可能不同——A15 芯片的 ANE 可能比 A14 快 30%。

量化的精度评估方法

不要盲目量化。Session 推荐的流程是:先用 Float16 量化,跑一遍测试集看精度下降是否可接受。如果可接受,再尝试混合量化(部分层 Float16、部分层 Int8)。coremltools 提供了 quantization_utils 模块,可以在量化前后对比模型的输出差异。

代码片段

// 比较不同 Compute Unit 的性能
import CoreML

func benchmarkModel() async {
    let modelURL = Bundle.main.url(forResource: "ImageClassifier", 
                                    withExtension: "mlmodelc")!
    
    let configurations: [(String, MLComputeUnits)] = [
        ("CPU Only", .cpuOnly),
        ("GPU Only", .gpuOnly),
        ("Neural Engine", .neuralEngine),
        ("All (Auto)", .all)
    ]
    
    for (name, computeUnits) in configurations {
        let config = MLModelConfiguration()
        config.computeUnits = computeUnits
        
        let model = try! MLModel(contentsOf: modelURL, configuration: config)
        
        // 预热模型
        let sampleInput = try! MLDictionaryFeatureProvider(
            dictionary: ["image": MLMultiArray(shape: [1, 224, 224, 3], 
                                                dataType: .float32)]
        )
        _ = try? model.prediction(from: sampleInput)
        
        // 计时 100 次推理
        let start = CFAbsoluteTimeGetCurrent()
        for _ in 0..<100 {
            _ = try? model.prediction(from: sampleInput)
        }
        let elapsed = CFAbsoluteTimeGetCurrent() - start
        
        print("\(name): \(elapsed / 100 * 1000)ms/次")
    }
}
# 在 Python 中使用 coremltools 量化模型
import coremltools as ct

# 加载原始模型
model = ct.models.MLModel('ImageClassifier.mlmodel')

# Float16 量化(推荐首选,几乎无精度损失)
model_fp16 = ct.models.neural_network.quantization_utils.quantize_weights(
    model, mode='linear', dtype=np.float16
)
model_fp16.save('ImageClassifier_fp16.mlmodel')

# 混合量化:对精度敏感的层保持 Float16,其余用 Int8
from coremltools.optimize.torch.quantization import LinearQuantizer

config = {
    "global_config": {
        "weight_dtype": "int8",
        "granularity": "per_channel"
    },
    "module_type_configs": {
        "Conv": {"weight_dtype": "float16"}  # 卷积层保持高精度
    }
}
// 使用批量预测提高吞吐量
func batchPredict(images: [CGImage]) throws -> [Int] {
    let model = try! VNCoreMLModel(for: ImageClassifier().model)
    let request = VNCoreMLRequest(model: model)
    
    // 使用 MLBatchProvider 批量处理
    let batchInputs = images.map { image -> MLFeatureProvider in
        let featureValue = try! MLFeatureValue(
            pixelBuffer: image.toCVPixelBuffer()!
        )
        return MLDictionaryFeatureProvider(
            dictionary: ["image": featureValue]
        )
    }
    
    let batch = MLArrayBatchProvider(array: batchInputs)
    let modelObj = try! MLModel(contentsOf: model.modelURL)
    let results = try modelObj.predictions(from: batch)
    
    // 提取所有预测结果
    var predictions: [Int] = []
    for i in 0..<results.count {
        let result = results.features(at: i)
        let label = result.featureValue(for: "classLabel")!.int64Value
        predictions.append(Int(label))
    }
    return predictions
}

最佳实践

在 app 中缓存 benchmark 结果。用 ProcessInfo.processInfo.physicalMemorysysctl 获取设备信息,针对不同设备选择不同的 compute unit 配置。不要每次启动都重新 benchmark——一次耗时可能好几秒。

如果你的模型需要频繁更新(比如从服务端下载新版本),在下载后立即对新模型做一次 benchmark,把结果和模型文件一起缓存。这样下次加载时直接用最优配置。

还有什么值得关注

  • coremltools 5.0 支持 PyTorch 和 TensorFlow 模型的直接转换,中间不再需要 ONNX 格式。
  • Core ML 的 prediction(from:options:) 支持 usesCPUOnly 选项,在后台任务中避免抢占 GPU 资源。
  • Xcode 的 Model Performance Dashboard 可以可视化不同 compute unit 的性能对比,不需要自己写 benchmark 代码。
WWDC 2021