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
Watch original video

Highlight

MPSGraph extends Metal Performance Shaders to multi-dimensional tensor calculations, allowing developers to use graphs, tensors, variables, and automatic differentiation to perform custom machine learning inference and training processes on GPUs.

Core Content

In the past, developers who wanted to implement a custom machine learning operator on the GPU needed to switch back and forth between data structures, memory flows, and Metal kernels. An activation function seems to be just a few multiplications, additions and error functions. When it falls on the GPU, each step may be written back to the global memory and then read out by the next step. The performance loss will be directly pressed into the inference loop.

The answer given by Apple in this session is MPSGraph. It writes computations as directed acyclic graphs: operations are nodes, tensors are edges, placeholders receive input at runtime, and target tensors specify the results to return. Developers describe mathematical relationships, and MPSGraph is responsible for handing the graph to the compiler for optimization.

The value of this model is clear in the GeLU activation function example. Developers use MPSGraph to combine constants, square roots, multiplications, error functions, additions, and multiplications to obtain a custom GPU calculation function; then the compiler stitches adjacent operations into a more compact Metal shader area, and the GeLU acceleration given by the session is close to 10 to 50 times.

The scope of MPSGraph also covers training. The session demonstrates variables, automatic differentiation, optimizer updates, asynchronous execution, and single run calls using numeric classifiers and generative adversarial networks. The point is not to hide the model in a black box, but to let developers express inference, loss functions, gradients and parameter updates directly.

Detailed Content

01:30 MPS first completes the basic capabilities of machine learning

At the beginning of the session, the base changes of MPS will be explained. MPS has added the EDLines edge detection primitive, which has resulted in a 10x speedup for the Measure app; the distance transformation filter has also been optimized, achieving a 15x speedup among the features of Final Cut Pro. In the machine learning part, convolution and fully connected primitives add complete FP32 support, and developers can implement it on the data source provider.kernelWeightsDataTypeand set tofloat32, which preserves accuracy and reduces the effort of recalculating hyper-parameters or refactoring code.

Key points:

  • FP32 support falls on convolution and fully connected primitives.
  • The access point explicitly given by the session is the data source provider.kernelWeightsDataType.
  • This part solves the problem of accuracy migration, not the problem of model structure expression.

02:05 MPSNDArray brings multi-dimensional tensors to MPS

MPSNDArray is a new MPS primitive that supports up to 16 dimensions. There are no additional restrictions on dimension size, as long as the data can fit into memory. The creation process is divided into two steps: first useMPSNDArrayDescriptorSpecify the shape and element data type, and then pass the descriptor and Metal device to the initializer.

Key points:

  • MPSNDArray is oriented to multi-dimensional data in machine learning and scientific computing.
  • It can be added to an existingMTLTexture based MPSImageBatchConvert between types.
  • It can also import and export data through Metal buffers to facilitate access to existing Metal data channels.

03:00 MPSGraph uses graphs to describe tensor data flow

MPSGraph is located on top of the MPS framework and is used to extend Metal’s compute capabilities to multidimensional tensors. It provides support on macOS, iOS, iPadOS and tvOS. The Metal Performance Shaders Graph framework needs to be added from Xcode’s Build Phases to the project.

Key points:

  • MPSGraphRepresents a DAG composed of operations and tensors.
  • operations are calculation nodes, such as multiply, add, subtract.
  • Tensors are data flow edges, including shape, data type, and operation references that create the tensor.
  • graph takes ownership of tensors and operations created through its own methods.

07:22 Execute GeLU with placeholder, operation and target

The GeLU example breaks the custom calculation into three steps. The first step is to use placeholder to create the input tensor; if a certain dimension is only determined at runtime, write the dimension as-1;If even the number of dimensions is unknown, placeholder shape can be passednil, MPSGraph propagates shape at runtime. The second step is to write GeLU as a set of MPSGraph operations, including constants,squareRoot, multiply, error function, add. The third step is to pass in the feeds dictionary when executing graph, specify target tensors, and optionally specify target operations.

Key points:

  • placeholder is the runtime data entry, provided by callerMPSGraphTensorDatareplace.
  • feeds dictionary maps placeholders to input data.
  • target tensors determine which results the graph will compute and return. -MPSGraphTensorDataCan be composed of Metal buffer,MPSImageBatch、MPSVector、MPSMatrixorMPSNDArrayinitialization.

13:00 stitching automatically fuses adjacent operations

In the GeLU example, a conventional implementation might have each operation write its output back to memory, where it can be read by the next operation. The MPSGraph compiler identifies adjacent operations and passes them to the Metal compiler to fuse them into a single optimized shader. For hand-tuned MPS kernels such as convolution, matrix multiplication, and reduction, MPSGraph can also identify surrounding stitchable math operations as regions and integrate them into these hand-tuned kernels.

Key points:

  • Stitching is automatically done by MPSGraph compiler and Metal compiler.
  • The GeLU results given by -session are close to 10 to 50 times faster and bring significant memory savings.
  • Developers do not need to write additional code for this optimization.

15:17 Inference graph can express circular text generation

The text generation demo uses the Long Short Term Memory based Text Generator model. The generation loop starts from a character index, obtains the rank-one tensor through the gather operation of the embedding layer, and then sends it to the LSTM layer. The LSTM results enter the un-embedding layer, which is implemented by matrix multiplication operation. Then the result is multiplied by the inverse temperature, MPSGraph automatically broadcasts the scalar to the vector length, and then enters the softmax layer to form the probability distribution of the next character.

Key points:

  • The trained parameters are fixed during the inference phase and can be represented by constants.
  • gather, matrix multiplication, broadcast, and softmax are all expressed by MPSGraph operation.
  • Temperature affects the distribution shape before softmax. High temperature is more random, and low temperature is closer to the training data.

20:21 Training graph uses variables and automatic differentiation to update parameters

The digital classifier training process first defines input and target label placeholders, and uses-1Represents variable length batch size. Trainable parameters are part of the graph and are saved through variables; variables retain their values ​​across graph runs, so they do not need to be passed in as feeds in each round. The convolution layer first creates the convolution descriptor and then calls the convolution method on the graph. layout passedMPSGraphTensorNamedDataLayoutExpression, supporting common formats such as NCHW, NHWC, OIHW, HWIO, etc.

Key points:

  • variables are created using graph’s variable method, including initial values, shape and element data type. -readVariableReturns the tensor value of the variable at the current execution point. When passing the variable tensor directly to the operation, the graph will implicitly add readVariable.
  • softmax axis is set to-1, pointing to the fastest changing axis.
  • loss Use Softmax Cross Entropy, thenreductionSumAggregate into a single scalar loss.

25:47 gradient method automatically writes backward pass

Training requires calculating the gradients of loss on weights and biases. MPSGraph provides the gradient method, which takes a tensor and a set of target tensors and returns the tensor to a dictionary of gradients. Session uses logarithm as an example to illustrate that graph will generate reciprocal operations as partial derivatives, pruning useless operations and folding constants.

Key points:

  • automatic differentiation automatically writes backward pass.
  • For numeric classifiers, the gradient method is used to calculate the gradients of the loss tensor with respect to weights and biases.
  • The stochastic gradient descent optimizer uses the learning rate, original weight value and gradient to obtain the update tensor. -assign variableoperation writes the update tensor back to variable.

27:28 Can be cached, asynchronously, and plugged into existing Metal workflows when performing training

The training loop takes a batch of input images and labels in each round, passes it into the graph through feeds, requests the loss tensor as the target, and puts the assignment operations into the target operations to ensure that weights and biases are updated. The graph is compiled once when it first encounters a unique set of input and output types, and subsequent iterations automatically reuse cached executables. The asynchronous run method will return immediately,executionDescriptorCompletion handlers can be registered; the session also demonstrates the use of semaphore to control CPU encoding and GPU execution overlap.

Key points:

  • Parameter updates for each round of training are triggered by target operations.
  • cached executables reduce repeated compilation costs.
  • Asynchronous execution is suitable for overlapping CPU encoding and GPU execution.
  • Existing Metal workflows can pass in their ownMTLCommandQueue, you can also encode graph toMPSCommandBuffer。

31:21 A single MPSGraph can train both the generator and the discriminator

The final GAN ​​demo trains two networks. The generator generates handwritten digits from random seeds, uses convolution transpose to expand the output size, and uses mean, variance, gamma, and beta to form batch normalization. The discriminator receives real or generated images, outputs realness, and uses dropout to randomly zero out some of the input values, while increasing the remaining values ​​according to the reciprocal of rate.

Key points:

  • convolution transpose is described by corresponding convolution descriptor and expected output shape.
  • Batch normalization determines the normalized dimensions such as batch, height, and width by selecting axes.
  • discriminator pass only requests gradients of discriminator variables.
  • generator pass will cause MPSGraph to automatically differentiate across the two networks, but only return the gradients of the generator variables.
  • The losses and update operations of the two networks can be executed in one run call, and the graph will be compiled and optimized across the two networks.

Core Takeaways

  • Make a local text generator with adjustable temperature: Use MPSGraph to express embedding, LSTM, matrix multiplication, broadcast and softmax, and let users use sliders to adjust the temperature. This feature comes directly from the session’s Shakespeare demo, and the entry point is placeholder feeds, constant trained parameters and target tensor execution.

  • Make a device-side handwritten digit training sandbox: use canvas input as placeholder, convolution, bias add, ReLU, reshape, softmax, Softmax Cross Entropy andreductionSumCompose a training graph. Use automatic differentiation to obtain the gradients of weights and biases, and then update the parameters through SGD update and assign variable.

  • Make a custom activation function experiment panel: Starting from GeLU, combine constants, unary math, binary math and error function into MPSGraph helper. The panel shows the output and time consumption under different input shapes, focusing on verifying the optimization effect of stitching on adjacent math operations.

  • Make a small GAN ​​visual trainer: use generator to generate digital images, use discriminator to determine realness, and put the two networks into the same MPSGraph. The interface displays generator loss, discriminator loss and generated results. During training, target operations are used to trigger updates of both networks at the same time.

  • Make a Metal data channel bridging tool: Combine the existing Metal buffer,MPSImageBatchorMPSNDArraypackaged intoMPSGraphTensorData, allowing existing image processing or scientific computing pipelines to be gradually migrated to MPSGraph without having to change all data sources at once.

Comments

GitHub Issues · utterances