Home > front end >  how to save tensorflow model to pickle file
how to save tensorflow model to pickle file

Time:03-16

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!

CodePudding user response:

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.

  • Related