Training LSTM Model With MLCompute on iOS or macOS

Build a prediction model for your iOS applications using the MLCompute framework

Source: Undraw

During WWDC 2020 Apple presented a new MLCompute framework. It provides a flexible API for training and inferencing neural networks on different computing units.

I recently had a task to train an LSTM model.

LSTM is an advanced RNN(Recurrent Neural Network) that has 2 states between predictions, unlike vanilla RNN.

This improves predictions results and protects us from vanishing gradients during training.

Time Series Forecasting

Our task was to predict future tendentious using the current state. We can train LSTM for solve this task:

Tensors

MLCompute used MLCTensor for calculations. We need to setup an array of MLCTensor objects for setup LSTM:

For working with LSTM, MLCompute provides 2 classes:

Here’s the code for the MLCLSTMDescriptor class:

And the code for MLCLSTMLayer is:

Build Pipeline

• For inferences, it uses: MLCInferenceGraph
• For train, it uses: MLCTrainingGraph
• To initialise the above two objects and setup our LSTM layer, we need to use MLCGraph.

Training

Before training, we need wrap our data using MLCTensorData object and then execute the train loop as shown below:

Results

You can download a sample project from my GitHub Repository:

GitHub - dhrebeniuk/MLCLSTMExample: Example how we can use LSTM with MLCompute

That’s all. Thanks for reading.