使用 Tensorflow 中保存的 model 进行推理 2：如何控制输入/输出？

Question

. Answers to this question are eligible for a +200 reputation bounty. user1981275 wants to to this question.

Adapting my code from TF1 to TF2.6 I run into trouble. I am trying to add some custom layers to an inception resnet, save the model, and then load and run it.

from tensorflow.keras.layers import Dense                                                                                                                       
from tensorflow.keras.models import Model                                                                                                                       
from tensorflow.keras.applications.inception_resnet_v2 import InceptionResNetV2                                                                                 
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D                                                                                               
import tensorflow as tf                                                                                                                                         
import numpy as np                                                                                                                                              
from PIL import Image                                                                                                                                           
                                                                                                                                                                
export_path = "./save_test"                                                                                                                                     
                                                                                                                                                                
# Get model without top and add two layers                                                                                                                      
base_model = InceptionResNetV2(weights='imagenet', input_tensor=None, include_top=False)                                                                        
out = base_model.output                                                                                                                                         
out = GlobalAveragePooling2D()(out)                                                                                                                             
predictions = Dense(7, activation='softmax', name="output")(out)                                                                                                
                                                                                                                                                                
# Make new model using inputs from base model and custom outputs                                                                                                
model = Model(inputs=base_model.input, outputs=[predictions])                                                                                                   
                                                                                                                                                                
# save model                                                                                                                                                    
tf.saved_model.save(model, export_path)                                                                                                                         
                                                                                                                                                                
# load model and run                                                                                                                                            
with tf.compat.v1.Session(graph=tf.Graph()) as sess:                                                                                                            
    tf.compat.v1.saved_model.loader.load(sess, ['serve'], export_path)                                                                                          
    graph = tf.compat.v1.get_default_graph()                                                                                                                    
                                                                                                                                                                
    img = Image.new('RGB', (299, 299))                                                                                                                          
    x = tf.keras.preprocessing.image.img_to_array(img)                                                                                                          
    x = np.expand_dims(x, axis=0)                                                                                                                               
    x = x[..., :3]                                                                                                                                              
    x /= 255.0                                                                                                                                                  
    x = (x - 0.5) * 2.0                                                                                                                                         
                                                                                                                                                                
    y_pred = sess.run('output/Softmax:0', feed_dict={'serving_default_input_1:0': x})                                                                           

Error: KeyError: "The name 'output/Softmax:0' refers to a Tensor which does not exist. The operation, 'output/Softmax', does not exist in the graph."

What I don't understand: predictions.name is 'output/Softmax:0' , but graph.get_tensor_by_name('output/Softmax:0') tells me it does not exist!

Note: I am aware that I can save and load with TF2's tf.keras.models.save and tf.keras.models.load_model and then run the model with model(x) . However, in my application I have multiple models in memory and I have found that the inference takes much longer than in my TF1 code using the session object. I would therefore like to use the TF1 approach with the session object in compatibility mode.

How can I control the names of input/output when saving? What am I missing?

Answer 1

If you haven't already, you could try something like this (running on TF 2.7) as I think you are referencing the wrong keys in SignatureDef :

from tensorflow.keras.layers import Dense                                                                                                                       
from tensorflow.keras.models import Model                                                                                                                       
from tensorflow.keras.applications.inception_resnet_v2 import InceptionResNetV2                                                                                 
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D                                                                                               
import tensorflow as tf                                                                                                                                         
import numpy as np                                                                                                                                              
from PIL import Image                                                                                                                                           
                                                                                                                                                                
export_path = "./save_test"                                                                                                                                     
                                                                                                                                                                
base_model = InceptionResNetV2(weights='imagenet', input_tensor=None, include_top=False)                                                                        
out = base_model.output                                                                                                                                         
out = GlobalAveragePooling2D()(out)                                                                                                                             
predictions = Dense(7, activation='softmax', name="output")(out)                                                                                                
model = Model(inputs=base_model.input, outputs=[predictions])                                                                                                   
                                                                                                                                                              
tf.saved_model.save(model, export_path)

with tf.compat.v1.Session(graph=tf.Graph()) as sess:                                                                                                            
    meta_graph = tf.compat.v1.saved_model.loader.load(sess, ["serve"], export_path)
    sig_def = meta_graph.signature_def[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    input_key = list(dict(sig_def.inputs).keys())[0]
    input_name = sig_def.inputs[input_key].name
    output_name = sig_def.outputs['output'].name
    img = Image.new('RGB', (299, 299))                                                                                                                          
    x = tf.keras.preprocessing.image.img_to_array(img)                                                                                                          
    x = np.expand_dims(x, axis=0)                                                                                                                               
    x = x[..., :3]                                                                                                                                              
    x /= 255.0                                                                                                                                                  
    x = (x - 0.5) * 2.0   
    y_pred = sess.run(output_name, feed_dict={input_name: x})        
    print(y_pred)  

INFO:tensorflow:Restoring parameters from ./save_test/variables/variables
[[0.14001141 0.13356228 0.14509581 0.22432518 0.16313255 0.11899492
  0.07487784]]

You could also take a look at the SignatureDef for input and output information:

print(meta_graph.signature_def)
{'serving_default': inputs {
  key: "input_2"
  value {
    name: "serving_default_input_2:0"
    dtype: DT_FLOAT
    tensor_shape {
      dim {
        size: -1
      }
      dim {
        size: -1
      }
      dim {
        size: -1
      }
      dim {
        size: 3
      }
    }
  }
}
outputs {
  key: "output"
  value {
    name: "StatefulPartitionedCall:0"
    dtype: DT_FLOAT
    tensor_shape {
      dim {
        size: -1
      }
      dim {
        size: 7
      }
    }
  }
}
method_name: "tensorflow/serving/predict"
, '__saved_model_init_op': outputs {
  key: "__saved_model_init_op"
  value {
    name: "NoOp"
    tensor_shape {
      unknown_rank: true
    }
  }
}
}

If you remove the first layer of your base_model and add a new Input layer, you can use static key names like this:

from tensorflow.keras.layers import Dense                                                                                                                       
from tensorflow.keras.models import Model                                                                                                                       
from tensorflow.keras.applications.inception_resnet_v2 import InceptionResNetV2                                                                                 
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D                                                                                               
import tensorflow as tf                                                                                                                                         
import numpy as np                                                                                                                                              
from PIL import Image                                                                                                                                           
                                                                                                                                                                
export_path = "./save_test"                                                                                                                                     
                                                                                                                                                                
base_model = InceptionResNetV2(weights='imagenet', input_tensor=None, include_top=False)
base_model.layers.pop(0)
new_input = tf.keras.layers.Input(shape=(299,299,3), name='input')
out = base_model(new_input)                                                                                                                                        
out = GlobalAveragePooling2D()(out)                                                                                                                             
predictions = Dense(7, activation='softmax', name="output")(out) 

model = Model(inputs=new_input, outputs=[predictions])                                                                                                   
tf.saved_model.save(model, export_path)

with tf.compat.v1.Session(graph=tf.Graph()) as sess:                                                                                                            
    meta_graph = tf.compat.v1.saved_model.loader.load(sess, ["serve"], export_path)
    sig_def = meta_graph.signature_def[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    input_name = sig_def.inputs['input'].name
    output_name = sig_def.outputs['output'].name
    img = Image.new('RGB', (299, 299))                                                                                                                          
    x = tf.keras.preprocessing.image.img_to_array(img)                                                                                                          
    x = np.expand_dims(x, axis=0)                                                                                                                               
    x = x[..., :3]                                                                                                                                              
    x /= 255.0                                                                                                                                                  
    x = (x - 0.5) * 2.0   
    y_pred = sess.run(output_name, feed_dict={input_name: x})        
    print(y_pred)   

INFO:tensorflow:Restoring parameters from ./save_test/variables/variables
[[0.21079363 0.10773096 0.07287834 0.06983061 0.10538215 0.09172108
  0.34166315]]