loader.py 文件源码

def __init__(self, model_base_dir, num_parallel_predictions=2):
        possible_dirs = os.listdir(model_base_dir)
        model_dir = os.path.join(model_base_dir, max(possible_dirs))
        print("Loading {}".format(model_dir))

        self.sess = tf.get_default_session()
        loaded_model = tf.saved_model.loader.load(self.sess, ['serve'], model_dir)
        assert 'serving_default' in list(loaded_model.signature_def)

        input_dict, output_dict = _signature_def_to_tensors(loaded_model.signature_def['serving_default'])
        self._input_tensor = input_dict['images']
        self._output_dict = output_dict
        self.sema = Semaphore(num_parallel_predictions)