WWDC Quick Look 💓 By SwiftGGTeam
中文 EN 日本語
Build customized ML models with the Metal Performance Shaders Graph

Build customized ML models with the Metal Performance Shaders Graph

WWDC 2020Graphics & GamesAI & Machine LearningMetalMPSGraph
观看原视频

Highlight

MPSGraph 把 Metal Performance Shaders 扩展到多维张量计算，开发者可以用图、张量、变量和自动微分在 GPU 上执行自定义机器学习推理与训练流程。

核心内容

以前想在 GPU 上实现一个定制机器学习算子，开发者需要在数据结构、内存流动和 Metal kernel 之间来回切换。一个激活函数看似只是几次乘法、加法和误差函数，落到 GPU 上时，每一步都可能写回全局内存，再被下一步读走，性能损耗会直接压到推理循环里。

Apple 在这场 session 里给出的答案是 MPSGraph。它把计算写成有向无环图：操作是节点，张量是边，placeholder 在运行时接收输入，target tensor 指定要返回的结果。开发者描述数学关系，MPSGraph 负责把图交给编译器做优化。

这套模型的价值在 GeLU 激活函数示例里很清楚。开发者用 MPSGraph 组合常量、平方根、乘法、误差函数、加法和乘法，得到一个自定义 GPU 计算函数；随后编译器把相邻操作 stitching 到一个更紧凑的 Metal shader 区域里，session 给出的 GeLU 加速幅度达到接近 10 到 50 倍。

MPSGraph 的范围还覆盖训练。session 用数字分类器和生成对抗网络展示了变量、自动微分、优化器更新、异步执行和单次 run 调用。重点不在于把模型藏进一个黑盒，而是让开发者直接表达推理、损失函数、梯度和参数更新。

详细内容

01:30 MPS 先补齐机器学习基础能力

session 开头先交代 MPS 的底座变化。MPS 新增了 EDLines 边缘检测 primitive，Measure app 因此获得 10 倍提速；距离变换 filter 也被调优，在 Final Cut Pro 的特性里达到 15 倍提速。机器学习部分，卷积和全连接 primitive 增加完整 FP32 支持，开发者通过在 data source provider 上实现 kernelWeightsDataType 并设为 float32，可以保留准确度，减少重新计算 hyper-parameters 或重构代码的工作。

关键点：

  • FP32 支持落在 convolution 和 fully connected primitives 上。
  • session 明确给出的接入点是 data source provider 的 kernelWeightsDataType。
  • 这部分解决的是精度迁移问题，不是模型结构表达问题。

02:05 MPSNDArray 给 MPS 带来多维张量

MPSNDArray 是新的 MPS primitive，支持最多 16 个维度。维度大小没有额外限制，只要数据能放进内存即可。创建流程分两步：先用 MPSNDArrayDescriptor 指定 shape 和 element data type，再把 descriptor 与 Metal device 传给 initializer。

关键点：

  • MPSNDArray 面向机器学习和科学计算里的多维数据。
  • 它可以在已有的 MTLTexture based MPSImageBatch 类型之间转换。
  • 它也可以通过 Metal buffers 导入和导出数据，方便接入现有 Metal 数据通道。

03:00 MPSGraph 用图描述张量数据流

MPSGraph 位于 MPS framework 之上，用来把 Metal 的 compute 能力扩展到 multidimensional tensors。它在 macOS、iOS、iPadOS 和 tvOS 上提供支持，项目里需要从 Xcode 的 Build Phases 添加 Metal Performance Shaders Graph framework。

关键点：

  • MPSGraph 表示 operations 与 tensors 组成的 DAG。
  • operations 是计算节点，例如 multiply、add、subtract。
  • tensors 是数据流边，包含 shape、data type，以及创建该 tensor 的 operation 引用。
  • graph 对通过自身方法创建的 tensors 与 operations 持有所有权。

07:22 用 placeholder、operation 和 target 执行 GeLU

GeLU 示例把自定义计算拆成三步。第一步，用 placeholder 创建输入 tensor；如果某个维度只在运行时确定，把该维度写成 -1；如果连维度数量也未知，placeholder shape 可以传 nil，MPSGraph 会在运行时传播 shape。第二步，把 GeLU 写成一组 MPSGraph operations，包括 constants、squareRoot、multiply、error function、add。第三步，执行 graph 时传入 feeds dictionary，指定 target tensors，并可选地指定 target operations。

关键点：

  • placeholder 是运行时数据入口，由 caller 提供的 MPSGraphTensorData 替换。
  • feeds dictionary 把 placeholder 映射到输入数据。
  • target tensors 决定 graph 要计算并返回哪些结果。
  • MPSGraphTensorData 可以由 Metal buffer、MPSImageBatch、MPSVector、MPSMatrix 或 MPSNDArray 初始化。

13:00 stitching 自动融合相邻操作

GeLU 示例里，常规实现可能让每个 operation 都把输出写回内存，再由下一个 operation 读取。MPSGraph compiler 会识别相邻 operation，把它们交给 Metal compiler 融合成单个优化 shader。对于 convolution、matrix multiplication、reduction 等手工调优的 MPS kernels，MPSGraph 还能把周边可 stitching 的 math operations 识别成区域，并融合进这些 hand-tuned kernels。

关键点：

  • stitching 由 MPSGraph compiler 与 Metal compiler 自动完成。
  • session 给出的 GeLU 结果是接近 10 到 50 倍提速，并带来显著内存节省。
  • 开发者不需要为这项优化额外写代码。

15:17 推理图可以表达循环式文本生成

文本生成 demo 使用 Long Short Term Memory based Text Generator model。生成循环从一个字符 index 开始，通过 embedding layer 的 gather operation 得到 rank-one tensor，再送入 LSTM layer。LSTM 结果进入 un-embedding layer，这里由 matrix multiplication operation 实现。随后结果乘以 inverse temperature，MPSGraph 自动把 scalar broadcast 到 vector 长度，再进入 softmax layer 形成下一个字符的概率分布。

关键点：

  • trained parameters 在推理阶段固定，可以用 constants 表示。
  • gather、matrix multiplication、broadcast、softmax 都由 MPSGraph operation 表达。
  • temperature 影响 softmax 前的分布形状，高温更随机，低温更贴近训练数据。

20:21 训练图用变量和自动微分更新参数

数字分类器训练流程先定义 input 与 target label placeholders，并用 -1 表示 variable length batch size。trainable parameters 是 graph 的一部分，通过 variables 保存；variables 会跨 graph runs 保留值，所以每一轮不用作为 feeds 传入。卷积层先创建 convolution descriptor，再调用 graph 上的 convolution method。layout 通过 MPSGraphTensorNamedDataLayout 表达，支持 NCHW、NHWC、OIHW、HWIO 等常见格式。

关键点：

  • variables 用 graph 的 variable method 创建，包含 initial values、shape 和 element data type。
  • readVariable 返回变量在当前执行点的 tensor 值，直接把 variable tensor 传给 operation 时，graph 会隐式添加 readVariable。
  • softmax axis 设为 -1，指向最快变化的 axis。
  • loss 使用 Softmax Cross Entropy，再用 reductionSum 聚合成单个 scalar loss。

25:47 gradient method 自动写出 backward pass

训练需要计算 loss 对 weights 和 biases 的 gradients。MPSGraph 提供 gradient method，输入一个 tensor 和一组目标 tensors，返回 tensor 到 gradients 的 dictionary。session 用 logarithm 示例说明，graph 会生成 reciprocal operation 作为 partial derivative，并把无用 operation pruning 掉，把常量折叠掉。

关键点：

  • automatic differentiation 自动写出 backward pass。
  • 对数字分类器，gradient method 被用于计算 loss tensor 对 weights 和 biases 的 gradients。
  • stochastic gradient descent optimizer 使用 learning rate、原 weight value 与 gradient 得到 update tensor。
  • assign variable operation 把 update tensor 写回 variable。

27:28 执行训练时可以缓存、异步和接入现有 Metal 工作流

训练 loop 每轮取一批 input images 和 labels，通过 feeds 传入 graph，请求 loss tensor 作为 target，并把 assignment operations 放进 target operations，保证 weights 和 biases 被更新。graph 在第一次遇到某组唯一输入输出类型时编译一次，后续 iteration 自动复用 cached executables。异步 run method 会立即返回，executionDescriptor 可以注册 completion handler；session 还展示了用 semaphore 控制 CPU encoding 与 GPU execution 重叠。

关键点：

  • 每轮训练的参数更新由 target operations 触发。
  • cached executables 减少重复编译成本。
  • 异步执行适合把 CPU 编码和 GPU 执行重叠。
  • 现有 Metal workflow 可以传入自己的 MTLCommandQueue，也可以把 graph encode 到 MPSCommandBuffer。

31:21 单个 MPSGraph 可以训练生成器和判别器

最后的 GAN demo 训练两个网络。generator 从随机 seed 生成手写数字，使用 convolution transpose 扩大输出尺寸，并用 mean、variance、gamma、beta 组成 batch normalization。discriminator 接收真实或生成的图像，输出 realness，并使用 dropout 随机置零部分输入值，同时按 rate 的倒数提升剩余值。

关键点：

  • convolution transpose 通过对应 convolution descriptor 和 expected output shape 描述。
  • batch normalization 通过选择 axes 决定 batch、height、width 等归一化维度。
  • discriminator pass 只请求 discriminator variables 的 gradients。
  • generator pass 会让 MPSGraph 穿过两个网络做自动微分，但只返回 generator variables 的 gradients。
  • 两个网络的 losses 与 update operations 可以在一次 run call 中执行，graph 会跨两个网络编译和优化。

核心启发

  • 做一个可调温度的本地文本生成器：用 MPSGraph 表达 embedding、LSTM、matrix multiplication、broadcast 和 softmax，让用户用 slider 调整 temperature。这个 feature 直接来自 session 的 Shakespeare demo，入口是 placeholder feeds、constant trained parameters 和 target tensor 执行。

  • 做一个设备端手写数字训练沙盒：把画布输入作为 placeholder，把 convolution、bias add、ReLU、reshape、softmax、Softmax Cross Entropy 和 reductionSum 组成训练图。用 automatic differentiation 获取 weights 和 biases 的 gradients，再通过 SGD update 与 assign variable 更新参数。

  • 做一个自定义激活函数实验面板：从 GeLU 开始，把 constants、unary math、binary math 和 error function 组合成 MPSGraph helper。面板显示不同输入 shape 下的输出和耗时，重点验证 stitching 对相邻 math operations 的优化效果。

  • 做一个小型 GAN 可视化训练器：用 generator 生成数字图像，用 discriminator 判断 realness，把两个网络放进同一个 MPSGraph。界面展示 generator loss、discriminator loss 和生成结果，训练时用 target operations 同时触发两个网络的 update。

  • 做一个 Metal 数据通道桥接工具：把现有 Metal buffer、MPSImageBatch 或 MPSNDArray 包装成 MPSGraphTensorData，让已有图像处理或科学计算管线能逐步迁移到 MPSGraph，而无需一次改掉所有数据来源。

关联 Session

评论

GitHub Issues · utterances