Question

我正在使用给定模型运行TF应用程序进行推理。但是，它不是在GPU上运行，而是在CPU上运行，尽管在启用CUDA的情况下构建了tensorflow库。要了解TF模型，张量流模型（.pb）是否有设备信息，如tf.device（/ cpu：0）或tf.device（/ gpu：0）???

Answer 1

来自docs（强调我的）：

有时导出的元图来自导入器没有的培训环境。例如，该模型可能已在GPU上或在具有副本的分布式环境中进行过培训。导入此类模型时，能够清除图形中的设备设置非常有用，这样我们就可以在本地可用的设备上运行它。 这可以通过在import_meta_graph设置为clear_devices option 的情况下致电True来实现。
with tf.Session() as sess:
  new_saver = tf.train.import_meta_graph('my-save-dir/my-model-10000.meta',
      clear_devices=True)
  new_saver.restore(sess, 'my-save-dir/my-model-10000')

Answer 2

将GraphDef加载到tf.Graph之后，使用_set_device API将所有操作移至CPU。 https://github.com/tensorflow/tensorflow/blob/r1.14/tensorflow/python/framework/ops.py#L2255

gf = tf.GraphDef()
gf.ParseFromString(open('graph.pb','rb').read()) 
with tf.Session() as sess: 
    tf.import_graph_def(gf, name='')  
    g = tf.get_default_graph() 
    ops = g.get_operations() 
    for op in ops: 
        op._set_device('/device:CPU:*')

Tensorflow模型（.pb）有设备信息吗？

2 个答案: