6

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!

shakhyar.codes
  • 218
  • 4
  • 14
  • 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. – NotAName Mar 16 '22 at 06:59

2 Answers2

6

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.

AloneTogether
  • 25,814
  • 5
  • 20
  • 39
  • 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 – shakhyar.codes Mar 16 '22 at 07:29
  • 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 – AloneTogether Mar 16 '22 at 07:31
  • 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? – shakhyar.codes Mar 16 '22 at 07:52
  • Which method are you experimenting with right now? – AloneTogether Mar 16 '22 at 07:54
  • training on colab and then loading the binary locally for inference – shakhyar.codes Mar 16 '22 at 07:56
  • Are you using joblib or pickle? – AloneTogether Mar 16 '22 at 07:57
  • i am using joblib – shakhyar.codes Mar 16 '22 at 07:58
  • 1
    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! – shakhyar.codes Mar 16 '22 at 08:05
1

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')
Curious Watcher
  • 580
  • 6
  • 12