TensorFlow模型部署ML Engine预测失败：contents需为标量

Hey there, let's sort out that prediction error you're hitting! The error message contents must be scalar, got shape [1] tells us exactly what's going wrong: your trained model expects a single scalar string tensor (one image's raw bytes) as input, but Google Cloud ML Engine automatically wraps your input into a batch tensor (shape [1], since you're sending one image at a time). These two shapes don't line up, causing the failure.

Here are two solutions to fix this—one quick for single-image predictions, and a more flexible long-term fix that supports batch processing:

Solution 1: Quick Fix - Adjust the Prediction Command #

If you only need to predict on single images, you can skip modifying your model entirely. Instead, use the --instance parameter to pass your input directly, which avoids ML Engine wrapping it into a batch:

# First, generate the base64-encoded string of your image
IMG_BASE64=$(python -c 'import base64, sys; print(base64.b64encode(open(sys.argv[1], "rb").read()).decode())' flower.jpg)

# Predict using the single instance parameter
gcloud ml-engine predict --model ${MODEL_NAME} --instance "{\"image_bytes\": {\"b64\": \"${IMG_BASE64}\"}}"

Solution 2: Long-Term Fix - Update Your Export Script for Batch Support #

This is the better approach if you ever want to predict on multiple images at once, or if you want to align with Cloud ML Engine's expected input format. We'll modify your export script to wrap the model so it accepts a batch of image bytes, then processes each one individually.

Here's the updated export script:

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

input_graph = 'retrained_graph.pb'
saved_model_dir = 'my_model'

with tf.Graph().as_default() as graph:
    # Load your pre-trained graph definition
    with tf.gfile.FastGFile(input_graph, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    # Define a placeholder that accepts a batch of image bytes (any number of images)
    batch_image_bytes = tf.placeholder(tf.string, shape=[None], name='batch_image_bytes')

    # Create a function to run prediction on a single image's bytes
    def predict_single_image(image_bytes):
        # Reuse your original model's logic for a single image
        single_prediction = tf.import_graph_def(
            graph_def,
            input_map={'DecodeJpeg/contents:0': image_bytes},
            return_elements=['final_result:0'],
            name=f'single_instance_{tf.get_default_graph().unique_name("")}'
        )[0]
        return single_prediction

    # Apply the single-image prediction to every element in the batch
    batch_predictions = tf.map_fn(
        predict_single_image,
        batch_image_bytes,
        dtype=tf.float32,
        back_prop=False
    )

    # Build the SavedModel signature for batch inference
    inputs = {'image_bytes': tf.saved_model.utils.build_tensor_info(batch_image_bytes)}
    outputs = {'predictions': tf.saved_model.utils.build_tensor_info(batch_predictions)}
    signature = tf.saved_model.signature_def_utils.build_signature_def(
        inputs=inputs,
        outputs=outputs,
        method_name=signature_constants.PREDICT_METHOD_NAME
    )

with tf.Session(graph=graph) as sess:
    # Initialize any variables in the graph
    sess.run(tf.global_variables_initializer())
    # Save the updated model that supports batch inputs
    builder = saved_model_builder.SavedModelBuilder(saved_model_dir)
    builder.add_meta_graph_and_variables(
        sess,
        [tag_constants.SERVING],
        signature_def_map={'serving_default': signature}
    )
    builder.save()

Deploy the Updated Model and Test #

1. Upload the new model to your Cloud Storage bucket, then deploy it to ML Engine:

# Create the model if you haven't already
gcloud ml-engine models create ${MODEL_NAME} --regions us-central1

# Deploy the updated version
gcloud ml-engine versions create v2 --model ${MODEL_NAME} --origin gs://your-bucket/my_model --runtime-version 1.15

(Replace your-bucket with your GCS bucket name, and use a TensorFlow runtime version that matches what you used for training.)

2. Now you can use your original prediction command, and it should work perfectly:

gsutil cp gs://cloud-ml-data/img/flower_photos/tulips/4520577328_a94c11e806_n.jpg flower.jpg
python -c 'import base64, sys, json; img = base64.b64encode(open(sys.argv[1], "rb").read()).decode(); print(json.dumps({"image_bytes": {"b64": img}}))' flower.jpg > request.json
gcloud ml-engine predict --model ${MODEL_NAME} --version v2 --json-instances request.json

Bonus: Batch Multiple Images #

With this updated model, you can predict on multiple images at once by adding more lines to request.json—each line is a separate image instance in the same JSON format.

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