Swap a TensorFlow Dataset input pipeline with a placeholder after training

后端 未结 1 1366
误落风尘
误落风尘 2020-12-30 15:05

I\'m working with the new tf.data.Dataset API and I can\'t seem to figure out how to perform inference. Ultimately, I want to convert my model to a TensorRT gra

1条回答
  •  孤城傲影
    2020-12-30 15:24

    Since you already have a trained graph saved in a checkpoint, in theory the simplest solution for you is to export the inference graph via optimize_for_inference.

    This tool works both for already-frozen graphs and, as is your case, for graphs with variables still defined. Assuming you go for the frozen graph way, the first step is to transform your graph's variables in constants via:

    python freeze_graph.py \
    --input_graph=temp/path/graph.pbtxt \
    --input_checkpoint=temp/path/your_model_name.ckpt \
    --output_graph=frozen_model.pb \
    --output_node_names=name_of_the_output_tensor_you_want_to_use
    

    This will generate a new binary file called frozen_model.pb that has the Variable operations replaced with Const ops with the values loaded from the checkpoint file.

    Then, you need to generate the inference graph with:

    python optimize_for_inference.py \
    --input=frozen_model.pb \
    --output=inference.pb \
    --frozen_graph=True \
    --input_names=IteratorGetNext
    --output_names=name_of_the_output_tensor_you_want_to_use
    

    This will replace the IteratorGetNext node with a float placeholder. You might want to choose another node, in which case just change the name. You can also change the type of the generated placeholder via the --placeholder_type_enum option. In that case, you need to provide an integer value matching the datatype you want from the DataType enum.

    NOTE: I said "in theory" because actually inspecting the generated inception graph from a test I made it seems there are still some weird ops in there that are not really necessary for inference. You might have to further process your graph via nvidia's Graph Surgeon or TF's graph transform tool

    0 讨论(0)
提交回复
热议问题