如何将TensorFlow模型保存到pickle文件

Question

如何将TensorFlow模型保存到pickle文件

6

我想保存Tensorflow模型以备后续部署使用。我不想使用model.save()进行保存，因为我的目的是以某种方式将其“序列化”，并在未安装Tensorflow的不同系统中使用，例如：

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

之前我使用sklearn进行KNN模型的序列化时，成功了，并且无需安装sklearn即可运行推理。

但是当我尝试序列化Tensorflow模型时，出现了以下错误：

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

我的模型长这样：

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: "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

这里有一个涉及此问题的StackOverflow问题，但是答案中的链接已经失效了。

还有另一个类似的问题在这里，但我并没有完全理解。

我有一个非常简单的模型，没有检查点，也没有太复杂的东西，所以是否有一些方法将TensorFlow模型对象保存到二进制文件中？即使是多个二进制文件，我也不介意，但它只是不需要使用tensoflow。如果numpy解决方案会有帮助的话，我会使用它，但我不知道如何在这里实现它。任何帮助都将不胜感激，谢谢！

- shakhyar.codes

你需要从模型中提取权重，将它们创建成一个数组并对该数组进行pickle。否则我不确定还有其他的方法可以完成这个任务。 - NotAName

2个回答

1

对于使用keras创建的TensorFlow模型，请按照此处类似答案中所述，将模型保存为.h5扩展名：

错误，未成功的TensorSliceReader构造函数：无法找到任何匹配的文件以取消pickle文件ram：

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

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- AloneTogether · Accepted Answer

使用joblib似乎可以在TF 2.8上工作，由于您有一个非常简单的模型，因此可以在Google Colab上对其进行训练，然后只需在其他系统上使用pickle文件：

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')

无需 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)

不知道您的系统规格，很难判断它是否真正“简单直接”。