TF2: добавить предварительную обработку в предварительно обученную сохраненную модель для обслуживания тензорного потока (расширение графика сохраненной модели)

Question

Я обновился до TensorFlow 2 и теперь сталкиваюсь с проблемой при расширении предварительно обученной модели с некоторой дополнительной предварительной обработкой.

У меня есть предварительно обученная модель обнаружения объектов (SSD ResNet50 FPN), которую я хочу развернуть в TensorFlow служить. Я хочу загрузить SavedModel и добавить необходимую предварительную обработку для непосредственного приема jpegs в формате base64. Я делал это раньше с 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()

Но как только я пытаюсь заставить его работать с SavedModel, загруженным с TF model = tf.keras.models.load_model("my_saved_model"), он выдает: TypeError: ' AutoTrackable 'объект не может быть вызван

Я думаю, он не поддерживает укладку модели поверх моего пользовательского тензора ввода, но я не нашел другого рабочего решения для него. Я также экспериментировал с подключением входного тензора из SavedModel напрямую с тензором img_uint8, но я не знаю, как их правильно подключить. Есть идеи?

Answer 1

Хорошо, я нашел решение, здесь мы go:

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})

Важная часть - НЕ создавать новый сеанс. Если вы это сделаете, это больше не будет работать. Здесь - более подробное описание.