Python 删除要在CPU上运行的操作图tensorflow
我已经训练了一个网络(使用GPU),现在我想在CPU上运行它(用于推理)。为此,我使用以下代码加载元图,然后加载网络的参数Python 删除要在CPU上运行的操作图tensorflow,python,tensorflow,gpu,Python,Tensorflow,Gpu,我已经训练了一个网络(使用GPU),现在我想在CPU上运行它(用于推理)。为此,我使用以下代码加载元图,然后加载网络的参数 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 =
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)
问题是,由于图形是为训练而定义的,所以我使用了一些不在CPU上运行的特定操作。例如,记录GPU活动的“MaxBytesInUse”
这就是为什么当我尝试运行此代码时,会出现以下错误:
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"]()]]
有没有一种简单的方法可以删除特定的GPU相关操作并在CPU上运行图形?我认为这样的方法应该可以解决您的问题
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内核,检查内核定义的
约束
字段,您可以使可以在cpu上运行
更加准确,但这对于您的情况可能已经足够了。我还在tf.train.import\u meta\u graph
中添加了一个clear\u devices=True
,它在操作中清除设备指令,强制它们在特定设备上运行(如果图形中有任何指令)。您可能会发现这篇相关文章很有用: