Firebase ML Kit Custom Models for iOS developers — Part 2: Implementing Tic Tac Toe

In Part 1, we build a model to add 3 floating point numbers to understand how TensorFlow Lite model works with the Firebase ML Kit.

That model was never going to be used in a real app, so let’s consider a more realistic example. Let’s say you are building an app for playing Tic Tac Toe (a.k.a. Noughts and Crosses or OXO) and your app needs to suggest hints to the player or implement an AI.

Let’s create a TensorFlow Lite model that finds the best moves for a particular board state and use it inside our iOS app to score all the possible next moves that the player can make.

Step 1: Inspect the model

The Tic Tac Toe model takes a board state (3x3 grid of cells) and outputs the next move for every cell (so the output is also 3x3). It is a simple network with 2 hidden layers, trained on a million games played by 2 random players. Note that the model is not a perfect player! (I leave it to ML readers to retrain the model to be a better player — send me your results please)

You can recreate the dataset and retrain the model yourself in his Tic Tac Toe Colab. Or you can do like I did and jump straight to download the tictactoe.tflite file.

Open the model in Netron again. You will see the layers and the inputs and outputs that we are interested in:

• Input type: float32[3,3]
• Output type: float32[3,3]

Inspecting the inputs/outputs of the tic tac toe model in Netron

How do we input X and O as floats? The way the model was built is that if the current player is O then the cells with O will be encoded as 1.0, and the opponent’s Xs will be encoded as -1.0. If the current player is X then this is reversed (Xs are 1.0 and Os are -1.0). The empty cells are always encoded as 0.0.​

For example, the board state:

X|O|
-+-+-
X|O|O
-+-+-
 | |X

If the current player is O so input to the model would be:

[[-1., 1., 0.],
 [-1., 1., 1.],
 [0., 0., -1.]]

The output of the model is the same shape, where the value of each cell is how good the move is for the current player. Clearly the best move for O is the bottom middle cell, so you would hope that the model outputs similar to:

[[0., 0., 0.],
 [0., 0., 0.],
 [0., 1., 0.]]

By inspecting the model and understanding its inputs and outputs, we now know things that we are required to do in the Swift code:

1. Convert the board state into a 2D array (3x3) where current player cells are 1.0, opponent cells are -1.0 and empty cells are 0.0 as input to the model.
2. Interpret the model output as 0.0 for a bad move and 1.0 for a best move.

Let’s get started!

Step 2: Integrate the model in your iOS project

The UI of the app is a Tic Tac Toe board that enables the user to setup any board state and a button to switch the current player (“Next move”). The background of the empty cells in the board are shaded according the output of the model — dark grey is a higher value and a better move than a light grey.

1.Add the model locally and on cloud. We already did this for a local only model in the previous example. In this example, let’s also add the model to the server — ready so that we can update our model in the future (without our users having to update the app). We can do that by placing the TensorFlow Lite model file in the Xcode project for local as well as upload your TensorFlow Lite model file into Firebase Console under the ML Kit Custom Models section to make it available on cloud.

2. Prepare an interpreter. As before, load the model by configure your model sources and create a ModelInterpreter object from one of them.

As you can see, the steps to configure a local model are the same. However there are some extra configuration steps for remote (on cloud) models. Firstly, create a CustomRemoteModel object by specifying the name of the model that you assigned when you published it via Firebase Console. Another part that we need before create a ModelInterpreter is starting a model download task by passing a CustomRemoteModel with the download conditions.

Then, we can create a ModelInterpreter from one of the models, local or remote by checking whether the remote model has been downloaded already before using it.

3.Specify the model’s input and output. The Tic Tac Toe model has [3,3] as input and output shapes. Input is a current state of the board which is a 2D arrays that include -1 (opponent cell), 0 (empty cell), 1 (current player cell). The output is also 2D arrays which is the scores of all the moves as described in the inspect step above.

4. Perform inference on input data. Here come the fun part! Every time the board state or the player changes in the user interface, we call the run function of the interpreter to get the best moves for the current player.

In the run function, there are quite a lot of checking to ensure that the input and output dimension are correct. If it is all good, then the output will be send back to show the scores in the view controller.

I hope these Part 1 and Part 2 together would help you to get started with custom models on Firebase ML Kit and will be able to add more complex models to your iOS apps in the future.

Firebase ML Kit is still in beta and therefore might change — keep an eye on the official documentation for Custom Models on iOS for the latest examples.

Get the full project code for these examples at github.com/naluinui.