如何重写张量流图以将CPU用于所有操作

Question

我已经在多GPU和CPU设置上训练了一个网络，并将结果模型保存为tensorflow SavedModel。然后，我还有另一个脚本，可以加载生成的模型并运行所需的操作以进行预测，即对模型进行推断。这适用于训练模型的相同设置。

但是，我需要将模型部署为在具有1个CPU且没有GPU的设备上运行。当我尝试运行相同的脚本时，出现以下错误：

InvalidArgumentError（请参阅上面的回溯）：从检查点还原失败。这很可能是由于当前图形与来自检查点的图形之间的不匹配。请确保您没有更改基于检查点的预期图形。原始错误：

无法为操作default_policy_1 / tower_1 / variable分配设备：已将节点default_policy_1 / tower_1 / variable（在restore.py:56定义）明确分配给/ device：GPU：0，但可用设备为[/ job：localhost / replica：0 / task：0 / device：CPU：0，/ job：localhost /副本：0 / task：0 / device：XLA_CPU：0]。确保设备规格引用的是有效设备。所请求的设备似乎是GPU，但是未启用CUDA。 [[node default_policy_1 / tower_1 / Variable（在restore.py:56处定义）]]

这看起来很有希望，但是代码根本没有改变我的图，删除了 0个节点- 删除了在CPU上运行的操作图tensorflow

通常，删除所有不在CPU上运行的操作似乎并不明智。

我尝试将所有内容包装在一个with tf.device('CPU:0')块中，以及使用config = tf.ConfigProto(device_count={'GPU': 0})但均未更改错误。

相关代码：

from tensorflow.python.saved_model import loader

input_tensor_key_feed_dict = {'observations': np.array([[23]]), 'prev_action': np.array([0]),
                              'prev_reward': np.array([0]), 'is_training': False}

config = tf.ConfigProto(device_count={'GPU': 0})
with tf.device('CPU:0'):
    with session.Session(None, graph=ops_lib.Graph(), config=config) as sess:

        loader.load(sess, tag_set.split(','), saved_model_dir) #error occurs here

        outputs = sess.run(output_tensor_names_sorted, feed_dict=inputs_feed_dict)
        for i, output in enumerate(outputs):
            output_tensor_key = output_tensor_keys_sorted[i]
            print('Result for output key %s:\t%s' % (output_tensor_key, output))

Answer 1

我将初始化一个没有设备规范的新模型，然后仅加载模型变量，因为它将是 tf.Saver() 的标准训练检查点。此时，您应该能够保存 SavedModel 的版本，tensorflow 可以决定将操作放置在何处。