Как сохранить модель Tensorflow 2.0, которая использует модель из контрольной точки .meta из Tensorflow 1.xx как часть?

Question

Я обучен модели с tensorflow 1.15 и сохранен как контрольная точка (с файлами .meta, .index и .data).

Мне нужно добавить некоторые дополнительные операции в начало идо конца этого графика. Некоторые из этих операций существуют только в tensorflow 2.0 и tensorflow_text 2.0. После этого я хочу сохранить эту модель для tensorflow-serving.

Что я пытался сделать: используя tensorflow 2.0, я сохранил ее как .pb файл, подобный этому.

trained_checkpoint_prefix = 'path/to/model'
export_dir = os.path.join('path/to/export', '0')

graph = tf.Graph()
with tf.compat.v1.Session(graph=graph) as sess:
    # Restore from checkpoint
    loader = tf.compat.v1.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
    loader.restore(sess, trained_checkpoint_prefix)

    # Export checkpoint to SavedModel
    builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)

    classification_signature = tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
        inputs={
            'token_indices': get_tensor_info('token_indices_ph:0'),
            'token_mask': get_tensor_info('token_mask_ph:0'),
            'y_mask': get_tensor_info('y_mask_ph:0'),
        },
        outputs={'probas': get_tensor_info('ner/Softmax:0'), 'seq_lengths': get_tensor_info('ner/Sum:0')},
        method_name='predict',
    )

    builder.add_meta_graph_and_variables(sess,
                                         [tf.saved_model.TRAINING, tf.saved_model.SERVING],
                                         strip_default_attrs=True, saver=loader,
                                         signature_def_map={'predict': classification_signature}) # , clear_devices=True)
    builder.save()  

Послечто я создал tf.keras.Model, который загружает .pb модель и выполняет весь необходимый персонал:

import os
from pathlib import Path

import tensorflow as tf
import tensorflow_text as tf_text


class BertPipeline(tf.keras.Model):
    def __init__(self):
        super().__init__()

        vocab_file = Path('path/to/vocab.txt')
        vocab = vocab_file.read_text().split('\n')[:-1]
        self.vocab_table = self.create_table(vocab)

        export_dir = 'path/to/pb/model'
        self.model = tf.saved_model.load(export_dir)

        self.bert_tokenizer = BertTokenizer(
            self.vocab_table,
            max_chars_per_token=15,
                token_out_type=tf.int64
            ,
            lower_case=True,
        )

        self.to_dense = tf_text.keras.layers.ToDense()

    def call(self, texts):
        tokens = self.bert_tokenizer.tokenize(texts)
        tokens = tf.cast(tokens, dtype=tf.int32)

        mask = self.make_mask(tokens)
        token_ids = self.make_token_ids(tokens)

        token_indices = self.to_dense(token_ids)
        token_mask = self.to_dense(tf.ones_like(mask))
        y_mask = self.to_dense(mask)

        res = self.model.signatures['predict'](
            token_indices=token_indices,
            token_mask=token_mask,
            y_mask=y_mask,
        )

        starts_range = tf.range(0, tf.shape(res['seq_lengths'])[0]) * tf.shape(res['probas'])[1]
        row_splits = tf.reshape(
            tf.stack(
                [
                    starts_range,
                    starts_range + res['seq_lengths'],
                ],
                axis=1,
            ),
            [-1],
        )

        row_splits = tf.concat(
            [
                row_splits,
                tf.expand_dims(tf.shape(res['probas'])[0] * tf.shape(res['probas'])[1], 0),
            ],
            axis=0,
        )

        probas = tf.RaggedTensor.from_row_splits(
            tf.reshape(res['probas'], [-1, 2]),
            row_splits,
        )[::2]

        probas

        return probas

    def make_mask(self, tokens):
        masked_suff = tf.concat(
            [
                tf.ones_like(tokens[:, :, :1], dtype=tf.int32),
                tf.zeros_like(tokens[:, :, 1:], dtype=tf.int32),
            ],
            axis=-1,
        )

        joined_mask = self.join_wordpieces(masked_suff)
        return tf.concat(
            [
                tf.zeros_like(joined_mask[:, :1], dtype=tf.int32),
                joined_mask,
                tf.zeros_like(joined_mask[:, :1], dtype=tf.int32),
            ],
            axis=-1,
        )

    def make_token_ids(self, tokens):
        joined_tokens = self.join_wordpieces(tokens)

        return tf.concat(
            [
                tf.fill(
                    [joined_tokens.nrows(), 1],
                    tf.dtypes.cast(
                        self.vocab_table.lookup(tf.constant('[CLS]')),
                        dtype=tf.int32,
                    )
                ),
                self.join_wordpieces(tokens),
                tf.fill(
                    [joined_tokens.nrows(), 1],
                    tf.dtypes.cast(
                        self.vocab_table.lookup(tf.constant('[SEP]')),
                        dtype=tf.int32,
                    )
                ),
            ],
            axis=-1,
        )


    def join_wordpieces(self, wordpieces):
        return tf.RaggedTensor.from_row_splits(
            wordpieces.flat_values, tf.gather(wordpieces.values.row_splits,
                                              wordpieces.row_splits))

    def create_table(self, vocab, num_oov=1):
        init = tf.lookup.KeyValueTensorInitializer(
            vocab,
            tf.range(tf.size(vocab, out_type=tf.int64), dtype=tf.int64),
            key_dtype=tf.string,
            value_dtype=tf.int64)
        return tf.lookup.StaticVocabularyTable(init, num_oov, lookup_key_dtype=tf.string)

Когда я вызываю этот код, он отлично работает:

bert_pipeline = BertPipeline()
print(bbert_pipeline(["Some test string", "another string"]))

---
<tf.RaggedTensor [[[0.17896245419979095, 0.8210375308990479], [0.8825045228004456, 0.11749550700187683], [0.9141901731491089, 0.0858098641037941]], [[0.2768123149871826, 0.7231876850128174], [0.9391192197799683, 0.060880810022354126]]]>

Но японятия не имею, как его сохранить. Если я правильно понимаю, tf.keras.Model не рассматривайте self.model и self.bert_tokenizer как часть модели. Если я вызываю bert_pipeline.summary(), то нет никаких операций:

bert_pipeline.build([])
bert_pipeline.summary()

---
Model: "bert_pipeline_3"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
to_dense (ToDense)           multiple                  0         
=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0
_________________________________________________________________

Кроме того, я попытался запустить его с tensorflow.compat.v1, используя явные Session и Graph, но в этом случае я просто не могу загрузитьмодель правильно. Тот же код с import tensorflow.compat.v1 as tf и шаблон для tensorflow 1.xx не могут инициализировать некоторые переменные:

# tf.saved_model.load(export_dir) changed to tf.saved_model.load_v2(export_dir) above

import tensorflow.compat.v1 as tf
graph = tf.Graph()
with tf.Session(graph=graph) as sess:
    bert_pipeline = BertPipeline()
    texts = tf.placeholder(tf.string, shape=[None], name='texts')

    res_tensor = bert_pipeline(texts)

    sess.run(tf.tables_initializer())
    sess.run(tf.global_variables_initializer())

    sess.run(res_tensor, feed_dict={texts: ["Some test string", "another string"]})

---
FailedPreconditionError                   Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/client/session.py in _do_call(self, fn, *args)
   1364     try:
-> 1365       return fn(*args)
   1366     except errors.OpError as e:

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/client/session.py in _run_fn(feed_dict, fetch_list, target_list, options, run_metadata)
   1349       return self._call_tf_sessionrun(options, feed_dict, fetch_list,
-> 1350                                       target_list, run_metadata)
   1351 

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/client/session.py in _call_tf_sessionrun(self, options, feed_dict, fetch_list, target_list, run_metadata)
   1442                                             fetch_list, target_list,
-> 1443                                             run_metadata)
   1444 

FailedPreconditionError: [_Derived_]{{function_node __inference_pruned_77348}} {{function_node __inference_pruned_77348}} Attempting to use uninitialized value bert/encoder/layer_3/attention/self/query/kernel
     [[{{node bert/encoder/layer_3/attention/self/query/kernel/read}}]]
     [[bert_pipeline/StatefulPartitionedCall]]

During handling of the above exception, another exception occurred:

FailedPreconditionError                   Traceback (most recent call last)
<ipython-input-15-5a0a45327337> in <module>
     21     sess.run(tf.global_variables_initializer())
     22 
---> 23     sess.run(res_tensor, feed_dict={texts: ["Some test string", "another string"]})
     24 
     25 #     print(res)

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/client/session.py in run(self, fetches, feed_dict, options, run_metadata)
    954     try:
    955       result = self._run(None, fetches, feed_dict, options_ptr,
--> 956                          run_metadata_ptr)
    957       if run_metadata:
    958         proto_data = tf_session.TF_GetBuffer(run_metadata_ptr)

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/client/session.py in _run(self, handle, fetches, feed_dict, options, run_metadata)
   1178     if final_fetches or final_targets or (handle and feed_dict_tensor):
   1179       results = self._do_run(handle, final_targets, final_fetches,
-> 1180                              feed_dict_tensor, options, run_metadata)
   1181     else:
   1182       results = []

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/client/session.py in _do_run(self, handle, target_list, fetch_list, feed_dict, options, run_metadata)
   1357     if handle is None:
   1358       return self._do_call(_run_fn, feeds, fetches, targets, options,
-> 1359                            run_metadata)
   1360     else:
   1361       return self._do_call(_prun_fn, handle, feeds, fetches)

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/client/session.py in _do_call(self, fn, *args)
   1382                     '\nsession_config.graph_options.rewrite_options.'
   1383                     'disable_meta_optimizer = True')
-> 1384       raise type(e)(node_def, op, message)
   1385 
   1386   def _extend_graph(self):
FailedPreconditionError: [_Derived_]  Attempting to use uninitialized value bert/encoder/layer_3/attention/self/query/kernel
     [[{{node bert/encoder/layer_3/attention/self/query/kernel/read}}]]
     [[bert_pipeline/StatefulPartitionedCall]]

Пожалуйста, если у вас есть какие-то мысли, как исправить мой подход к сохранению графика или, может быть, вы знаете, как это сделатьэто лучше - скажи мне. Спасибо!

Answer 1

Я решил это. Прежде всего, я не смог сделать это с tf.keras. Я использовал

import tensorflow.compat.v1 as tf

Кроме того, я использовал .meta, .index и контрольную точку бла бла, не обращаясь к '.pb'.

Основное, что я использовал, описано здесь: Tensorflow: Как заменить узел в графе вычислений?

Я сделал 2 разных графика, и после этого объединил их, как в этой части кода:

def _build_model(self):
    with tf.Graph().as_default() as g_1:
        self.lookup_table = self._make_lookup_table()

        init_table = tf.initialize_all_tables()

        self.bert_tokenizer = BertTokenizer(
            self.lookup_table,
            max_chars_per_token=15,
            token_out_type=tf.int64,
            lower_case=True,
        )

        self.texts_ph = tf.placeholder(tf.string, shape=(None,), name="texts_ph")  # input

        words_without_name, tokens_int_64 = self.bert_tokenizer.tokenize(self.texts_ph)
        words = words_without_name.to_tensor(default_value='', name='tokens')

        tokens = tf.cast(tokens_int_64, dtype=tf.int32)

        mask = self._make_mask(tokens)
        token_ids = self._make_token_ids(tokens)

        self.token_indices = token_ids.to_tensor(default_value=0, name='token_indices')  # output 1
        self.token_mask = tf.ones_like(mask).to_tensor(default_value=0, name='token_mask') # output 2
        self.y_mask = mask.to_tensor(default_value=0, name='y_mask') # output 3

    with tf.Graph().as_default() as g_2:
        sess = tf.Session()
        path_to_model = 'path/to/model'
        self._load_model(sess, path_to_model)

        token_indices_2 = g_2.get_tensor_by_name('token_indices_ph:0'),
        token_mask_2 = g_2.get_tensor_by_name('token_mask_ph:0'),
        y_mask_2 = g_2.get_tensor_by_name('y_mask_ph:0'),

        probas = g_2.get_tensor_by_name('ner/Softmax:0')
        seq_lengths = g_2.get_tensor_by_name('ner/Sum:0')

        exclude_scopes = ('Optimizer', 'learning_rate', 'momentum', 'EMA/BackupVariables')
        all_vars = variables._all_saveable_objects()
        self.vars_to_save = [var for var in all_vars if all(sc not in var.name for sc in exclude_scopes)]
        self.saver = tf.train.Saver(self.vars_to_save

    g_1_def = g_1.as_graph_def()
    g_2_def = g_2.as_graph_def()

    with tf.Graph().as_default() as g_combined:
        self.texts = tf.placeholder(tf.string, shape=(None,), name="texts")

        y1, y2, y3, self.init_table, self.words = tf.import_graph_def(
           g_1_def, input_map={"texts_ph:0": self.texts},
           return_elements=["token_indices/GatherV2:0", "token_mask/GatherV2:0", "y_mask/GatherV2:0", 'init_all_tables', 'tokens/GatherV2:0'],
           name='',
        )

        self.dense_probas, self.lengths = tf.import_graph_def(
            g_2_def, input_map={"token_indices_ph:0": y1, "token_mask_ph:0": y2, "y_mask_ph:0": y3},
            return_elements=["ner/Softmax:0", "ner/Sum:0"],
            name='',
        )

        self.sess = tf.Session(graph=g_combined)
        self.graph = g_combined

        self.sess.run(self.init_table)

        vars_dict_to_save = {v.name[:-2]: g_2.get_tensor_by_name(v.name) for v in self.vars_to_save}
        self.saver.restore(self.sess, path_to_model)

Вы можете заметить, что я вызываю self._load_model(sess, path_to_model) для загрузки модели, создаю saver с необходимыми переменными, а после этого снова загружаю модель с self.saver.save(sess, path_to_model). Первая загрузка необходима для чтения предварительно сохраненного графика и доступа к его тензорам. Второе необходимо для загрузки весов в другом сеансе с объединенным графом g_combined. Я думаю, что есть способ сделать это без загрузки данных с диска два раза, но он работает, и я не хочу его ломать: -).

Еще одна важная вещь - vars_dict_to_save. Этот дикт необходим для отображения между загруженными весами и тензорами в графах.

После этого у вас есть полный граф со всеми операциями, поэтому вы можете назвать его так:

def __call__(self, texts):
    lengths, words, probs = self.sess.run(
        [self.lengths, self.words, self.dense_probas],
        feed_dict={
            self.texts: texts
        },
    )
    return lengths, words, probs

Обратите вниманиек реализации метода __call__. Он использует сессию, которую я создал с объединенным графиком.

После того, как у вас есть полный график с загруженными весами, легко экспортировать график для обслуживания:

def export(self, export_dir):
    with self.graph.as_default():
        builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

        predict_signature = tf.saved_model.signature_def_utils.predict_signature_def(
            inputs={
                'texts': self.texts,
            },
            outputs={
                'lengths': self.lengths,
                'tokens': self.words,
                'probs': self.dense_probas,
            },
        )

        builder.add_meta_graph_and_variables(
            self.sess,
            [tf.saved_model.SERVING],
            strip_default_attrs=True,
            signature_def_map={'predict': predict_signature},
            saver=self.saver,
            main_op=self.init_table,
        )
        builder.save()

Есть несколько важных моментов: -Использовать объединенный граф .as_default() - Используйте те же сеансы, которые вы использовали с объединенным графом. - Используйте ту же заставку, которую вы использовали для загрузки весов в объединенном графике. - Добавьте main main_op, если у вас есть таблицы, которые нужно инициализировать.

Я буду рад, если это кому-нибудь поможет :-). Это было не тривиально для меня, и я потратил много времени, чтобы заставить его работать.

PS BertTokenizer в этом коде немного отличается от такого класса от tensorflow_text, но это не связано с проблемой.