如何将tensorflow.keras模型移动到GPU [英] How to move a tensorflow.keras model to GPU

问题描述

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

Let's say I have a keras model like this:

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.

I would like to move this model to GPU.

我尝试这样做:

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

但是问题在于，变量名称会更改.例如:

But the problem with this is that, the variable names change. For example:

第一层重量的名称为model是:从model.layers[0].weights[0].name获得

The first layer's weight's name of model is: Got from model.layers[0].weights[0].name

'dense/kernel:0'

'dense/kernel:0'

但是第一层权重的名称为gpu_model是:从gpu_model.layers[0].weights[0].name

But the first layer's weight's name of gpu_model is: Got from gpu_model.layers[0].weights[0].name

'dense_3/kernel:0'

'dense_3/kernel:0'

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

How can I do this GPU transformation while also preserving the names of the variables?

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

I don't want to save the model to disk and load again

推荐答案

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

I am answering my own question. If someone has a better solution. Kindly post it

这是我发现的解决方法:

This is a work around I found:

1. 创建一个像PyTorch这样的state_dict
2. 获取模型架构为JSON
3. 清除Keras会话并删除模型实例
4. 在tf.device上下文中使用JSON创建新模型
5. 从state_dict加载先前的权重

1. Create a state_dict like PyTorch
2. Get the model architecture as JSON
3. Clear the Keras session and delete the model instance
4. Create a new model from the JSON within tf.device context
5. Load the previous weights from 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)

这篇关于如何将tensorflow.keras模型移动到GPU的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！