PB转TFLite遇ValueError：需指定import/image_feed输入形状

Hey there! Let's break down why you're running into this ValueError and how to fix it. The core issue here is that the TensorFlow Lite converter can't automatically infer fixed shapes for your input tensors from the frozen PB model—many pre-trained captioning models use dynamic input shapes (like variable batch sizes) which TFLite doesn't support natively for conversion.

Step 1: Know Your Input Tensors' Required Shapes #

First, let's map out what each input needs based on the Show-and-Tell model architecture:

• import/image_feed: This is the input image tensor. The original model uses InceptionV3, which expects 299x299x3 RGB images with a batch size of 1 (for inference).
• import/input_feed: This feeds single tokens to the LSTM during caption generation. For inference, this should be a batch of 1, so shape [1].
• import/lstm/state_feed: This is the initial state of the LSTM. The Show-and-Tell model uses a 512-dimensional LSTM state, so shape [1, 512].

Step 2: Modify Your Code to Specify Input Shapes #

You have two straightforward ways to resolve this:

Option 1: Update the Session-Based Conversion

Add lines to set fixed shapes for your input tensors right after fetching them:

g = tf.get_default_graph()
input_tensors = [g.get_tensor_by_name(x) for x in input_names]
output_tensors = [g.get_tensor_by_name(x) for x in output_names]

# Add these lines to define fixed input shapes
input_tensors[0].set_shape([1, 299, 299, 3])  # Image input: batch 1, 299x299 RGB
input_tensors[1].set_shape([1])               # Token input: single token per batch
input_tensors[2].set_shape([1, 512])          # LSTM state: batch 1, 512 units

# Proceed with conversion as before
converter = tf.lite.TFLiteConverter.from_session(sess, input_tensors, output_tensors)
model = converter.convert()

Option 2: Use from_frozen_graph (Simpler, No Session Needed)

This method lets you directly specify input shapes without creating a TensorFlow Session, which is cleaner:

# Replace the Session and tensor fetching code with this block
converter = tf.lite.TFLiteConverter.from_frozen_graph(
    GRAPH_LOCATION,
    input_arrays=['import/image_feed', 'import/input_feed', 'import/lstm/state_feed'],
    output_arrays=['import/softmax', 'import/lstm/state', 'import/lstm/initial_state'],
    input_shapes={
        'import/image_feed': [1, 299, 299, 3],
        'import/input_feed': [1],
        'import/lstm/state_feed': [1, 512]
    }
)
model = converter.convert()

Quick Checks to Avoid Issues #

• If your model uses a different input image size (not 299x299), adjust the shape to match what the original PB model expects.
• If the LSTM state dimension isn't 512, verify the original model's LSTM configuration and update that shape accordingly.

Once you make these changes, the converter will have all the shape information it needs to successfully convert your PB model to TFLite.

内容的提问来源于stack exchange，提问作者Tarini Ramsewak