如何将 tensorflow.keras 模型移动到 GPU答案

【问题标题】：How to move a tensorflow.keras model to GPU如何将 tensorflow.keras 模型移动到 GPU
【发布时间】：2020-04-24 06:33:47
【问题描述】：

假设我有一个这样的 keras 模型：

with tf.device("/CPU"):
    model = tf.keras.Sequential([
    # Adds a densely-connected layer with 64 units to the model:
    tf.keras.layers.Dense(64, activation='relu', input_shape=(32,)),
    # Add another:
    tf.keras.layers.Dense(64, activation='relu'),
    # Add a softmax layer with 10 output units:
    tf.keras.layers.Dense(10, activation='softmax')])

我想将此模型移至 GPU。

我试过这样做：

with tf.device("/GPU:0"):
    gpu_model = tf.keras.models.clone_model(model)

但问题在于，变量名称发生了变化。例如：

第一层权重model的名字是： Got from model.layers[0].weights[0].name

'密集/内核：0'

但是第一层的权重名称gpu_model是： Got from gpu_model.layers[0].weights[0].name

'dense_3/kernel:0'

如何在保留变量名称的同时进行这种 GPU 转换？

我不想将模型保存到磁盘并再次加载

【问题讨论】：

是我的想法还是变量名相同？
@MatiasValdenegro 对不起，我复制粘贴错了。我现在编辑了它

标签： python tensorflow keras namespaces

【解决方案1】：

我正在回答我自己的问题。如果有人有更好的解决方案。请贴出来

这是我找到的解决方法：

像 PyTorch 一样创建 state_dict
以 JSON 格式获取模型架构
清除 Keras 会话并删除模型实例
从 tf.device 上下文中的 JSON 创建一个新模型
从 state_dict 加载之前的权重

state_dict = {}
for layer in model.layers:
    for weight in layer.weights:
        state_dict[weight.name] = weight.numpy()

model_json_config = model.to_json()
tf.keras.backend.clear_session() # this is crucial to get previous names again
del model

with tf.device("/GPU:0"):
    new_model = tf.keras.models.model_from_json(model_json_config)

for layer in new_model.layers:
    current_layer_weights = []
    for weight in layer.weights:
        current_layer_weights.append(state_dict[weight.name])
    layer.set_weights(current_layer_weights)

【讨论】：