如何加载经过训练的 TensorFlow 模型以使用不同的批量大小进行预测?

问题描述

这应该是一个非常简单的问题，但确实吸引了我.

This shall be a quite simple question but indeed caught me.

我的目的是加载一个经过训练的 TF 模型，然后创建一个会话来运行预测/推理，并为我的设备环境评估合适的 batch_size.

My purpose is to load a trained TF model, then create a session to run the predict/inference, and evaluate an appropriate batch_size for my device environment.

我准备了两种模型:冷冻模型和保存模型.

I prepared two types of model: frozen model and saved model.

对于冻结模型，我通过.ParseFromString()成功加载并获取到GraphDef，并通过TensorRT优化了GraphDef，但输入节点的batch_size固定为1(1*299*299*3).好像在导出然后冻结模型时无法配置batch_size，之后也无法更改，因为它是append-only.

For the frozen model, I successfully loaded and get the GraphDef by .ParseFromString(), and optimized the GraphDef by TensorRT, but the batch_size of input node is fixed to 1 (1*299*299*3). Seems like the batch_size can't be configured when exporting then freeze model, and can't be changed afterwards because it's append-only.

对于保存的模型，输入节点的维度是?(?*299*299*3).似乎它应该能够接受任何大小的批次.但是在我加载保存的模型后，然后输入一个 299*299*3 的图像，并得到以下错误:

For the saved model, the dimension of input node is ? (?*299*299*3). Seems it shall be able to accept any size of batch. But after I loaded the saved model, then input a 299*299*3 image, and get below ERROR:

ValueError: Cannot feed value of shape (299, 299, 3) for Tensor u'input:0', which has shape '(?, 299, 299, 3)'

我尝试将轴=0 扩展到输入图像，但仍然出错:

I tried to expand a axis=0 to the input image, but still get error:

InvalidArgumentError (see above for traceback): You must feed a value for placeholder tensor 'import/input' with dtype float and shape [?,299,299,3]

 [[Node: import/input = Placeholder[dtype=DT_FLOAT, shape=[?,299,299,3], _device="/job:localhost/replica:0/task:0/device:GPU:0"]()]]

总的来说，这应该是一个非常简单的问题，但我不知道如何解决.

Overall, this is supposed to be a quite simple question, but I'm not sure how to solve it.

欢迎提出任何想法.

谢谢，

有冻结模型的脚本，占位符固定为(1*299*299*3)，所以只接受args.batch_size"为1.

The script with frozen model, the placeholder is fixed to (1*299*299*3), so it only accepts "args.batch_size" to be 1.

    g = ops.Graph()
with g.as_default():
    inp, out = importer.import_graph_def(graph_def=f32_graph, return_elements=['Placeholder','InceptionV3/Logits/SpatialSqueeze'])
    inp = inp.outputs[0]
    out = out.outputs[0]
    print tf.shape(inp)
    if args.image:
        import numpy as np
        func = TestKit.preprocess_func['tensorflow'][args.network]
        img = func(args.image)
        img = np.expand_dims(img, axis = 0)
    batch_input = img
    for i in range(int(args.batch_size)-1):
        batch_input = np.concatenate((batch_input, img), axis=0)
    print len(batch_input)
    gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.625)
    with csess.Session(config=cpb2.ConfigProto(gpu_options=gpu_options), graph=g) as sess:
        t0 = time.time()
        for _ in range(2000):
            val = sess.run(out, {inp: batch_input})
            predict = np.squeeze(val)
            top_indices = predict.argsort()[-5:][::-1]
            result = [(i, predict[i]) for i in top_indices]
        t1 = time.time()
        print result
        print 'Duration:', str(t1-t0)

保存模型的脚本，input:0"节点为(?*299*299*3)，但不能接受上述任何形状的输入图像.

---

The script of saved model, "input:0" node is (?*299*299*3) but can't accept any shape of input image as mentioned above.

    with tf.Graph().as_default():
    with tf.Session(graph=tf.Graph()) as sess:
        tf.saved_model.loader.load(sess, ['serve'], './')
        for op in sess.graph.get_operations():
            print str(op.name)

        if args.image:
            import numpy as np
            func = TestKit.preprocess_func['tensorflow'][args.network]
            img = func(args.image)
            img = np.expand_dims(img, axis = 0)
        inp, out = importer.import_graph_def(graph_def=sess.graph.as_graph_def(), return_elements=['input','InceptionV3/Logits/SpatialSqueeze'])
        t0 = time.time()
        for _ in range(1000):
            val = sess.run(out, {'input:0':img})
            predict = np.squeeze(val)
            top_indices = predict.argsort()[-5:][::-1]
            result = [(i, predict[i]) for i in top_indices]
        t1 = time.time()
        print result
        print 'Duration:', str(t1-t0)

推荐答案

我终于解决了这个问题.重点是冻结模型不会改变批量大小，因为图是仅附加的，所以一旦原始图使用输入占位符(1 * 299 * 299 * 3)生成，它将永远不会改变.

I finally solved the problem.The point is freeze model will not change the batch size because the graph is append-only, so once the original graph is generated with input placeholder (1 * 299 * 299 * 3), it will never be changed.

最后，我尝试使用输入占位符(无 * 299 * 299 * 3)重新生成原始二进制图(二进制 .pbtxt).然后将此二元图转换为冻结模型，冻结模型就可以有维度为(?* 299 * 299 * 3)的输入占位符.现在支持批量输入.

Finally, I tried to regenerate the original binary graph (binary .pbtxt) with input placeholder (None * 299 * 299 * 3). Then convert this binary graph to freeze model, the freeze model then can have the input placeholder with dimension (? * 299 * 299 * 3). Batch input is supported now.

谢谢，

这篇关于如何加载经过训练的 TensorFlow 模型以使用不同的批量大小进行预测?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！