使用Pickle加载TensorFlow模型失败（Windows环境）

尝试用Pickle序列化TensorFlow模型，以下是保存模型的代码（dump.py）：

import tensorflow as tf
import pickle
import numpy as np

tf.random.set_seed(42)

input_x = np.random.randint(0, 50000, (10000,1))
input_y = np.random.randint(0, 50000, (10000,1))
output = input_x + input_y
input = np.concatenate((input_x, input_y), axis=1)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(2, activation = tf.keras.activations.relu, input_shape=[2]),   
    tf.keras.layers.Dense(2, activation = tf.keras.activations.relu),
    tf.keras.layers.Dense(1),
])

model.compile(loss = tf.keras.losses.mae,
              optimizer=tf.optimizers.Adam(learning_rate=0.00001),
              metrics = ['mse'])
          
model.fit(input, output, epochs = 1000)

fl = open('D:/tf/tf.pkl', 'wb')
pickle.dump(model, fl)
fl.close()

加载模型的代码（load.py）：

import pickle

fl = open('D:/tf/tf.pkl', 'rb')
model = pickle.load(fl)
print(model.predict([[2.2, 5.1]]))
fl.close()

代码在Linux下运行正常，但Windows环境中dump.py执行成功，load.py报错，错误信息如下：

2022-08-09 19:48:30.078245: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cudart64_110.dll'; dlerror: cudart64_110.dll not found
2022-08-09 19:48:30.078475: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
2022-08-09 19:48:32.847626: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'nvcuda.dll'; dlerror: nvcuda.dll not found
2022-08-09 19:48:32.847804: W tensorflow/stream_executor/cuda/cuda_driver.cc:269] failed call to cuInit: UNKNOWN ERROR (303)
2022-08-09 19:48:32.851014: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:169] retrieving CUDA diagnostic information for host: DEVELOPER
2022-08-09 19:48:32.851211: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:176] hostname: DEVELOPER
2022-08-09 19:48:32.851607: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
Traceback (most recent call last):
  File "D:\tf\create_model.py", line 29, in <module>
    model = pickle.load(fl)
  File "C:\Users\developer\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\saving\pickle_utils.py", line 48, in deserialize_model_from_bytecode
    model = save_module.load_model(temp_dir)
  File "C:\Users\developer\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\developer\AppData\Local\Programs\Python\Python39\lib\site-packages\tensorflow\python\saved_model\load.py", line 977, in load_internal
    raise FileNotFoundError(
FileNotFoundError: Unsuccessful TensorSliceReader constructor: Failed to find any matching files for ram://5488f35a-e52b-472b-b9d6-110c8b5a3aaf/variables/variables
 You may be trying to load on a different device from the computational device. Consider setting the `experimental_io_device` option in `tf.saved_model.LoadOptions` to the io_device such as '/job:localhost'.

核心原因 #

TensorFlow的Keras模型并非设计为直接用Pickle序列化，虽然Keras做了兼容处理，但跨平台（尤其是Windows和Linux）时会出现临时文件路径、设备IO的兼容性问题，报错中的ram://路径加载失败就是典型表现。

推荐方案：使用TensorFlow官方保存/加载方式 #

放弃Pickle，改用TensorFlow官方支持的模型持久化方法，这是最稳定的跨平台方案：

修改后的dump.py（保存模型） #

import tensorflow as tf
import numpy as np

tf.random.set_seed(42)

input_x = np.random.randint(0, 50000, (10000,1))
input_y = np.random.randint(0, 50000, (10000,1))
output = input_x + input_y
input = np.concatenate((input_x, input_y), axis=1)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(2, activation = tf.keras.activations.relu, input_shape=[2]),   
    tf.keras.layers.Dense(2, activation = tf.keras.activations.relu),
    tf.keras.layers.Dense(1),
])

model.compile(loss = tf.keras.losses.mae,
              optimizer=tf.optimizers.Adam(learning_rate=0.00001),
              metrics = ['mse'])
          
model.fit(input, output, epochs = 1000)

# 官方方法保存，会生成一个包含模型参数和配置的文件夹
model.save('D:/tf/tf_model')

修改后的load.py（加载模型） #

import tensorflow as tf

# 从官方保存的文件夹加载模型
model = tf.keras.models.load_model('D:/tf/tf_model')
print(model.predict([[2.2, 5.1]]))

兼容方案：必须用Pickle时的处理 #

如果因特殊限制必须使用Pickle，可以通过将模型先序列化到内存字节流，再用Pickle保存该字节流，加载时反向操作：

修改后的dump.py #

import tensorflow as tf
import pickle
import numpy as np
import io

tf.random.set_seed(42)

input_x = np.random.randint(0, 50000, (10000,1))
input_y = np.random.randint(0, 50000, (10000,1))
output = input_x + input_y
input = np.concatenate((input_x, input_y), axis=1)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(2, activation = tf.keras.activations.relu, input_shape=[2]),   
    tf.keras.layers.Dense(2, activation = tf.keras.activations.relu),
    tf.keras.layers.Dense(1),
])

model.compile(loss = tf.keras.losses.mae,
              optimizer=tf.optimizers.Adam(learning_rate=0.00001),
              metrics = ['mse'])
          
model.fit(input, output, epochs = 1000)

# 将模型保存到内存字节流
buffer = io.BytesIO()
tf.keras.models.save_model(model, buffer)
buffer.seek(0)

# 用Pickle保存字节流数据
with open('D:/tf/tf.pkl', 'wb') as fl:
    pickle.dump(buffer.getvalue(), fl)

修改后的load.py #

import tensorflow as tf
import pickle
import io

with open('D:/tf/tf.pkl', 'rb') as fl:
    model_bytes = pickle.load(fl)

# 从字节流加载模型
buffer = io.BytesIO(model_bytes)
buffer.seek(0)
model = tf.keras.models.load_model(buffer)

print(model.predict([[2.2, 5.1]]))

内容的提问来源于stack exchange，提问作者arthurbesse