在GPU上训练模型并在CPU上进行预测时如何设置设备

Question

我在GPU上训练了一个模型，并将其保存为以下格式（export_path是我的输出目录）

builder = tf.saved_model.builder.SavedModelBuilder(export_path)

tensor_info_x = tf.saved_model.utils.build_tensor_info(self.Xph)
tensor_info_y = tf.saved_model.utils.build_tensor_info(self.predprob)
tensor_info_it = tf.saved_model.utils.build_tensor_info(self.istraining)
tensor_info_do = tf.saved_model.utils.build_tensor_info(self.dropout)

prediction_signature = (
       tf.saved_model.signature_def_utils.build_signature_def(
              inputs={'myx': tensor_info_x, 'istraining': tensor_info_it, 'dropout': tensor_info_do},
              outputs={'ypred': tensor_info_y},
              method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

builder.add_meta_graph_and_variables(
       net, [tf.saved_model.tag_constants.SERVING],
       signature_def_map={
           tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
           prediction_signature },)
builder.save()

现在我正在尝试加载并运行预测。如果我使用GPU，它可以正常工作，但是如果没有GPU，我会得到以下错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation 'rnn/while/rnn/multi_rnn_cell/cell_0/cell_0/layer_norm_basic_lstm_cell/dropout/add/Enter': Operation was explicitly assigned to /device:GPU:0 but available devices are [ /job:localhost/replica:0/task:0/device:CPU:0 ]. Make sure the device specification refers to a valid device.

现在我正在阅读有关tf.train.import_meta_graph和clear_device选项的内容，但我无法使其工作。我是这样加载我的模型的：

predict_fn = predictor.from_saved_model(modelname)

在哪个时刻会抛出上述错误？modelname是pb文件的完整文件名。有没有一种方法可以遍历图中的节点并手动设置设备（或执行类似操作）？我正在使用tensorflow 1.8.0。

我看到Can a model trained on gpu used on cpu for inference and vice versa?，我不认为我重复了那个问题。与那个问题的区别在于，我想知道在训练之后要做什么。

- hadron

Answer 1

最终我在我的GPU机器上使用“clear_devices=True”重新保存了模型，然后将保存的模型移动到我的CPU机器上。我无法找到任何具体的解决方案，因此在下面发布我的脚本：

import tensorflow as tf

with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, tf.saved_model.tag_constants.SERVING], m)
    loaded_graph = tf.get_default_graph()
    x = loaded_graph.get_tensor_by_name('myx:0')
    dropout = loaded_graph.get_tensor_by_name('mydropout:0')
    y = loaded_graph.get_tensor_by_name('myy:0')
    export_path = 'somedirectory'
    builder = tf.saved_model.builder.SavedModelBuilder(export_path + '/mymodel')
    tensor_info_x = tf.saved_model.utils.build_tensor_info(x)
    tensor_info_y = tf.saved_model.utils.build_tensor_info(y)
    tensor_info_do = tf.saved_model.utils.build_tensor_info(dropout)
    prediction_signature = (
                            tf.saved_model.signature_def_utils.build_signature_def(
                            inputs={'myx': tensor_info_x, 'mydropout': tensor_info_do},
                            outputs={'myy': tensor_info_y}, 
                            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

    builder.add_meta_graph_and_variables(sess, [tf.saved_model.tag_constants.SERVING],
                   signature_def_map={tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                       prediction_signature }, clear_devices=True)
    builder.save()

- hadron