Training LSTM Model With MLCompute on iOS or macOS
Build a prediction model for your iOS applications using the MLCompute framework

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:
That’s all. Thanks for reading.