CODEWITHSUNDEEP
pythontensorflowgpu
Stringer Bell
Stringer Bell

Reputation: 161

Remove operation graph tensorflow to run on CPU

I have trained a network (using GPU) and now I want to run it (for inference) on a CPU. To do so, I use the following code which loads the meta graph and then the parameters of the network.

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)

Problem is that since the graph has been defined for training, I have used some specific operations that do not run on CPU. For example "MaxBytesInUse" https://www.tensorflow.org/api_docs/python/tf/contrib/memory_stats/MaxBytesInUse which records the GPU activity.

Thats is why, when I try to run this code, I get the following error :

InvalidArgumentError: No OpKernel was registered to support Op 'MaxBytesInUse' with these attrs.  Registered devices: [CPU], Registered kernels:
  device='GPU'

     [[Node: PeakMemoryTracker/MaxBytesInUse = MaxBytesInUse[_device="/device:GPU:0"]()]]

Is there a simple way to remove the specific GPU related operations and to run the graph on a CPU ?

Upvotes: 1

Views: 1854

Answers (1)

javidcf
javidcf

Reputation: 59741

I think something like this should solve your problem

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)

The idea is that you iterate through all the nodes in the graph definition and remove those without a CPU kernel. In reality, you could make can_run_on_cpu more accurate by checking that there is a CPU kernel that works for the node operation and input types, checking the constraint field of the kernel definition, but this will probably be good enough for your case. I also added a clear_devices=True to tf.train.import_meta_graph, which clears device directives in operations that force them to run on a specific device (in case you had any of those in your graph).

Upvotes: 1

Related Questions