Accelerate machine learning with Metal Performance Shaders Graph
Machine Learning & AI 高级 25m

用 Metal Performance Shaders Graph 加速机器学习

Accelerate machine learning with Metal Performance Shaders Graph

2021年6月9日

在 Apple 官方观看视频

一句话判断

MPSGraph 已经从”Metal 上的神经网络推理引擎”进化成了 Apple Silicon 上的通用计算图框架——如果你在做自定义的 ML 推理管线，它比直接写 Metal compute shader 省几个数量级的开发时间。

这场 Session 讲了什么

Session 介绍了 Metal Performance Shaders Graph（MPSGraph）在 iOS 15 和 macOS Monterey 中的重大更新。MPSGraph 是 Apple 提供的基于计算图（Computational Graph）的 GPU 加速框架，支持构建和执行复杂的数学运算图。

今年的重点改进包括：更丰富的算子（operator）库（新增了控制流、排序、稀疏操作等）、与 Core ML 的更深整合、以及训练（training）能力的增强。过去 MPSGraph 主要用于推理（inference），今年新增了对反向传播（backpropagation）的原生支持——你可以用 MPSGraph 构建前向图和梯度图，直接在 GPU 上训练自定义模型。Session 还展示了如何将 PyTorch 或 TensorFlow 的模型导出为 ONNX 格式，再通过 MPSGraph 加载并在 Apple GPU 上运行。

值得深挖的点

自动梯度计算（Automatic Gradient Computation）

MPSGraph 现在可以为任何前向计算图自动生成对应的梯度图。通过 gradient(for:) 方法，你指定需要求梯度的输入和损失函数输出，MPSGraph 自动构建反向传播图。它支持链式法则自动求导、高阶梯度（梯度的梯度），并且可以和自定义算子配合使用。这意味着你不需要手写反向传播代码——GPU 上训练自定义模型变得和推理一样简单。

MPSGraph 与 Core ML 的协作模式

Core ML 适合标准的模型推理（图像分类、目标检测等），但某些自定义算子 Core ML 不支持。现在你可以在 Core ML 模型中嵌入 MPSGraph 的自定义计算：Core ML 负责标准的卷积/全连接层，MPSGraph 负责自定义算子。两者共享同一个 GPU command buffer，没有中间数据拷贝。这种”Core ML 为主、MPSGraph 为辅”的架构是目前 Apple 平台上最灵活的推理方案。

代码片段

// 构建一个简单的神经网络推理图
func buildInferenceGraph(device: MTLDevice) -> MPSGraph {
    let graph = MPSGraph()
    
    // 输入占位符
    let input = MPSGraphTensor(
        shape: [1, 224, 224, 3], dataType: .float32)
    
    // 卷积层
    let weights = graph.constant(
        randomTensorWithShape: [3, 3, 3, 32], dataType: .float32)
    let conv = graph.convolution2D(
        input, weights: weights,
        strides: [1, 1], dilationRates: [1, 1],
        padding: .same, name: "conv1")
    
    // ReLU 激活
    let relu = graph.reLU(conv, name: "relu1")
    
    // 全局平均池化
    let pooled = graph.reduceMean(relu, axes: [1, 2], name: "gap")
    
    // 全连接层
    let fcWeights = graph.constant(
        randomTensorWithShape: [32, 10], dataType: .float32)
    let output = graph.matrixMultiplication(
        primary: pooled, secondary: fcWeights, name: "fc")
    
    return graph
}
// 使用自动梯度功能训练模型
func buildTrainingGraph(device: MTLDevice) -> (MPSGraph, MPSGraphTensor, MPSGraphTensor) {
    let graph = MPSGraph()
    
    // 前向图
    let input = MPSGraphTensor(shape: [32, 784], dataType: .float32)
    let label = MPSGraphTensor(shape: [32, 10], dataType: .float32)
    
    let weights = graph.variable(
        with: MPSDataType.float32, shape: [784, 10])
    let bias = graph.variable(
        with: MPSDataType.float32, shape: [10])
    
    let logits = graph.matrixMultiplication(
        primary: input, secondary: weights, name: "logits")
    let output = graph.addition(logits, bias, name: "output")
    
    // 损失函数：交叉熵
    let loss = graph.softMaxCrossEntropy(
        source: output, labels: label, name: "loss")
    
    // 自动构建梯度图
    let gradients = graph.gradient(
        for: loss, with respectTo: [weights, bias])
    
    return (graph, loss, gradients)
}
// 执行计算图
func runGraph(
    graph: MPSGraph,
    device: MTLDevice,
    commandQueue: MTLCommandQueue,
    input: MPSGraphTensor,
    output: MPSGraphTensor,
    inputData: MPSVector
) {
    guard let commandBuffer = commandQueue.makeCommandBuffer() else { return }
    
    // 创建输入数据字典
    let feeds: [MPSGraphTensor: MPSGraphTensorData] = [
        input: MPSGraphTensorData(inputData)
    ]
    
    // 执行并获取结果
    let results = graph.run(
        feed: feeds,
        targetTensors: [output],
        targetOperations: nil,
        commandBuffer: commandBuffer
    )
    
    commandBuffer.commit()
    commandBuffer.waitUntilCompleted()
    
    // 读取输出
    if let outputData = results[output] {
        print("推理完成，输出形状: \(outputData.shape)")
    }
}

最佳实践

  • 模型推理优先用 Core ML——只有当 Core ML 不支持你的自定义算子时才用 MPSGraph。Core ML 的图优化和内存管理比手动使用 MPSGraph 更成熟。
  • 训练场景下用 MPSGraph 的自动梯度功能，不要手写反向传播。手动写梯度容易出错，且每次前向图修改都需要同步更新反向图。
  • 用 MPSGraph 的 profiling 工具检查每个算子的执行时间，瓶颈通常在数据搬运（CPU-GPU 传输）而不是计算本身——减少中间结果的回读。

还有什么值得关注

  • MPSGraph 现在支持从 ONNX 模型直接导入计算图，省去了手动逐层构建的麻烦。
  • 新增的 MPSGraphConvolution2DOpDescriptor 支持分组卷积（grouped convolution）和深度可分离卷积（depthwise convolution），适合 MobileNet 类架构。
  • iOS 15 的 MPSGraph 支持了 float16 混合精度训练，在 Apple GPU 上比 float32 快约 1.5-2 倍，精度损失可忽略。
WWDC 2021