Video • Slides • Code below!

Did a recap of the machine learning and mobile talk (transcript there!) for the swift-sig group.

Errata: For the original presentation, I used ‘identity’ as the activation function for the output layer in the swift code, rather than the ‘softmax’ I used on this one (see note below). I thought it would make things more consistent, but instead it seemed to make the model more brittle (eg. why it acted up during this demo). If you’re trying this, I would suggest you use identity instead, which produces nicer results (but is probably not technically correct). Having said all that, keeping a 1:1 correlation between this code and the swift-sig presentation to avoid confusion.

Code:

Checkout swift-models repository.

Open Custom-CIFAR10 demo.

Remove data normalization step from the cifar data, eg change this line in CIFAR10.swift:

let imagesNormalized = ((imageTensor / 255.0) - mean) / std

to be:

let imagesNormalized = imageTensor / 255.0

This is just to match our eventual input from the default ios device.

Replace model with vgg-esque Swift model

struct CIFARModel: Layer {
    typealias Input = Tensor<Float>
    typealias Output = Tensor<Float>

    var conv1a = Conv2D<Float>(filterShape: (3, 3, 3, 32), padding: .same, activation: relu)
    var conv1b = Conv2D<Float>(filterShape: (3, 3, 32, 32), activation: relu)
    var pool1 = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))

    var conv2a = Conv2D<Float>(filterShape: (3, 3, 32, 64), padding: .same, activation: relu)
    var conv2b = Conv2D<Float>(filterShape: (3, 3, 64, 64), activation: relu)
    var pool2 = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))

    var flatten = Flatten<Float>()
    var dense1 = Dense<Float>(inputSize: 6 * 6 * 64, outputSize: 512, activation: relu)
    var dense2 = Dense<Float>(inputSize: 512, outputSize: 512, activation: relu)
    var output = Dense<Float>(inputSize: 512, outputSize: 10, activation: softmax) // use identity here!

    @differentiable
    func callAsFunction(_ input: Input) -> Output {
        let conv1 = input.sequenced(through: conv1a, conv1b, pool1)
        let conv2 = conv1.sequenced(through: conv2a, conv2b, pool2)
        return conv2.sequenced(through: flatten, dense1, dense2, output)
    }
}

Add s4tf –> numpy –> keras.save(“cifar.h5”)

    func save() {
      let keras = Python.import("keras")
      let model = keras.Sequential()

      model.add(keras.layers.Conv2D(filters: 32, kernel_size: [3, 3],
	            padding: "same", activation: "relu", input_shape: [32, 32, 3]))
      model.add(keras.layers.Conv2D(filters: 32, kernel_size: [3, 3], activation: "relu"))
      model.add(keras.layers.MaxPooling2D(pool_size: [2, 2], strides: [2, 2]))

      model.add(keras.layers.Conv2D(filters: 64, kernel_size: [3, 3],
	            padding: "same", activation: "relu"))
      model.add(keras.layers.Conv2D(filters: 64, kernel_size: [3, 3], activation: "relu"))
      model.add(keras.layers.MaxPooling2D(pool_size: [2, 2], strides: [2, 2]))

      model.add(keras.layers.Flatten())
      model.add(keras.layers.Dense(512, activation: "relu", input_shape: [2304]))
      model.add(keras.layers.Dense(512, activation: "relu"))
      model.add(keras.layers.Dense(10, activation: "softmax"))
	 
	  print (model.summary())

      model.layers[0].set_weights([conv1a.filter.makeNumpyArray(), conv1a.bias.makeNumpyArray()])
      model.layers[1].set_weights([conv1b.filter.makeNumpyArray(), conv1b.bias.makeNumpyArray()])

      model.layers[3].set_weights([conv2a.filter.makeNumpyArray(), conv2a.bias.makeNumpyArray()])
      model.layers[4].set_weights([conv2b.filter.makeNumpyArray(), conv2b.bias.makeNumpyArray()])

      model.layers[7].set_weights([dense1.weight.makeNumpyArray(), dense1.bias.makeNumpyArray()])
      model.layers[8].set_weights([dense2.weight.makeNumpyArray(), dense2.bias.makeNumpyArray()])
      model.layers[9].set_weights([output.weight.makeNumpyArray(), output.bias.makeNumpyArray()])
	  
      model.save("cifar.h5")
    }

will need to add a model.save() step after the main.swift loop has run to call this!

Run your model to produce a .h5 file.

python .h5 –> freeze graph –> protobuf

Copy/paste this to a .py file and run it.

import numpy as np
import tensorflow as tf

def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    """
    Freezes the state of a session into a pruned computation graph.

    Creates a new computation graph where variable nodes are replaced by
    constants taking their current value in the session. The new graph will be
    pruned so subgraphs that are not necessary to compute the requested
    outputs are removed.
    @param session The TensorFlow session to be frozen.
    @param keep_var_names A list of variable names that should not be frozen,
                          or None to freeze all the variables in the graph.
    @param output_names Names of the relevant graph outputs.
    @param clear_devices Remove the device directives from the graph for better portability.
    @return The frozen graph definition.
    """
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ''
        frozen_graph = tf.graph_util.convert_variables_to_constants(
            session, input_graph_def, output_names, freeze_var_names)
        return frozen_graph


from tensorflow import keras
from tensorflow.keras import layers

model = keras.models.load_model('/home/skoonce/swift/cifar-mlir-demo/swift-models/cifar.h5')

frozen_graph = freeze_session(tf.keras.backend.get_session(), output_names=[out.op.name for out in model.outputs])

tf.io.write_graph(frozen_graph, './', 'cifar.pbtxt', as_text=True)
tf.io.write_graph(frozen_graph, './', 'cifar.pb', as_text=False)

Bazel/MLIR compile/build/run

Warning: this might take a while 😉

git clone https://github.com/tensorflow/tensorflow
git reset --hard a2520cbd578590b142342910520b551541a82471
bazel build -c opt tensorflow/compiler/mlir/...

Protobuf –> mlir –> tf-lite file

./tensorflow/bazel-bin/tensorflow/compiler/mlir/lite/tf_tfl_translate \
-tf-input-shapes=1,32,32,3 \
-tf-input-data-types=DT_FLOAT \
-tf-output-arrays=dense_3/Softmax \
/home/skoonce/swift/cifar-mlir-demo/cifar.pb \
--tf-input-arrays=conv2d_1_input \
-o cifar-mlir-demo/cifar-mlir.tflite

Run on device:

Checkout tensorflow-lite ios image classification example.

Run pod install step.

Add .tflite model to demo, add cifar data to labels. Override the model in “ModelDataHandler.swift”:

    static let modelInfo: FileInfo = (name: "cifar-mlir", extension: "tflite")
    static let labelsInfo: FileInfo = (name: "labels-cifar", extension: "txt")

change input sizes to match new model:

    let inputWidth = 32
    let inputHeight = 32

and run! 🥳