Question

我已经训练了一个网络（使用GPU），现在我想在CPU上运行它（以进行推理）。为此，我使用以下代码加载元图，然后加载网络的参数。

for i in list(content):
    s = "\\x" + format(ord(i), 'X')

问题在于，由于已为训练定义了图表，因此我使用了一些无法在CPU上运行的特定操作。例如，“ MaxBytesInUse” https://www.tensorflow.org/api_docs/python/tf/contrib/memory_stats/MaxBytesInUse记录了GPU的活动。

这就是为什么当我尝试运行此代码时，出现以下错误：

config = tf.ConfigProto(
        device_count = {'GPU': 0}
    )
sess = tf.Session(config=config)
meta_graph=".../graph-0207-190023.meta"
model=".../model.data-00000-of-00001"

new_saver = tf.train.import_meta_graph(meta_graph)
new_saver.restore(sess, model)

是否有一种简单的方法可以删除与GPU相关的特定操作并在CPU上运行图形？

Answer 1

我认为这样可以解决您的问题

import tensorflow as tf

def remove_no_cpu_ops(graph_def):
    # Remove all ops that cannot run on the CPU
    removed = set()
    nodes = list(graph_def.node)
    for node in nodes:
        if not can_run_on_cpu(node):
            graph_def.node.remove(node)
            removed.add(node.name)
    # Recursively remove ops depending on removed ops
    while removed:
        removed, prev_removed = set(), removed
        nodes = list(graph_def.node)
        for node in graph_def.node:
            if any(inp in prev_removed for inp in node.input):
                graph_def.node.remove(node)
                removed.add(node.name)

def can_run_on_cpu(node):
    # Check if there is a CPU kernel for the node operation
    from tensorflow.python.framework import kernels
    for kernel in kernels.get_registered_kernels_for_op(node.op).kernel:
        if kernel.device_type == 'CPU':
            return True
    return False

config = tf.ConfigProto(
        device_count = {'GPU': 0}
    )
sess = tf.Session(config=config)
meta_graph = ".../graph-0207-190023.meta"
model = ".../model.data-00000-of-00001"
# Load metagraph definition
meta_graph_def = tf.MetaGraphDef()
with open(meta_graph, 'rb') as f:
    meta_graph_def.MergeFromString(f.read())
# Remove GPU ops
remove_no_cpu_ops(meta_graph_def.graph_def)
# Make saver from modified metagraph definition
new_saver = tf.train.import_meta_graph(meta_graph_def, clear_devices=True)
new_saver.restore(sess, model)

这个想法是您遍历图定义中的所有节点，并删除没有CPU内核的那些节点。实际上，您可以通过检查是否存在适用于节点操作和输入类型的CPU内核，检查内核定义的can_run_on_cpu字段来使constraint更加准确，但这可能很好足以应付您的情况。我还向clear_devices=True添加了tf.train.import_meta_graph，以清除强制其在特定设备上运行的操作中的设备指令（如果您的图形中有任何这些指令）。

删除操作图张量流以在CPU上运行

1 个答案: