python - 卡住具有不同数据格式的图形

Question

我使用 NCHW 数据格式在我的 GPU 上训练了一个小型 CNN，现在我想导出一个 .pb 文件，然后我可以用它在其他应用程序中进行推理.

我写了一个小的辅助函数来使用默认值调用 Tensorflow 的 freeze_graph 函数，给定一个包含检查点文件和 graph.pbtxt 的目录:

import os
import argparse
#os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
dir(tf.contrib) #fix for tf.contrib undefined ops bug
from tensorflow.python.tools.freeze_graph import freeze_graph 

def my_freeze_graph_2(model_dir, output_node_names):
"""Extract the sub graph defined by the output nodes and convert 
  all its variables into constant 
  Args:
      model_dir: the root folder containing the checkpoint state file
      output_node_names: a string, containing all the output node's names, 
                          comma separated
"""
if not tf.gfile.Exists(model_dir):
    raise AssertionError(
        "Export directory doesn't exists. Please specify an export "
        "directory: %s" % model_dir)

if not output_node_names:
    print("You need to supply the name of a node to --output_node_names.")
    return -1

# We retrieve our checkpoint fullpath
checkpoint = tf.train.get_checkpoint_state(model_dir)
input_checkpoint = checkpoint.model_checkpoint_path

# We precise the file fullname of our freezed graph
absolute_model_dir = os.path.abspath(model_dir)
output_graph = os.path.join(absolute_model_dir, "frozen_model.pb")

freeze_graph(input_graph=os.path.join(model_dir, 'graph.pbtxt'),
             input_saver='',
             input_binary=False,
             input_checkpoint=input_checkpoint,
             output_node_names=output_node_names,
             restore_op_name="save/restore_all",
             filename_tensor_name="save/Const:0",
             output_graph=output_graph,
             clear_devices=True,
             initializer_nodes='')

然后我有一个小脚本试图从 frozen_model.pb 构建图表来测试卡住是否有效:

import os
#os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import argparse
import tensorflow as tf
from freeze_graph import load_graph
import cv2

if __name__ == '__main__':
    # Let's allow the user to pass the filename as an argument
    parser = argparse.ArgumentParser()
parser.add_argument("--frozen_model_filename", default="model-multiple_starts/frozen_model.pb", type=str, help="Frozen model file to import")
args = parser.parse_args()

# We use our "load_graph" function
graph = load_graph(args.frozen_model_filename)

# We can verify that we can access the list of operations in the graph
for op in graph.get_operations():
    print(op.name)

# We access the input and output nodes 
x = graph.get_tensor_by_name('prefix/Reshape:0')
y = graph.get_tensor_by_name('prefix/softmax_tensor:0')

# We launch a Session
with tf.Session(graph=graph, config=tf.ConfigProto(log_device_placement=True)) as sess:
    # Note: we don't nee to initialize/restore anything
    # There is no Variables in this graph, only hardcoded constants 

    # Load an image to use as test
    im = cv2.imread('57_00000000.png', cv2.IMREAD_GRAYSCALE)
    im = im.T
    im = im / 255 - 0.5
    im = im[None,:,:,None]


    y_out = sess.run(y, feed_dict={
        x: im 
    })
    print(y_out)

如果我尝试运行我的测试脚本，我会收到以下错误:

InvalidArgumentError: CPU BiasOp only supports NHWC. [[Node: prefix/conv2d/BiasAdd = BiasAdd[T=DT_FLOAT, data_format="NCHW", _device="/job:localhost/replica:0/task:0/cpu:0"](prefix/conv2d/convolution, prefix/conv2d/bias/read)]]

我尝试了不同的配置:

• 从 CPU-only 脚本生成 .pb 文件，在 CPU-only 上运行
• 从脚本生成 GPU 可见的 .pb 文件，在 GPU 可见的情况下运行
• 从纯 CPU 脚本生成 .pb 文件，在 GPU 可见的情况下运行

它们都引发相同的错误。

问题在于我要卡住的检查点具有使用 data_format='NCHW' 定义的操作。如何使用 NHWC 数据格式卡住检查点？

更新:

翻阅这些文件，我发现在 graph.pbtxt 中，对于许多操作，data_format 被硬编码为 NCHW。我想，然后，我需要制作一个具有 NHWC 格式的新模型，有选择地从检查点加载层的权重并使用该图手动保存 .pb 文件... 我假设已经有一个过程可以做到这一点，但我找不到关于这个的任何文档，也找不到示例。

更新 2:

尝试在 OpenCV 的 DNN 模块中导入 .pb 文件后，我发现了以下内容:

• 将数据格式为 NCHW 的训练中的检查点和数据格式为 NHWC 的 graph.pbtxt 一起卡住会导致无法使用 .pb 文件。我还没有找到确切的原因，但是将 .pb 转换为 .pbtxt 并将其与工作卡住图进行比较，文件仅在存储的值上有所不同权重和偏差常数。
• 将来自训练的检查点和数据格式为 NHWC 的 graph.pbtxt 一起卡住，生成一个工作卡住图。

那么，检查点似乎不能在具有不同数据格式的图之间转移(即使在卡住过程中没有出现错误或警告)。

Answer 1

通常，您会希望将图形构建包装在函数中，以便您可以有条件地为预测案例重建图形，因为通常有相当多的图形从训练到预测会发生变化。正如您所发现的，NCHW 和 NWHC 版本，例如卷积层，实际上是图形原型(prototype)中的不同 Ops，并且它们以这种方式进行硬编码，因为 GPU 优化是仅适用于其中一种格式。

正确编辑图形原型(prototype)非常困难，这就是为什么大多数执行此操作的 TensorFlow 代码都遵循我上面描述的模式。在一个非常高的层次上:

def build_graph(data_format='NCHW'):
   # Conditionally use proper ops based on data_format arg

training_graph = tf.Graph()
with training_graph.as_default():
   build_graph(data_format='NCHW')

with tf.Session() as sess:
   # train
   # checkpoint session

prediction_graph = tf.Graph()
with prediction_graph.as_default():
   build_graph(data_format='NHWC')
   # load checkpoint
   # freeze graph

请注意，tf.estimator.Estimator 框架使这变得相对容易。您可以在 model_fn 中使用 mode 参数来决定数据格式，然后使用两个不同的 input_fn 进行训练和预测，以及框架会做剩下的。你可以在这里找到一个端到端的例子:https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10_estimator/cifar10_main.py#L77 (我已经链接到相关行)