[TensorFlow] Encrypting/decrypting a pre-trained model in mobile applications

Every mobile application with Machine Learning functions should contain encrypted model weights and layers data.

Let’s explore the most popular ML frameworks for mobile devices and encryption implementation for each.

TensorFlow

Model encryption for TensorFlow is quite simple.

Every TensorFlow model is serialized like a protobuff object into a file via Google Protocol Buffers Library.

Let’s see the content of TensorFlow proto (*.pb) file in non-binary representation:

tensorflow non-binary proto (*.pb) file

It demonstrates the structure of saved pb-model.
Every “Node” contains some params, such as “name”, “op”, “attr”, etc.

All weights of the neural net are stored in each “tensor_content” property of each node-variable from the graph.

Model Graph is represented by GraphDef class in TensorFlow, which inherits “google::protobuf::Message” class from the protobuf, so we can restore this object from the byte array or some string.

How to do that using Python API:

1. Load graph from the file path:

2. Get the content of it and serialize to string.

So we can encrypt model data (for example via AES), save that into a file, load this file from the app, decrypt and restore GraphDef object from file bytes.

It’s very easy!

I’ve developed a small library for that — TFSecured (with iOS demo), https://github.com/dneprDroid/TFSecured

Encrypt pb-model via python script:

$ python encrypt_model.py <INPUT_PB_MODEL>  \   # input path
                          <OUTPUT_PB_MODEL> \   # output path (for saving)
                          <KEY>

Sample usage — decryption (C++):

You may use a random string with random length as a key, then library calculates sha256 hash of it and uses as internal AES key with size 256 bits.

So you may send to your backend some internal data of the device and/or user (device id, etc.), then the backend calculates from that the key and sends the model encrypted by this key to the mobile device, next — you can generate in the same way the key in the mobile app and decrypt the model.

Important: Don’t send the key through the network, because it’s very easy to intercept all transferred data.