Tensorflow,恢复特定设备中的变量

时间:2016-11-06 16:58:18

标签: tensorflow

也许我的问题有点天真,但我在tensorflow文档中找不到任何内容。

我有一个训练有素的张量流模型,其中变量放在GPU中。现在我想恢复此模型并使用CPU进行测试。

如果我通过示例中的' tf.train.Saver.restore`执行此操作: saver = tf.train.import_meta_graph("/tmp/graph.meta") saver.restore(session, "/tmp/model.ckp")

我有以下重点:

InvalidArgumentError: Cannot assign a device to node 'b_fc8/b_fc8/Adam_1': Could not satisfy explicit device specification '/device:GPU:0' because no devices matching that specification are registered in this process; available devices: /job:localhost/replica:0/task:0/cpu:0

如何在CPU

中恢复这些变量

由于

2 个答案:

答案 0 :(得分:5)

使用clear_devices标志,即

saver = tf.train.import_meta_graph("/tmp/graph.meta", clear_devices=True)

答案 1 :(得分:1)

我正在使用tensorflow 0.12和clear_devices=True并且tf.device('/cpu:0')没有与我合作(saver.restore仍在尝试将变量分配给/ gpu:0)。

我真的需要强制一切到/ cpu:0因为我加载了几个不适合GPU内存的模型。以下是强制所有内容的两种选择:/ cpu:0

  1. 设置os.environ['CUDA_VISIBLE_DEVICES']=''
  2. 使用ConfigProto的device_count,例如tf.Session(config=tf.ConfigProto(device_count={"GPU": 0, "CPU": 1}))