Блог

Мне нужно загрузить обученную модель в Tensorflow и сохранить ее в формате SavedModel для использования в обслуживании Tensorflow. Я следую этому руководству. Я загружаю модель следующим образом:

     self.graph = tf.Graph()

    def _graph_hourglass(self, inputs):
        ### All layers are here ###
    ...

    def generate_model(self):
        with tf.device(self.cpu):
            with tf.name_scope('inputs'):
                self.img = tf.compat.v1.placeholder(dtype=tf.float32,
                 shape=(None, 256, 256, 3), name='input_img')
            self.output = self._graph_hourglass(self.img)

    def model_init(self):
        with self.graph.as_default():
            self.generate_model()

    def restore(self, load='path'):
        with tf.name_scope('Session'):
            with tf.device(self.cpu):
                self.Session = tf.compat.v1.Session()
                saver = tf.compat.v1.train.Saver()
                saver.restore(self.Session, load)

    def load_model(self, load = 'path'):
        with self.graph.as_default():
            self.restore(load)

    ...

    self.model_init()
    self.load_model(load='path/to/model')
 

Затем я использую класс SavedModelBuilder из учебника:

     export_path = 'new/path'
    print('Exporting trained model to', export_path)
    builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_path)
    tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(self.img)
    tensor_info_y = tf.compat.v1.saved_model.utils.build_tensor_info(self.output)

    prediction_signature = (
        tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
            inputs={'image': tensor_info_x},
            outputs={'joints': tensor_info_y},
            method_name=tf.compat.v1.saved_model.signature_constants.PREDICT_METHOD_NAME))

    builder.add_meta_graph_and_variables(
        self.Session, [tf.compat.v1.saved_model.tag_constants.SERVING],
        signature_def_map={
            tf.compat.v1.saved_model.signature_constants
                .DEFAULT_SERVING_SIGNATURE_DEF_KEY:
            prediction_signature,
        },
        strip_default_attrs=True)

    builder.save()
    print('Done exporting!')
 

Модель сохраняется без ошибок, я вижу «Готово к экспорту!» в консоли, но файл model .pb пуст! Это всего 252 байта. В чем здесь может быть проблема? Буду признателен за любую помощь.