TF2：为张量流服务添加预处理到预训练的已保存模型（扩展已保存模型的图）

Question

我升级到 TensorFlow 2，现在在扩展预训练模型并进行一些额外预处理时遇到问题。

我有一个预训练的对象检测模型 (SSD ResNet50 FPN)，我想将其部署到 TensorFlow 服务。我想加载 SavedModel 并添加必要的预处理以直接接受 base64 编码的 jpeg。我之前用 TF 1.x 和另一个 Keras 模型做过这个，它可以工作：

string_inp = tf.placeholder(tf.string, shape=(None,), name='base64_in')
imgs_map = tf.map_fn(
    tf.image.decode_image,
    string_inp,
    dtype=tf.uint8
)
imgs_map.set_shape((None, None, None, 3))
imgs = tf.image.resize_images(imgs_map, [456, 456], method=tf.image.ResizeMethod.BILINEAR)
imgs = tf.reshape(imgs, (-1, 456, 456, 3))
img_uint8 = tf.image.convert_image_dtype(imgs, dtype=tf.uint8, saturate=False)

pretrained_model= load_model('my-keras-model.h5', compile=False)
ouput_tensor= pretrained_model(img_uint8)

signature = tf.saved_model.signature_def_utils.predict_signature_def(                                                                        
    inputs={'jpegbase64': string_inp}, outputs={'probabilities': ouput_tensor})

builder.add_meta_graph_and_variables(                                                                                                        
    sess=K.get_session(),                                                                                                                    
    tags=[tf.saved_model.tag_constants.SERVING],                                                                                             
    signature_def_map={                                                                                                                      
        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:                                                                
            signature                                                                                                                        
    })                                                                                                                                       
builder.save()

但是一旦我尝试让它与加载了 TF 的 SavedModel 一起工作 model = tf.keras.models.load_model("my_saved_model") 它抛出： TypeError: 'AutoTrackable' 对象不可调用

我猜它不支持将模型堆叠在我的自定义输入张量之上，但我没有找到任何其他可行的解决方案。我还尝试将 SavedModel 中的输入张量直接与 img_uint8 张量连接，但我不知道如何正确连接它们。有什么想法吗？

Answer 1

好的，我找到了解决方案，我们开始吧：

graph_model = tf.Graph()
sess = tf.Session(graph=graph_model)
sess.as_default()
graph_model.as_default()
model = tf.saved_model.load(sess, export_dir="myModel", tags=['serve'])
graph_model_def = graph_model.as_graph_def()

# here is the important step, create a new graph and DON'T create a new session explicity
graph_base64 = tf.Graph()
graph_base64.as_default()

string_inp = tf.placeholder(tf.string, shape=(None,), name='base64_in')
imgs_map = tf.map_fn(
    tf.image.decode_image,
    string_inp,
    dtype=tf.uint8
)
imgs_map.set_shape((None, None, None, 3))
imgs = tf.image.resize_images(imgs_map, [300, 300], method=tf.image.ResizeMethod.BILINEAR)
imgs = tf.reshape(imgs, (-1, 300, 300, 3))
img_uint8 = tf.image.convert_image_dtype(imgs, dtype=tf.uint8, saturate=False)

# import the model graph with the new input
tf.import_graph_def(graph_model_def, name='', input_map={"image_tensor:0": img_uint8})

重要的是不要创建新会话。如果你这样做，它将不再起作用。 Here 是更详细的描述。