TensorFlow GPU版本序列化模型触发FileNotFoundError求助

TensorFlow CPU版本下运行以下代码可正常执行，但切换到GPU版本后，用dill、joblib或pickle加载序列化模型时均出现FileNotFoundError。

复现代码 #

import tensorflow as tf
import dill

def Generator():

    z_dim = 60

    FEATURES_LIST = ["aaa", "bbb", "ccc" ]
    ME_FEATURES_LIST = ["ddd", "eee", "fff" ]

    NUM_FEATURES = len(FEATURES_LIST)
    NUM_ME_FEATURES = len(ME_FEATURES_LIST)

    z = tf.keras.layers.Input(shape=(z_dim,), dtype='float32')
    y = tf.keras.layers.Input(shape=(NUM_ME_FEATURES,), dtype='float32')
    tr = tf.keras.layers.Input(shape=(1,), dtype='bool')
  
    x = tf.keras.layers.concatenate([z, y])
    x = tf.keras.layers.Dense(z_dim * NUM_ME_FEATURES, activation="relu")(x)
    out = tf.keras.layers.Dense(NUM_FEATURES, activation='sigmoid')(x)
    model = tf.keras.Model(inputs=[z, y, tr], outputs=(out, y))

    return model

G = Generator()

with open("dill_functional", 'wb') as file:
  dill.dump(G, file)

with open("dill_functional", 'rb') as file:
  G = dill.load(file)  # <--- 报错位置

print(str(G))

报错信息 #

执行输出及报错：

C:\Users\igor-.cloned\gan> python .\dill_test.py 2023-02-09
22:42:28.379108: 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 AVX2 To enable them in other
operations, rebuild TensorFlow with the appropriate compiler flags.

2023-02-09 22:42:29.759547: I
tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device
/job:localhost/replica:0/task:0/device:GPU:0 with 9426 MB memory:  ->
device: 0, name: NVIDIA GeForce RTX 3080 Ti, pci bus id: 0000:01:00.0,
compute capability: 8.6

WARNING:tensorflow:Compiled the loaded model,
but the compiled metrics have yet to be built. `model.compile_metrics`
will be empty until you train or evaluate the model.

Traceback (most recent call last):
 File "C:\Users\igor-\.cloned\gan\dill_test.py", line 32, in <module>
  G = dill.load(file)   
 File "C:\Users\igor-\anaconda3\envs\ai\lib\site-packages\dill\_dill.py", line 272, in load
  return Unpickler(file, ignore=ignore, **kwds).load()   
 File "C:\Users\igor-\anaconda3\envs\ai\lib\site-packages\dill\_dill.py", line 419, in load
  obj = StockUnpickler.load(self)   
 File "C:\Users\igor-\anaconda3\envs\ai\lib\site-packages\keras\saving\pickle_utils.py", line 47, in deserialize_model_from_bytecode
  model = save_module.load_model(temp_dir)   
 File "C:\Users\igor-\anaconda3\envs\ai\lib\site-packages\keras\utils\traceback_utils.py", line 70, in error_handler
  raise e.with_traceback(filtered_tb) from None   
 File "C:\Users\igor-\anaconda3\envs\ai\lib\site-packages\tensorflow\python\saved_model\load.py", line 933, in load_partial
  raise FileNotFoundError(
FileNotFoundError: Unsuccessful TensorSliceReader constructor: Failed to find any matching files for  ram://fc47ea82-4f6b-4736-9394-980cc1f14358/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'.

1. 优先使用Keras原生模型保存/加载方法 #

Keras官方不推荐用pickle、dill这类通用序列化工具保存模型，它们无法妥善处理TensorFlow的设备依赖和内部状态。改用原生方法：

保存模型 #

# 替换原dill.dump代码
G.save("generator_model")  # 保存为文件夹格式

加载模型 #

# 替换原dill.load代码
G = tf.keras.models.load_model("generator_model")

2. 若必须用pickle/dill，指定IO设备 #

如果因特殊需求必须使用pickle/dill，在加载时强制指定IO设备为本地主机，规避GPU设备的IO冲突：

修改加载代码如下：

import tensorflow as tf
import dill
from keras.saving import save_module

# 定义加载选项
load_options = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')

# 临时patch keras的load_model函数，注入加载选项
original_load_model = save_module.load_model
def patched_load_model(path, **kwargs):
    kwargs['options'] = load_options
    return original_load_model(path, **kwargs)
save_module.load_model = patched_load_model

# 执行加载
with open("dill_functional", 'rb') as file:
    G = dill.load(file)

# 恢复原函数
save_module.load_model = original_load_model

3. 临时强制使用CPU加载（调试用） #

如果只是临时验证功能，可以在加载前强制TensorFlow使用CPU：

import tensorflow as tf
# 隐藏所有GPU设备，强制操作在CPU执行
tf.config.set_visible_devices([], 'GPU')

# 执行加载代码
with open("dill_functional", 'rb') as file:
    G = dill.load(file)

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