Tensoflow: Как я могу добавить узел декодера изображения на мой график?

Question

У меня есть модель тензорного потока в виде замороженного графа, который принимает тензор изображения в качестве входных данных. Однако я хотел бы добавить новый узел декодера входного изображения к этому графику, чтобы модель также принимала закодированную строку байтов изображения jpg и в конечном итоге декодировала изображение сама. Я пробовал этот подход до сих пор:

model = './frozen_graph.pb'

with tf.gfile.FastGFile(model, 'rb') as f:

    # read graph
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name="")
    g = tf.get_default_graph()

    # fetch old input
    old_input = g.get_tensor_by_name('image_tensor:0')

    # define new input
    new_input = graph_def.node.add()
    new_input.name = 'encoded_image_string_tensor'
    new_input.op = 'Substr'
    # add new input attr
    image = tf.image.decode_image(new_input, channels=3)

    # link new input to old input
    old_input.input = 'encoded_image_string_tensor'  #  must match with the name above

Приведенный выше код возвращает это исключение:

Expected string passed to parameter 'input' of op 'Substr', got name: "encoded_image_string_tensor" op: "Substr"  of type 'NodeDef' instead.

Я не совсем уверен, смогу ли я использовать tf.image.decode_image в графике, так что, возможно, есть другой способ решения этой проблемы. У кого-нибудь есть подсказка?

Answer 1

Благодаря jdehesa , который дал мне хороший совет, я смог решить эту проблему. С параметром input_map я успешно отобразил новый график, который просто декодирует изображения jpg, на вход моего исходного графика (здесь: node.name='image_tensor:0'). Просто убедитесь, что вы переименовали name_scope графа декодера (здесь: decoder). После этого вы можете сохранить новый составной граф, используя тензор потока SavedModelBuilder.

Вот пример сети обнаружения объектов:

import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants


# The export path contains the name and the version of the model
model = 'path/to/model.pb'
export_path = './output/dir/'

sigs = {}

with tf.gfile.FastGFile(model, 'rb') as f:
        with tf.name_scope('decoder'):
                image_str_tensor = tf.placeholder(tf.string, shape=[None], name= 'encoded_image_string_tensor')
                # The CloudML Prediction API always "feeds" the Tensorflow graph with
                # dynamic batch sizes e.g. (?,).  decode_jpeg only processes scalar
                # strings because it cannot guarantee a batch of images would have
                # the same output size.  We use tf.map_fn to give decode_jpeg a scalar
                # string from dynamic batches.
                def decode_and_resize(image_str_tensor):
                        """Decodes jpeg string, resizes it and returns a uint8 tensor."""
                        image = tf.image.decode_jpeg(image_str_tensor, channels=3)

                        # do additional image manipulation here (like resize etc...)

                        image = tf.cast(image, dtype=tf.uint8)
                        return image

                image = tf.map_fn(decode_and_resize, image_str_tensor, back_prop=False, dtype=tf.uint8)

        with tf.name_scope('net'):
                # load .pb file
                graph_def = tf.GraphDef()
                graph_def.ParseFromString(f.read())

                # concatenate decoder graph and original graph
                tf.import_graph_def(graph_def, name="", input_map={'image_tensor:0':image})
                g = tf.get_default_graph()

with tf.Session() as sess:
        # load graph into session and save to new .pb file

        # define model input
        inp = g.get_tensor_by_name('decoder/encoded_image_string_tensor:0')

        # define model outputs
        num_detections = g.get_tensor_by_name('num_detections:0')
        detection_scores = g.get_tensor_by_name('detection_scores:0')
        detection_boxes = g.get_tensor_by_name('detection_boxes:0')
        out = {'num_detections': num_detections, 'detection_scores': detection_scores, 'detection_boxes': detection_boxes}


        builder = tf.saved_model.builder.SavedModelBuilder(export_path)

        tensor_info_inputs = {
                'inputs': tf.saved_model.utils.build_tensor_info(inp)}
        tensor_info_outputs = {}
        for k, v in out.items():
                tensor_info_outputs[k] = tf.saved_model.utils.build_tensor_info(v)

        # assign detection signature for tensorflow serving
        detection_signature = (
        tf.saved_model.signature_def_utils.build_signature_def(
                inputs=tensor_info_inputs,
                outputs=tensor_info_outputs,
                method_name=signature_constants.PREDICT_METHOD_NAME))

        # "build" graph
        builder.add_meta_graph_and_variables(
                sess, [tf.saved_model.tag_constants.SERVING],
                signature_def_map={
                'detection_signature':
                        detection_signature,
                signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                        detection_signature,
                },
                main_op=tf.tables_initializer()
        )
        # save graph
        builder.save()

Источники, которые я использовал, чтобы найти правильное решение:

Пример кодирования

Объяснение имен областей действия

Tensorflow Github Issue # 22162

Дополнительно: Если вам сложно найти правильные входные и выходные узлы, вы можете запустить это для отображения графика:

graph_op = g.get_operations()
for i in graph_op:
    print(i.node_def)