Выберите переменные для сохранения в модели при использовании tf.saved_model.simple_save - PullRequest
Новая
       4

Выберите переменные для сохранения в модели при использовании tf.saved_model.simple_save

0 голосов
/ 05 ноября 2019

Я загружаю сохраненную модель из tf-hub, затем настраиваю ее и сохраняю для будущего использования. На этом шаге размер модуля увеличивается с 800 Мб до 2,3 Гб. Из-за переменных оптимизатора Adam.
Итак - вопрос в том, могу ли я каким-то образом сохранить точно настроенную модель с тем же списком переменных, как она была импортирована?

Сохранение модели: 

export_dir = "../model/1"
with tf.Session(graph=tf.Graph()) as sess:
    module = hub.Module("https://tfhub.dev/google/universal-sentence-encoder-large/3", trainable=True)
    text_input = tf.placeholder(dtype=tf.string, shape=[None])
    embeddings = module(text_input)
    sess.run([tf.global_variables_initializer(), tf.tables_initializer()])
    tf.saved_model.simple_save(sess,
        export_dir=export_dir,
        inputs={'text': text_input},
        outputs={'embeddings': embeddings},
        legacy_init_op=tf.tables_initializer())

Тонкая настройка: 

raph=tf.Graph()
with tf.Session(graph=graph) as sess:
    model = tf.saved_model.loader.load(export_dir=export_dir, sess=sess,
                                           tags=[tag_constants.SERVING]) 

    # universal sentence encoder input/output
    input_tensor_name = model.signature_def[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].inputs['text'].name
    in_tensor = tf.get_default_graph().get_tensor_by_name(input_tensor_name)

    embedd_tensor_name = model.signature_def[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].outputs[
        'embeddings'].name
    out_tensor = tf.get_default_graph().get_tensor_by_name(embedd_tensor_name)

    # classification on top of universal sentence encoder
    input_y = tf.placeholder(tf.int32, shape=(None))
    labels = tf.one_hot(input_y, len(set(label)))
    logits = tf.layers.dense(out_tensor, len(set(label)))
    loss = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels)
    optimizer = tf.train.AdamOptimizer(learning_rate=0.001).minimize(loss)

    #init optimizer variables
    uninitialized_vars = []
    for var in tf.all_variables():
        try:
            sess.run(var)
        except tf.errors.FailedPreconditionError:
            uninitialized_vars.append(var)
    sess.run(tf.initialize_variables(uninitialized_vars))

    #train
    for epoch in range(2):
        feed_dict = {
            in_tensor: text,
            input_y: label
        }
        sess.run(optimizer, feed_dict)

    #Save tuned model
    save_path = export_dir+'_mod'   
    tf.saved_model.simple_save(sess, save_path, inputs={'text':in_tensor}, outputs={'embeddings':out_tensor})

Я пришел к выводу - я могу отфильтровать список переменных без дополнительных переменных оптимизатора. Но как сохранить модель в .pb только с этими переменными? 

    #Exclude optimizer vars from save
    all_vars = tf.global_variables()
    exclude_vars = tf.contrib.framework.get_variables_by_suffix('Adam')
    vars = [item for item in all_vars if item not in exclude_vars]

Кажется, это не проблема с tf.train.Saver(tf.trainable_variables()) Но мне нужен формат .pb.

...