how to save tensorflow model to pickle file

Question

I want to save a Tensorflow model and then later use it for deployment purposes. I dont want to use model.save() to save it because my purpose is to somehow 'pickle' it and use it in a different system where tensorflow is not installed, like:

model = pickle.load(open(path, 'rb'))
model.predict(prediction_array)

Earlier with sklearn, when i was pickling a KNN model, it was successful and i was able to run inference without installing sklearn.

But when I tried to pickle my Tensorflow model, I got this error:

Traceback (most recent call last):
  File "e:/VA_nlu_addition_branch_lite/nlu_stable2/train.py", line 21, in <module>
pickle.dump(model, open('saved/model.p', 'wb'))
TypeError: can't pickle _thread.RLock objects

My model looks like this:

model = keras.Sequential([
            keras.Input(shape=(len(x[0]))),
            keras.layers.Dense(units=16, activation='elu'),
            keras.layers.Dense(units=8, activation='elu'),
            keras.layers.Dense(units=len(y[0]), activation='softmax'),
        ])

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(x, y, epochs=200, batch_size=8)
pickle.dump(model, open('saved/model.p', 'wb'))

Model summary

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
dense (Dense)                (None, 16)                1680
_________________________________________________________________
dense_1 (Dense)              (None, 8)                 136
_________________________________________________________________
dense_2 (Dense)              (None, 20)                180
=================================================================
Total params: 1,996
Trainable params: 1,996
Non-trainable params: 0

Here is a StackOverflow question regarding this problem, but the link in the answer was expired.

Also here is another similar question, but i didn't quite get it.

I have a very simple model, no checkpoints, nothing much complicated, so is there some way to save the Tensorflow model object to a binary file? Or even if its multiple binary files, i dont mind, but it just doesn't need to use tensoflow, if the numpy solution would help, i would use that, but i dont know how to implement it here. Any help would be appreciated, Thanks!

Tonry

You need to extract weights from the model, create an array out of them and pickle that array. I'm not sure how else it can be done otherwise. – Haynes

Answer 1

Using joblib seems to work on TF 2.8 and since you have a very simple model, you can train it on Google Colab and then just use the pickled file on your other system:

import joblib
import tensorflow as tf

model = tf.keras.Sequential([
            tf.keras.layers.Input(shape=(5,)),
            tf.keras.layers.Dense(units=16, activation='elu'),
            tf.keras.layers.Dense(units=8, activation='elu'),
            tf.keras.layers.Dense(units=5, activation='softmax'),
        ])

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
x = tf.random.normal((20, 5))
y = tf.keras.utils.to_categorical(tf.random.uniform((20, 1), maxval=5, dtype=tf.int32))
model.fit(x, y, epochs=200, batch_size=8)
joblib.dump(model, 'model.pkl')

Load model without tf:

import joblib
import numpy as np
print(joblib.__version__)

model = joblib.load("/content/model.pkl")
print(model(np.random.random((1,5))))

1.1.0
tf.Tensor([[0.38729233 0.04049021 0.06067584 0.07901421 0.43252742]], shape=(1, 5), dtype=float32)

But it is hard to tell if it is really that "straight-forward" without knowing your system specs.

Repairer

hey, its not working, i was using tf 2.0, now it shows TypeError: can't pickle _thread._local objects. So does that means its wrong with tensorflow versions? For my specs, i have an intel i5 x64 bit processor, and using tensorflow cpu version 2.0.0 – Tonry

Yeah, I think, starting from version 2.6.0, keras models can be pickled..A tf version upgrade should solve the problem..but maybe try what I suggested: train your model on google colab and just load the binary on your system – Repairer

so when i loaded it locally, it says Using TensorFlow backend. (although i haven't imported tensorflow in my script... and it throws an exception: ModuleNotFoundError: No module named 'keras.saving'. The exception is coming from: File "C:\Users\user\AppData\Local\Programs\Python\Python36\lib\pickle.py", line 1388, in find_class __import__(module, level=0), so this error still belongs to pickle right? or should i have to upgrade my python as well to that of Colab's version? – Tonry

Which method are you experimenting with right now? – Repairer

training on colab and then loading the binary locally for inference – Tonry

Are you using joblib or pickle? – Repairer

i am using joblib – Tonry

Yes so the difference in my joblib and python versions to that of Colab was causing the mess. Both the system's joblib and python versions should be same. Thanks! – Tonry

Answer 2

For TensorFlow models created using keras, please save models in .h5 extension as explained in a similar answer here:

Error , Unsuccessful TensorSliceReader constructor: Failed to find any matching files for ram to unpickle a file:

from tensorflow.keras.models import load_model as tfk__load_model
# ...
# fit model:
# model.fit(x, y, ...)

# save fitted model to file
model.save('model.h5')

# load fitted model
fitted_model = tfk__load_model('model.h5')

Touchy